Waveform

class LALSimulationBackend(WaveformBackend):
    """Time-domain waveform backend implemented with LALSimulation."""

    def available_approximants(self) -> list[str]:
        """Return all implemented LAL time-domain approximants."""
        return [
            lalsimulation.GetStringFromApproximant(i)
            for i in range(lalsimulation.NumApproximants)
            if lalsimulation.SimInspiralImplementedTDApproximants(i)
        ]

    def generate_td_waveform(
        self,
        approximant: str,
        tc: float,
        sampling_frequency: float,
        minimum_frequency: float,
        **params: object,
    ) -> dict[str, TimeSeries]:
        """Generate plus/cross polarizations with ``SimInspiralChooseTDWaveform``."""
        remaining = dict(params)
        mass1 = float(_pop_alias(remaining, "detector_frame_mass_1", "mass1"))
        mass2 = float(_pop_alias(remaining, "detector_frame_mass_2", "mass2"))
        distance = float(_pop_alias(remaining, "luminosity_distance", "distance"))
        spin_1x = float(_pop_alias(remaining, "spin_1x", "spin1x", default=0.0))
        spin_1y = float(_pop_alias(remaining, "spin_1y", "spin1y", default=0.0))
        spin_1z = float(_pop_alias(remaining, "spin_1z", "spin1z", default=0.0))
        spin_2x = float(_pop_alias(remaining, "spin_2x", "spin2x", default=0.0))
        spin_2y = float(_pop_alias(remaining, "spin_2y", "spin2y", default=0.0))
        spin_2z = float(_pop_alias(remaining, "spin_2z", "spin2z", default=0.0))
        inclination = float(_pop_alias(remaining, "inclination", default=0.0))
        coa_phase = float(_pop_alias(remaining, "coa_phase", default=0.0))
        if remaining:
            extras = ", ".join(sorted(remaining))
            raise ValueError(f"Unsupported LAL waveform parameters: {extras}")
        if sampling_frequency <= 0:
            raise ValueError("sampling_frequency must be > 0")

        approx_enum = lalsimulation.GetApproximantFromString(approximant)
        lal_params = lal.CreateDict()
        hp, hc = lalsimulation.SimInspiralChooseTDWaveform(
            mass1 * MSUN,
            mass2 * MSUN,
            spin_1x,
            spin_1y,
            spin_1z,
            spin_2x,
            spin_2y,
            spin_2z,
            distance * MPC,
            inclination,
            coa_phase,
            0.0,
            0.0,
            0.0,
            1.0 / sampling_frequency,
            minimum_frequency,
            minimum_frequency,
            lal_params,
            approx_enum,
        )
        t0 = float(hp.epoch) + tc
        dt = hp.deltaT
        return {
            "plus": TimeSeries(hp.data.data, t0=t0, dt=dt),
            "cross": TimeSeries(hc.data.data, t0=t0, dt=dt),
        }

class PyCBCBackend(WaveformBackend):
    """Time-domain waveform backend implemented with PyCBC."""

    def __init__(self) -> None:
        """Require PyCBC only when this backend is instantiated."""
        try:
            self._pycbc_waveform = importlib.import_module("pycbc.waveform")
        except ImportError as exc:
            raise ImportError(_PYCBC_IMPORT_ERROR) from exc

    def available_approximants(self) -> list[str]:
        """Return all PyCBC time-domain approximants."""
        return list(self._pycbc_waveform.td_approximants())

    def generate_td_waveform(
        self,
        approximant: str,
        tc: float,
        sampling_frequency: float,
        minimum_frequency: float,
        **params: object,
    ) -> dict[str, TimeSeries]:
        """Generate plus/cross polarizations through ``pycbc_waveform_wrapper``."""
        pycbc_waveform_wrapper = importlib.import_module("gwmock_signal.waveform.pycbc_wrapper").pycbc_waveform_wrapper
        remaining = dict(params)
        translated = {
            "mass1": _pop_alias(remaining, "detector_frame_mass_1", "mass1"),
            "mass2": _pop_alias(remaining, "detector_frame_mass_2", "mass2"),
            "distance": _pop_alias(remaining, "luminosity_distance", "distance"),
            "spin1x": _pop_alias(remaining, "spin_1x", "spin1x", default=0.0),
            "spin1y": _pop_alias(remaining, "spin_1y", "spin1y", default=0.0),
            "spin1z": _pop_alias(remaining, "spin_1z", "spin1z", default=0.0),
            "spin2x": _pop_alias(remaining, "spin_2x", "spin2x", default=0.0),
            "spin2y": _pop_alias(remaining, "spin_2y", "spin2y", default=0.0),
            "spin2z": _pop_alias(remaining, "spin_2z", "spin2z", default=0.0),
            "inclination": _pop_alias(remaining, "inclination", default=0.0),
            "coa_phase": _pop_alias(remaining, "coa_phase", default=0.0),
        }
        translated.update(remaining)
        return pycbc_waveform_wrapper(
            tc=tc,
            sampling_frequency=sampling_frequency,
            minimum_frequency=minimum_frequency,
            waveform_model=approximant,
            **translated,
        )

