Waveform

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

    Args:
        f_ref: Reference frequency in Hz. Defaults to ``minimum_frequency`` of each
            call when ``None``.
        ringdown_fraction: Fraction of the analysis segment reserved after
            coalescence. Must be in ``(0, 1)``.
        segment_duration: Optional fixed analysis-segment length in seconds. When
            ``None`` (default) the length is estimated from the post-Newtonian chirp
            time so the full inspiral fits without wraparound.
    """

    def __init__(
        self,
        *,
        f_ref: float | None = None,
        ringdown_fraction: float = conditioning.DEFAULT_RINGDOWN_FRACTION,
        segment_duration: float | None = None,
    ) -> None:
        """Validate the placement configuration shared with the ripple backend."""
        if not 0.0 < ringdown_fraction < 1.0:
            raise ValueError("ringdown_fraction must be in (0, 1)")
        if segment_duration is not None and segment_duration <= 0:
            raise ValueError("segment_duration must be > 0")
        self._f_ref = f_ref
        self._ringdown_fraction = ringdown_fraction
        self._segment_duration = segment_duration

    def available_approximants(self) -> list[str]:
        """Return every LAL approximant this backend can generate.

        ``generate_td_waveform`` produces FD-native approximants via
        ``SimInspiralChooseFDWaveform`` and time-domain approximants via
        ``SimInspiralFD``, so the advertised set is the union of both. Iterating
        over approximant indices yields each name once, in a stable order.
        """
        return [
            lalsimulation.GetStringFromApproximant(i)
            for i in range(lalsimulation.NumApproximants)
            if lalsimulation.SimInspiralImplementedTDApproximants(i)
            or lalsimulation.SimInspiralImplementedFDApproximants(i)
        ]

    @staticmethod
    def _resolve_parameters(
        sampling_frequency: float, minimum_frequency: float, **params: object
    ) -> _ResolvedParameters:
        """Validate inputs and translate canonical parameters to backend-native ones."""
        remaining = dict(params)
        resolved = _ResolvedParameters(
            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)),
            lambda_1=float(_pop_alias(remaining, "lambda_1", "tidal_1", default=0.0)),
            lambda_2=float(_pop_alias(remaining, "lambda_2", "tidal_2", 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")
        if minimum_frequency <= 0:
            raise ValueError("minimum_frequency must be > 0")
        if resolved.lambda_1 < 0:
            raise ValueError("lambda_1 must be >= 0")
        if resolved.lambda_2 < 0:
            raise ValueError("lambda_2 must be >= 0")
        return resolved

    def generate_td_waveform(
        self,
        approximant: str,
        tc: float,
        sampling_frequency: float,
        minimum_frequency: float,
        **params: object,
    ) -> dict[str, TimeSeries]:
        """Generate plus/cross polarizations, conditioned from frequency to time domain."""
        p = self._resolve_parameters(sampling_frequency, minimum_frequency, **params)

        chirp_mass = (p.mass1 * p.mass2) ** 0.6 / (p.mass1 + p.mass2) ** 0.2
        n_samples = conditioning.segment_sample_count(
            chirp_mass,
            minimum_frequency,
            sampling_frequency,
            ringdown_fraction=self._ringdown_fraction,
            segment_duration=self._segment_duration,
        )
        delta_f = sampling_frequency / n_samples
        f_max = sampling_frequency / 2.0
        f_ref = self._f_ref if self._f_ref is not None else minimum_frequency

        approx_enum = lalsimulation.GetApproximantFromString(approximant)
        lal_params = lal.CreateDict()
        lalsimulation.SimInspiralWaveformParamsInsertTidalLambda1(lal_params, p.lambda_1)
        lalsimulation.SimInspiralWaveformParamsInsertTidalLambda2(lal_params, p.lambda_2)

        wf_args = (
            p.mass1 * MSUN,
            p.mass2 * MSUN,
            p.spin_1x,
            p.spin_1y,
            p.spin_1z,
            p.spin_2x,
            p.spin_2y,
            p.spin_2z,
            p.distance * MPC,
            p.inclination,
            p.coa_phase,
            0.0,  # longitude of ascending nodes
            0.0,  # eccentricity
            0.0,  # mean periastron anomaly
            delta_f,
            minimum_frequency,
            f_max,
            f_ref,
            lal_params,
            approx_enum,
        )
        # Follow bilby: FD-native approximants come back already referenced to the FD
        # phase; a TD approximant routed through SimInspiralFD carries an epoch from
        # the internal time-domain conditioning, undone by a dt = T + epoch shift.
        if lalsimulation.SimInspiralImplementedFDApproximants(approx_enum):
            hp, hc = lalsimulation.SimInspiralChooseFDWaveform(*wf_args)
            epoch_shift = 0.0
        else:
            hp, hc = lalsimulation.SimInspiralFD(*wf_args)
            epoch_shift = 1.0 / hp.deltaF + (hp.epoch.gpsSeconds + hp.epoch.gpsNanoSeconds * 1e-9)

        n_freq = n_samples // 2 + 1
        freqs = np.arange(n_freq) * delta_f
        hp_f = _to_onesided(hp.data.data, n_freq)
        hc_f = _to_onesided(hc.data.data, n_freq)
        if epoch_shift:
            time_shift = np.exp(-2j * np.pi * freqs * epoch_shift)
            hp_f = hp_f * time_shift
            hc_f = hc_f * time_shift
        in_band = freqs >= minimum_frequency
        hp_f = np.nan_to_num(np.where(in_band, hp_f, 0.0))
        hc_f = np.nan_to_num(np.where(in_band, hc_f, 0.0))

        hp_t, hc_t, epoch = conditioning.condition_fd_to_td(
            hp_f, hc_f, n_samples, sampling_frequency, self._ringdown_fraction
        )
        dt = 1.0 / sampling_frequency
        t0 = epoch + tc
        return {
            "plus": TimeSeries(hp_t, t0=t0, dt=dt),
            "cross": TimeSeries(hc_t, 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),
            "lambda1": _pop_alias(remaining, "lambda_1", "lambda1", "tidal_1", default=0.0),
            "lambda2": _pop_alias(remaining, "lambda_2", "lambda2", "tidal_2", 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 RippleBackend(WaveformBackend):
    """Time-domain waveform backend implemented with ripple (JAX).

    Args:
        f_ref: Reference frequency in Hz. Defaults to ``minimum_frequency`` of each
            call when ``None``.
        ringdown_fraction: Fraction of the analysis segment reserved after
            coalescence. Must be in ``(0, 1)``.
        segment_duration: Optional fixed analysis-segment length in seconds. When
            ``None`` (default) the length is estimated from the post-Newtonian
            chirp time so the full inspiral fits without wraparound.
    """

    def __init__(
        self,
        *,
        f_ref: float | None = None,
        ringdown_fraction: float = _DEFAULT_RINGDOWN_FRACTION,
        segment_duration: float | None = None,
    ) -> None:
        """Require ripple/JAX only when this backend is instantiated."""
        try:
            self._jax = importlib.import_module("jax")
            self._jnp = importlib.import_module("jax.numpy")
            self._ripplegw = importlib.import_module("ripplegw")
            self._conversions = importlib.import_module("ripplegw.conversions")
            self._constants = importlib.import_module("ripplegw.constants")
        except ImportError as exc:
            raise ImportError(_RIPPLE_IMPORT_ERROR) from exc
        # ripple needs double precision for waveform phase accuracy over long
        # inspirals. Importing ripplegw already enables this globally; set it
        # explicitly so correctness does not depend on import order.
        self._jax.config.update("jax_enable_x64", True)
        if not 0.0 < ringdown_fraction < 1.0:
            raise ValueError("ringdown_fraction must be in (0, 1)")
        if segment_duration is not None and segment_duration <= 0:
            raise ValueError("segment_duration must be > 0")
        self._f_ref = f_ref
        self._ringdown_fraction = ringdown_fraction
        self._segment_duration = segment_duration

    def available_approximants(self) -> list[str]:
        """Return the ripple approximants supported by this backend."""
        return list(_SUPPORTED_APPROXIMANTS)

    @property
    def segment_duration(self) -> float | None:
        """The fixed analysis-segment length in seconds, or ``None`` if auto-sized."""
        return self._segment_duration

    def with_segment_duration(self, segment_duration: float) -> RippleBackend:
        """Return a copy of this backend pinned to a fixed ``segment_duration``.

        Same ``f_ref`` and ``ringdown_fraction``; useful for forcing one shared grid
        across several batched calls (e.g. count-chunked catalogue generation).
        """
        return RippleBackend(
            f_ref=self._f_ref,
            ringdown_fraction=self._ringdown_fraction,
            segment_duration=segment_duration,
        )

    def segment_duration_for(
        self, chirp_mass_solar: float, minimum_frequency: float, sampling_frequency: float
    ) -> float:
        """Worst-case segment duration (seconds) the batch path uses for this chirp mass."""
        return self._segment_samples(chirp_mass_solar, minimum_frequency, sampling_frequency) / sampling_frequency

    def generate_td_waveform(
        self,
        approximant: str,
        tc: float,
        sampling_frequency: float,
        minimum_frequency: float,
        **params: object,
    ) -> dict[str, TimeSeries]:
        """Generate plus/cross polarizations from ripple, conditioned to time domain."""
        fd = self.generate_fd_polarizations(
            approximant,
            sampling_frequency=sampling_frequency,
            minimum_frequency=minimum_frequency,
            **params,
        )
        hp_t, hc_t, epoch = self._to_time_domain(fd)
        t0 = epoch + tc
        dt = 1.0 / sampling_frequency
        return {
            "plus": TimeSeries(hp_t, t0=t0, dt=dt),
            "cross": TimeSeries(hc_t, t0=t0, dt=dt),
        }

    def generate_fd_polarizations(
        self,
        approximant: str,
        *,
        sampling_frequency: float,
        minimum_frequency: float,
        **params: object,
    ) -> FrequencyDomainPolarizations:
        """Generate ripple's frequency-domain plus/cross polarizations (on-device).

        This is the building block the on-device (GPU) projection path consumes: the
        polarizations stay as JAX arrays and are not conditioned to the time domain.
        ``generate_td_waveform`` calls this and then inverse-FFTs the result.

        Args:
            approximant: A supported ripple approximant name.
            sampling_frequency: Sample rate in Hz; sets the Nyquist frequency.
            minimum_frequency: Low-frequency cutoff in Hz; bins below it are zeroed.
            **params: CBC source parameters (gwmock-pop canonical names or aliases).

        Returns:
            A :class:`FrequencyDomainPolarizations` with coalescence at ``t = 0``.
        """
        resolved = self._resolve_parameters(approximant, sampling_frequency, minimum_frequency, **params)
        return self._evaluate_fd(approximant, resolved, sampling_frequency, minimum_frequency)

    def generate_fd_polarizations_batch(
        self,
        approximant: str,
        *,
        sampling_frequency: float,
        minimum_frequency: float,
        parameters: Mapping[str, object],
    ) -> FrequencyDomainPolarizations:
        """Generate ripple FD polarizations for a batch of events on one shared grid.

        Evaluates ripple under ``jax.vmap`` over the catalogue, so all events share a
        single frequency grid. Because ``vmap`` needs a fixed shape, the grid is sized
        (worst case) for the longest inspiral in the batch — the smallest chirp mass,
        via the same post-Newtonian estimate as the per-event path — unless a fixed
        ``segment_duration`` was set on the backend. This is the on-device entry point
        for catalogue-scale (GPU) generation.

        Args:
            approximant: A supported ripple approximant name.
            sampling_frequency: Sample rate in Hz.
            minimum_frequency: Low-frequency cutoff in Hz; bins below it are zeroed.
            parameters: Mapping of **canonical** gwmock-pop parameter names (no aliases)
                to 1-D arrays of equal length ``n_events`` (e.g. ``detector_frame_mass_1``,
                ``spin_1z``, ``inclination``). Omitted optional parameters default to zero.

        Returns:
            A :class:`FrequencyDomainPolarizations` whose ``plus`` and ``cross`` are
            ``(n_events, n_samples // 2 + 1)`` JAX arrays (coalescence at ``t = 0``).
        """
        ripple_params, n_samples = self._resolve_batch(approximant, sampling_frequency, minimum_frequency, parameters)
        jnp = self._jnp
        delta_f = sampling_frequency / n_samples
        freqs = jnp.arange(n_samples // 2 + 1) * delta_f
        in_band = freqs >= minimum_frequency
        f_ref = self._f_ref if self._f_ref is not None else minimum_frequency
        waveform = self._ripplegw.waveform_preset[approximant](f_ref=f_ref)

        def _one(event: dict) -> tuple:
            polarizations = waveform(freqs, event)
            return (
                jnp.nan_to_num(jnp.where(in_band, polarizations["p"], 0.0)),
                jnp.nan_to_num(jnp.where(in_band, polarizations["c"], 0.0)),
            )

        # jit-compile the vmapped evaluation so the whole batch fuses into one
        # kernel (the main GPU win); compiled once per (n_events, n_samples) shape.
        plus, cross = self._jax.jit(self._jax.vmap(_one))(ripple_params)
        return FrequencyDomainPolarizations(
            frequencies=freqs,
            plus=plus,
            cross=cross,
            sampling_frequency=sampling_frequency,
            n_samples=n_samples,
        )

    def _resolve_batch(
        self,
        approximant: str,
        sampling_frequency: float,
        minimum_frequency: float,
        parameters: Mapping[str, object],
    ) -> tuple[dict, int]:
        """Validate a batch of canonical parameters and build ripple-native arrays.

        Returns ``(ripple_params, n_samples)`` where ``ripple_params`` is a dict of
        equal-length JAX arrays ready for ``vmap`` and ``n_samples`` is the shared,
        worst-case segment length.
        """
        if approximant not in _SUPPORTED_APPROXIMANTS:
            raise ValueError(
                f"RippleBackend does not support approximant {approximant!r}. "
                f"Available: {list(_SUPPORTED_APPROXIMANTS)}."
            )
        if sampling_frequency <= 0:
            raise ValueError("sampling_frequency must be > 0")
        if minimum_frequency <= 0:
            raise ValueError("minimum_frequency must be > 0")

        jnp = self._jnp
        mass1 = self._batch_array(parameters, "detector_frame_mass_1")
        n_events = mass1.shape[0]
        mass2 = self._batch_array(parameters, "detector_frame_mass_2", n_events)
        distance = self._batch_array(parameters, "luminosity_distance", n_events)
        inclination = self._batch_array(parameters, "inclination", n_events, default=0.0)
        coa_phase = self._batch_array(parameters, "coa_phase", n_events, default=0.0)
        spins = {
            name: self._batch_array(parameters, name, n_events, default=0.0)
            for name in ("spin_1x", "spin_1y", "spin_1z", "spin_2x", "spin_2y", "spin_2z")
        }
        lambda_1 = self._batch_array(parameters, "lambda_1", n_events, default=0.0)
        lambda_2 = self._batch_array(parameters, "lambda_2", n_events, default=0.0)

        is_precessing = approximant in _PRECESSING_MODELS
        if not is_precessing:
            for name in ("spin_1x", "spin_1y", "spin_2x", "spin_2y"):
                if bool(jnp.any(spins[name] != 0.0)):
                    raise ValueError(f"{approximant} is an aligned-spin model; {name} must be zero for all events.")
        is_tidal = approximant in _TIDAL_MODELS
        if not is_tidal and (bool(jnp.any(lambda_1 != 0.0)) or bool(jnp.any(lambda_2 != 0.0))):
            raise ValueError(f"{approximant} does not support tidal parameters; use an NRTidal approximant.")
        if bool(jnp.any(lambda_1 < 0.0)) or bool(jnp.any(lambda_2 < 0.0)):
            raise ValueError("lambda_1 and lambda_2 must be >= 0")

        chirp_mass, eta = self._jax.vmap(self._conversions.ms_to_Mc_eta)(jnp.stack([mass1, mass2], axis=-1))
        n_samples = self._segment_samples(float(jnp.min(chirp_mass)), minimum_frequency, sampling_frequency)

        ripple_params = {
            "M_c": chirp_mass,
            "eta": eta,
            "s1_z": spins["spin_1z"],
            "s2_z": spins["spin_2z"],
            "d_L": distance,
            "phase_c": coa_phase,
            "iota": inclination,
        }
        if is_precessing:
            ripple_params["s1_x"] = spins["spin_1x"]
            ripple_params["s1_y"] = spins["spin_1y"]
            ripple_params["s2_x"] = spins["spin_2x"]
            ripple_params["s2_y"] = spins["spin_2y"]
        if is_tidal:
            ripple_params["lambda_1"] = lambda_1
            ripple_params["lambda_2"] = lambda_2
        return ripple_params, n_samples

    def _batch_array(
        self,
        parameters: Mapping[str, object],
        name: str,
        n_events: int | None = None,
        *,
        default: float | None = None,
    ) -> Array:
        """Return one parameter as a 1-D float64 JAX array, validating its length."""
        jnp = self._jnp
        if name not in parameters:
            if default is None:
                raise ValueError(f"Missing required batch parameter: {name!r}")
            return jnp.full(n_events, default, dtype=jnp.float64)
        values = jnp.asarray(parameters[name], dtype=jnp.float64)
        if values.ndim != 1:
            raise ValueError(f"Batch parameter {name!r} must be 1-D; got shape {values.shape}.")
        if n_events is not None and values.shape[0] != n_events:
            raise ValueError(f"Batch parameter {name!r} has length {values.shape[0]}, expected {n_events}.")
        return values

    def _resolve_parameters(
        self,
        approximant: str,
        sampling_frequency: float,
        minimum_frequency: float,
        **params: object,
    ) -> _ResolvedParameters:
        """Validate inputs and translate canonical parameters to backend-native ones."""
        if approximant not in _SUPPORTED_APPROXIMANTS:
            raise ValueError(
                f"RippleBackend does not support approximant {approximant!r}. "
                f"Available: {list(_SUPPORTED_APPROXIMANTS)}."
            )
        if sampling_frequency <= 0:
            raise ValueError("sampling_frequency must be > 0")
        if minimum_frequency <= 0:
            raise ValueError("minimum_frequency must be > 0")

        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"))
        spins = {
            "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))

        is_precessing = approximant in _PRECESSING_MODELS
        if not is_precessing:
            in_plane = ("spin_1x", "spin_1y", "spin_2x", "spin_2y")
            nonzero_in_plane = sorted(name for name in in_plane if spins[name] != 0.0)
            if nonzero_in_plane:
                raise ValueError(
                    f"{approximant} is an aligned-spin model; "
                    f"in-plane spins must be zero: {', '.join(nonzero_in_plane)}"
                )
        lambda_1 = float(_pop_alias(remaining, "lambda_1", "tidal_1", default=0.0))
        lambda_2 = float(_pop_alias(remaining, "lambda_2", "tidal_2", default=0.0))
        is_tidal = approximant in _TIDAL_MODELS
        if not is_tidal and (lambda_1 or lambda_2):
            raise ValueError(f"{approximant} does not support tidal parameters; use an NRTidal approximant.")
        if lambda_1 < 0:
            raise ValueError("lambda_1 must be >= 0")
        if lambda_2 < 0:
            raise ValueError("lambda_2 must be >= 0")
        if remaining:
            extras = ", ".join(sorted(remaining))
            raise ValueError(f"Unsupported ripple waveform parameters: {extras}")

        return _ResolvedParameters(
            mass1=mass1,
            mass2=mass2,
            spins=spins,
            distance=distance,
            inclination=inclination,
            coa_phase=coa_phase,
            lambda_1=lambda_1,
            lambda_2=lambda_2,
            is_tidal=is_tidal,
            is_precessing=is_precessing,
            f_ref=self._f_ref if self._f_ref is not None else minimum_frequency,
        )

    def _segment_samples(self, chirp_mass_solar: float, minimum_frequency: float, sampling_frequency: float) -> int:
        """Return an even sample count whose duration contains the full inspiral.

        The duration is rounded up to a power of two seconds so the inspiral
        (estimated from the leading-order post-Newtonian chirp time) fits in the pre-coalescence
        portion of the buffer without cyclic wraparound.
        """
        if self._segment_duration is not None:
            seconds = self._segment_duration
        else:
            mc_seconds = chirp_mass_solar * self._constants.MTSUN
            # Leading-order (Newtonian) chirp time from minimum_frequency to merger.
            tau0 = (5.0 / 256.0) * (np.pi * minimum_frequency) ** (-8.0 / 3.0) * mc_seconds ** (-5.0 / 3.0)
            inspiral_room = 1.0 - self._ringdown_fraction
            seconds = max((tau0 + _SEGMENT_BUFFER_SECONDS) / inspiral_room, _MIN_SEGMENT_SECONDS)
        seconds_pow2 = float(2.0 ** np.ceil(np.log2(seconds)))
        n_samples = round(seconds_pow2 * sampling_frequency)
        if n_samples % 2:
            n_samples += 1
        return n_samples

    def _evaluate_fd(
        self,
        approximant: str,
        resolved: _ResolvedParameters,
        sampling_frequency: float,
        minimum_frequency: float,
    ) -> FrequencyDomainPolarizations:
        """Evaluate ripple on the analysis frequency grid (coalescence at t=0)."""
        jnp = self._jnp
        spins = resolved.spins
        chirp_mass, eta = self._conversions.ms_to_Mc_eta(jnp.array([resolved.mass1, resolved.mass2]))

        n_samples = self._segment_samples(float(chirp_mass), minimum_frequency, sampling_frequency)
        delta_f = sampling_frequency / n_samples
        freqs = jnp.arange(n_samples // 2 + 1) * delta_f

        # ripple's class interface fixes its internal tc=0; coalescence is placed
        # in the time grid by _to_time_domain.
        ripple_params = {
            "M_c": chirp_mass,
            "eta": eta,
            "s1_z": spins["spin_1z"],
            "s2_z": spins["spin_2z"],
            "d_L": resolved.distance,
            "phase_c": resolved.coa_phase,
            "iota": resolved.inclination,
        }
        if resolved.is_precessing:
            ripple_params["s1_x"] = spins["spin_1x"]
            ripple_params["s1_y"] = spins["spin_1y"]
            ripple_params["s2_x"] = spins["spin_2x"]
            ripple_params["s2_y"] = spins["spin_2y"]
        if resolved.is_tidal:
            ripple_params["lambda_1"] = resolved.lambda_1
            ripple_params["lambda_2"] = resolved.lambda_2
        waveform = self._ripplegw.waveform_preset[approximant](f_ref=resolved.f_ref)
        polarizations = waveform(freqs, ripple_params)

        # Zero out-of-band bins (including DC, where the amplitude diverges) and
        # guard against any non-finite values, keeping everything on device.
        in_band = freqs >= minimum_frequency
        hp_f = jnp.nan_to_num(jnp.where(in_band, polarizations["p"], 0.0))
        hc_f = jnp.nan_to_num(jnp.where(in_band, polarizations["c"], 0.0))
        return FrequencyDomainPolarizations(
            frequencies=freqs,
            plus=hp_f,
            cross=hc_f,
            sampling_frequency=sampling_frequency,
            n_samples=n_samples,
        )

    def coalescence_placement(self, n_samples: int, sampling_frequency: float) -> tuple[int, float]:
        """Return ``(merger_index, epoch)`` for placing coalescence in a segment.

        ``merger_index`` is the sample at which coalescence sits after the
        time-domain roll (near the segment end, leaving a small ringdown pad), and
        ``epoch`` is the time of the first sample relative to coalescence (negative),
        so a caller places coalescence at ``epoch + tc``. Shared by the time-domain
        backend and the batched device path so both use the same convention.
        """
        merger_index = round((1.0 - self._ringdown_fraction) * n_samples)
        return merger_index, -merger_index / sampling_frequency

    def _to_time_domain(self, fd: FrequencyDomainPolarizations) -> tuple[np.ndarray, np.ndarray, float]:
        """Inverse-FFT frequency-domain polarizations and place coalescence in the segment.

        Returns ``(hp, hc, epoch)`` where ``epoch`` is the time of the first sample
        relative to coalescence (negative), so the caller places coalescence at
        ``epoch + tc``.
        """
        dt = 1.0 / fd.sampling_frequency
        # Inverse real FFT: h(t) = irfft(h(f)) / dt (continuous-transform normalization).
        hp_t = np.fft.irfft(np.asarray(fd.plus), n=fd.n_samples) / dt
        hc_t = np.fft.irfft(np.asarray(fd.cross), n=fd.n_samples) / dt

        # With tc=0 coalescence lands at sample 0 and the inspiral wraps to the tail.
        # Roll it forward so coalescence sits near the segment end, leaving the
        # inspiral contiguous before it and a small ringdown pad after.
        merger_index, epoch = self.coalescence_placement(fd.n_samples, fd.sampling_frequency)
        hp_t = np.roll(hp_t, merger_index)
        hc_t = np.roll(hc_t, merger_index)
        return hp_t, hc_t, epoch

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)