Time series

@dataclass(frozen=True, eq=False)
class TimeSeries:
    """A uniformly sampled series: sample ``k`` along the last axis is at ``t0 + k / sample_rate``.

    Args:
        data: The samples, an array with at least one axis; time is the last. A NumPy or a JAX
            array is kept as given. Anything else is converted with :func:`numpy.asarray`.
        t0: GPS time of the first sample, in seconds.
        sample_rate: Samples per second, in Hz.
        name: Optional label, such as a detector or channel name. It labels the whole series.
        unit: Optional unit of the samples, such as ``"strain"``.

    Raises:
        ValueError: If ``data`` has no axis, or ``t0`` or ``sample_rate`` is not finite, or
            ``sample_rate`` is not positive.
        TypeError: If ``t0`` is not a single number.

    Two series are equal only if they are the same object: comparing arrays elementwise has no
    single truth value, and a JAX array cannot be compared without a device round trip.
    """

    data: Any
    t0: float
    sample_rate: float
    name: str | None = None
    unit: str | None = None

    def __post_init__(self) -> None:
        """Validate the timing and normalise the samples to an array."""
        data = self.data
        if not hasattr(data, "__array_namespace__"):
            data = np.asarray(data)
            object.__setattr__(self, "data", data)
        if len(data.shape) == 0:
            raise ValueError("data must have at least one axis; time is the last.")
        sample_rate = float(self.sample_rate)
        if not (math.isfinite(sample_rate) and sample_rate > 0.0):
            raise ValueError(f"sample_rate must be positive and finite; got {self.sample_rate!r}.")
        object.__setattr__(self, "sample_rate", sample_rate)
        try:
            t0 = float(self.t0)
        except TypeError as exc:
            # A traced value inside a JAX transformation has no number yet to check; it was checked
            # when the series entered the transformation. Matched by name so JAX is not imported.
            if type(exc).__name__ == "ConcretizationTypeError":
                return
            raise TypeError(f"t0 must be a single number; got {self.t0!r}.") from exc
        if not math.isfinite(t0):
            raise ValueError(f"t0 must be finite; got {self.t0!r}.")
        object.__setattr__(self, "t0", t0)

    # ----------------------------------------------------------------------------------------
    # Derived quantities

    @property
    def value(self) -> Any:
        """The samples; an alias of :attr:`data`, the name GWpy uses."""
        return self.data

    @property
    def xp(self) -> ModuleType:
        """The array namespace of :attr:`data`, such as ``numpy`` or ``jax.numpy``."""
        return _namespace(self.data)

    @property
    def n_samples(self) -> int:
        """Number of samples along the time axis."""
        return int(self.data.shape[-1])

    @property
    def dt(self) -> float:
        """Seconds between samples."""
        return 1.0 / self.sample_rate

    @property
    def duration(self) -> float:
        """Seconds spanned by the samples: ``n_samples * dt``."""
        return self.n_samples * self.dt

    @property
    def end(self) -> float:
        """GPS time of the sample after the last, on the series' lattice: ``t0 + n_samples * dt``."""
        return self._time_of(self.n_samples)

    @property
    def grid(self) -> SamplingGrid:
        """The sample lattice of this series."""
        return SamplingGrid(epoch=self.t0, sampling_frequency=self.sample_rate)

    @property
    def times(self) -> Any:
        """GPS time of each sample, in the namespace of :attr:`data`, as float64.

        Raises:
            RuntimeError: If the namespace cannot hold float64, as JAX cannot outside its 64-bit
                mode. In single precision a GPS time near 1.4e9 s resolves only to 128 s, so every
                sample of a typical series would get the same time.
        """
        xp = self.xp
        times = self.t0 + xp.arange(self.n_samples, dtype=xp.float64) * self.dt
        if times.dtype != xp.float64:
            raise RuntimeError(_NO_FLOAT64)
        return times

    # ----------------------------------------------------------------------------------------
    # New series from this one

    def with_data(self, data: Any) -> TimeSeries:
        """Return a series with the same timing, name and unit, holding ``data``."""
        return dataclasses.replace(self, data=data)

    def copy(self) -> TimeSeries:
        """Return a series whose samples are a copy of these."""
        return self.with_data(self.xp.asarray(self.data, copy=True))

    def crop(self, start: float | None = None, end: float | None = None) -> TimeSeries:
        """Return the samples at GPS times in ``[start, end)``.

        The bounds decide the result's length, so ``crop`` needs a concrete ``t0``: it cannot run
        inside a JAX transformation over the series.

        The result starts at the first sample at or after ``start`` and stops before the first at
        or after ``end``. A bound that falls on a sample, to within the float64 resolution of GPS
        times (:meth:`~gwmock_data.sampling_grid.SamplingGrid.lattice_tolerance_samples`), counts
        as on it, so cropping at a sample's own timestamp keeps that sample however large the GPS
        time. Bounds outside the series are clipped to it.

        GWpy differs for a bound between samples: it rounds down, keeping the sample before
        ``start``. Here every kept sample is at or after ``start``.

        Args:
            start: First GPS time to keep; ``None`` keeps from the first sample.
            end: GPS time to stop before; ``None`` keeps to the last sample.

        Returns:
            The cropped series, possibly with no samples. Its samples are a view where the array
            library makes one.

        Raises:
            ValueError: If a bound is not finite, or ``end`` is before ``start``.
        """
        for bound_name, bound in (("start", start), ("end", end)):
            if bound is not None and not math.isfinite(bound):
                raise ValueError(f"{bound_name} must be finite; got {bound!r}.")
        if start is not None and end is not None and end < start:
            raise ValueError(f"end ({end}) is before start ({start}).")
        first = 0 if start is None else self._first_index_at_or_after(start)
        stop = self.n_samples if end is None else self._first_index_at_or_after(end)
        first = min(max(first, 0), self.n_samples)
        stop = min(max(stop, first), self.n_samples)
        return dataclasses.replace(self, data=self.data[..., first:stop], t0=self._time_of(first))

    def pad(self, before: int = 0, after: int = 0) -> TimeSeries:
        """Return the series with ``before`` zeros prepended and ``after`` appended, in samples.

        ``t0`` moves back by ``before`` samples, so the existing samples keep their times. The
        counts decide the result's length, so inside a JAX transformation they must be Python
        integers, not traced values.

        Raises:
            ValueError: If either count is negative.
        """
        if before < 0 or after < 0:
            raise ValueError(f"pad counts must not be negative; got before={before}, after={after}.")
        xp = self.xp
        lead = self.data.shape[:-1]
        parts = [
            xp.zeros((*lead, before), dtype=self.data.dtype),
            self.data,
            xp.zeros((*lead, after), dtype=self.data.dtype),
        ]
        return dataclasses.replace(self, data=xp.concatenate(parts, axis=-1), t0=self._time_of(-before))

    def is_compatible(self, other: TimeSeries) -> bool:
        """Return whether ``other`` has the same sample rate and unit, so the two can be combined."""
        return self.sample_rate == other.sample_rate and self.unit == other.unit

    def _time_of(self, index: int) -> Any:
        # The arithmetic of SamplingGrid.time_of, written out so it also works on a traced t0.
        return self.t0 + index * self.dt

    def _first_index_at_or_after(self, gps_time: float) -> int:
        grid = self.grid
        exact = float(grid.index_of(gps_time))
        return math.ceil(exact - grid.lattice_tolerance_samples(gps_time))

    def __array__(self, dtype: Any = None, copy: bool | None = None) -> np.ndarray:
        """Return the samples as a NumPy array, so ``numpy.asarray(series)`` works as for GWpy.

        Follows NumPy's protocol: ``copy=True`` always copies, ``None`` copies only when the dtype
        or the array library requires it, and ``False`` raises ``ValueError`` when a copy would be
        required -- a dtype change, or samples that are not a NumPy array.
        """
        needs_copy = not isinstance(self.data, np.ndarray) or (dtype is not None and np.dtype(dtype) != self.data.dtype)
        if copy is False and needs_copy:
            raise ValueError("these samples cannot be returned as a NumPy array without a copy.")
        array = np.asarray(self.data, dtype=dtype)
        return np.array(array, copy=True) if copy else array

    # ----------------------------------------------------------------------------------------
    # Moving data and converting

    def to_numpy(self) -> TimeSeries:
        """Return the series with its samples as a NumPy array on the host."""
        return self.with_data(np.asarray(self.data))

    def to_jax(self, device: Any = None) -> TimeSeries:
        """Return the series with its samples as a JAX array, on ``device`` if given.

        Raises:
            ModuleNotFoundError: If JAX is not installed.
            RuntimeError: If JAX is not in 64-bit mode, in which it would round double-precision
                samples to single precision.
        """
        jax = require("jax", "jax", "TimeSeries.to_jax")
        _require_x64(jax)
        return self.with_data(jax.device_put(jax.numpy.asarray(self.data), device))

    def to_gwpy(self) -> gwpy.timeseries.TimeSeries:
        """Return a GWpy ``TimeSeries`` with the same samples, timing, name and unit.

        NumPy samples are shared, not copied, as :meth:`from_gwpy` shares GWpy's, so a write through
        either object is seen by the other. JAX samples are brought to the host with
        :func:`numpy.asarray`; for a CPU array that is a read-only view of JAX's buffer.

        Raises:
            ValueError: If the series has more than one axis; GWpy's is one-dimensional.
            ModuleNotFoundError: If GWpy is not installed.
        """
        if len(self.data.shape) != 1:
            raise ValueError(f"a GWpy TimeSeries is one-dimensional; this series has shape {tuple(self.data.shape)}.")
        timeseries = require("gwpy.timeseries", "gwpy", "TimeSeries.to_gwpy")
        return timeseries.TimeSeries(
            np.asarray(self.data),
            t0=self.t0,
            sample_rate=self.sample_rate,
            name=self.name,
            unit=self.unit,
            copy=False,
        )

    @classmethod
    def from_gwpy(cls, series: gwpy.timeseries.TimeSeries) -> TimeSeries:
        """Return a series with the samples, timing, name and unit of a GWpy ``TimeSeries``.

        The samples are GWpy's array itself, not a copy. A dimensionless GWpy series gets
        ``unit=None``.

        Raises:
            ValueError: If the GWpy series is irregularly sampled.
        """
        try:
            sample_rate = series.sample_rate.to_value("Hz")
        except AttributeError as exc:
            # GWpy has no sample rate for a series built from an explicit, irregular times array.
            raise ValueError(
                "a TimeSeries is uniformly sampled; this GWpy series has an irregular time index."
            ) from exc
        unit = series.unit.to_string() if series.unit is not None else ""
        return cls(
            series.value,
            t0=series.t0.to_value("s"),
            sample_rate=sample_rate,
            name=None if series.name is None else str(series.name),
            unit=unit or None,
        )

    def to_pycbc(self) -> pycbc.types.TimeSeries:
        """Return a PyCBC ``TimeSeries`` with the same samples and timing.

        PyCBC series carry neither a name nor a unit.

        Raises:
            ValueError: If the series has more than one axis; PyCBC's is one-dimensional.
            ModuleNotFoundError: If PyCBC is not installed.
        """
        if len(self.data.shape) != 1:
            raise ValueError(f"a PyCBC TimeSeries is one-dimensional; this series has shape {tuple(self.data.shape)}.")
        types = require("pycbc.types", "pycbc", "TimeSeries.to_pycbc")
        return types.TimeSeries(np.asarray(self.data), delta_t=self.dt, epoch=self.t0)

    @classmethod
    def from_pycbc(
        cls, series: pycbc.types.TimeSeries, *, name: str | None = None, unit: str | None = None
    ) -> TimeSeries:
        """Return a series with the samples and timing of a PyCBC ``TimeSeries``."""
        return cls(
            np.asarray(series.numpy()),
            t0=float(series.start_time),
            sample_rate=float(series.sample_rate),
            name=name,
            unit=unit,
        )