class WaveformBackend(ABC):
    """Abstract interface for time-domain waveform generators."""

    @abstractmethod
    def available_approximants(self) -> list[str]:
        """Return supported time-domain approximant names."""

    @abstractmethod
    def generate_td_waveform(
        self,
        approximant: str,
        tc: float,
        sampling_frequency: float,
        minimum_frequency: float,
        **params: object,
    ) -> dict[str, TimeSeries]:
        """Generate ``plus`` and ``cross`` GWpy time series."""

class WaveformFactory:
    """Registry and dispatcher for time-domain waveform generators.

    On construction, every name returned by the configured backend is registered
    and mapped to that backend's ``generate_td_waveform`` implementation.
    You may register additional names pointing at custom callables. See package docs for examples.
    """

    def __init__(self, backend: WaveformBackend | None = None) -> None:
        """Build the registry of built-in backend approximants.

        Note:
            Enumerating approximants can be slow; reuse one factory
            instance in tight loops instead of creating many factories.
        """
        self._backend = backend or LALSimulationBackend()
        self._models: dict[str, Callable[..., dict[str, TimeSeries]]] = {
            name: self._wrap_backend_call(name) for name in self._backend.available_approximants()
        }

    def _wrap_backend_call(self, default_approximant: str) -> Callable[..., dict[str, TimeSeries]]:
        """Adapt the backend interface to the factory's callable registry contract."""

        def _call_backend(
            *,
            waveform_model: str | None = None,
            approximant: str | None = None,
            tc: float,
            sampling_frequency: float,
            minimum_frequency: float,
            **params: Any,
        ) -> dict[str, TimeSeries]:
            for supplied_name in (waveform_model, approximant):
                if supplied_name is not None and supplied_name != default_approximant:
                    raise ValueError(
                        f"Registered model {default_approximant!r} cannot be called with conflicting "
                        f"approximant {supplied_name!r}."
                    )
            model_name = default_approximant
            return self._backend.generate_td_waveform(
                approximant=model_name,
                tc=tc,
                sampling_frequency=sampling_frequency,
                minimum_frequency=minimum_frequency,
                **params,
            )

        return _call_backend

    def register_model(self, name: str, factory_func: Callable[..., Any] | str) -> None:
        """Register or overwrite a waveform model under ``name``.

        Args:
            name: Key used with ``WaveformFactory.generate`` and ``WaveformFactory.get_model``.
            factory_func: Callable that accepts merged waveform kwargs (including
                ``waveform_model``, ``tc``, ``sampling_frequency``, ``minimum_frequency``)
                and returns a dict of GWpy ``plus``/``cross`` series, **or** an import
                string: either ``module.path:callable`` (colon before the name) or
                ``package.module.callable`` (split on the last ``.`` for attribute lookup).

        Raises:
            ImportError: If a string path does not refer to an importable module.
            AttributeError: If the imported module has no such callable attribute.
            ValueError: If factory_func string is neither 'module.path:callable' nor 'package.module.callable'.
            TypeError: Registered model is not callable.
        """
        if isinstance(factory_func, str):
            if ":" in factory_func:
                module_path, func_name = factory_func.split(":", 1)
            else:
                if "." not in factory_func:
                    raise ValueError("factory_func string must be 'module.path:callable' or 'package.module.callable'")
                module_path, func_name = factory_func.rsplit(".", 1)
            module = importlib.import_module(module_path)
            factory_func = getattr(module, func_name)

        if not callable(factory_func):
            raise TypeError(f"Registered model '{name}' is not callable")

        self._models[name] = factory_func
        logger.info("Registered waveform model: %s", name)

    def get_model(self, name: str) -> Callable[..., dict[str, TimeSeries]]:
        """Look up the generator function registered for ``name``.

        Args:
            name: Registered model name (built-in approximant or custom).

        Returns:
            The callable registered for this name.

        Raises:
            ValueError: If ``name`` is not registered.
        """
        if name in self._models:
            return self._models[name]
        raise ValueError(f"Waveform model '{name}' not found. Available: {list(self._models.keys())}.")

    def list_models(self) -> list[str]:
        """Return every registered waveform model name, in dict iteration order.

        Returns:
            List of keys (backend approximants plus any custom registrations).
        """
        return list(self._models.keys())

    def generate(
        self,
        waveform_model: str,
        parameters: dict[str, Any],
        **extra_params: Any,
    ) -> dict[str, TimeSeries]:
        """Generate polarizations by calling the registered model with merged parameters.

        The callable is invoked with ``waveform_model``, then entries from ``parameters``,
        then ``extra_params`` (later keys override earlier ones).

        Args:
            waveform_model: Name of the registered model to run.
            parameters: Injection parameters (e.g. ``tc``, masses, spins) merged first.
            **extra_params: Additional fixed settings (e.g. ``sampling_frequency``,
                ``minimum_frequency``) merged after ``parameters``; later keys override.

        Returns:
            Dict whose keys are the strings ``plus`` and ``cross``, each mapping to a
            GWpy [`TimeSeries`](https://gwpy.github.io/docs/latest/api/gwpy.timeseries.TimeSeries/).

        Raises:
            ValueError: If ``waveform_model`` is not registered.
            TypeError: If the underlying generator is called with invalid arguments.
        """
        waveform_func = self.get_model(waveform_model)
        if "waveform_model" in parameters or "waveform_model" in extra_params:
            raise ValueError("Do not pass 'waveform_model' in parameters/extra_params.")
        all_params: dict[str, Any] = {**parameters, **extra_params, "waveform_model": waveform_model}
        return waveform_func(**all_params)