pysmo

@runtime_checkable
class Seismogram(Protocol):
    """Protocol class to define the `Seismogram` type.

    Examples:
        Usage for a function that takes a Seismogram compatible class instance as
        argument and returns the begin time in isoformat:

        ```python
        >>> from pysmo import Seismogram
        >>> from pysmo.classes import SAC  # SAC is a class that "speaks" Seismogram
        >>>
        >>> def example_function(seis_in: Seismogram) -> str:
        ...     return seis_in.begin_time.isoformat()
        ...
        >>> sac = SAC.from_file("example.sac")
        >>> seismogram = sac.seismogram
        >>> example_function(seismogram)
        '2010-02-27T06:44:06.069538+00:00'
        >>>
        ```
    """

    begin_time: pd.Timestamp
    """Seismogram begin time."""

    data: np.ndarray
    """Seismogram data."""

    delta: pd.Timedelta
    """The sampling interval.

    Should be a positive `pd.Timedelta` instance.
    """

    @property
    def end_time(self) -> pd.Timestamp:
        """Seismogram end time."""
        ...

@runtime_checkable
class Station(Location, Protocol):
    """Protocol class to define the `Station` type."""

    name: str
    """Station name or identifier.

    A 1-5 character identifier for the station recording the data.
    """

    network: str
    """Network name or identifier.

    A 1-2 character code identifying the network/owner of the data.
    """

    location: str
    """Location ID.

    A two character code used to uniquely identify different data streams
    at a single station.
    """

    channel: str
    """Channel code.

    A three character combination used to identify:

    1. Band and general sample rate.
    2. Instrument type.
    3. Orientation of the sensor.
    """

    elevation: int | float | None
    """Station elevation in metres."""

@runtime_checkable
class Event(LocationWithDepth, Protocol):
    """Protocol class to define the `Event` type."""

    time: pd.Timestamp
    """Event origin time."""

@runtime_checkable
class Location(Protocol):
    """Protocol class to define the `Location` type."""

    latitude: float
    """Latitude in degrees."""

    longitude: float
    """Longitude in degrees."""

@runtime_checkable
class LocationWithDepth(Location, Protocol):
    """Protocol class to define the `LocationWithDepth` type."""

    depth: float
    """Location depth in metres (positive downward from the surface)."""

@runtime_checkable
class Response(Protocol):
    """Protocol class to define the `Response` type.

    Represents an analog instrument response (Laplace domain), equivalent to
    a SAC PZ file: poles/zeros plus the total system sensitivity.
    """

    poles: list[complex]
    """Response poles, in radians/second (SAC PZ / `LAPLACE (RADIANS/SECOND)` convention)."""

    zeros: list[complex]
    """Response zeros, in radians/second."""

    overall_sensitivity: NonZeroNumber
    """Scale factor combined with `poles`/`zeros` to reconstruct the full,
    frequency-dependent transfer function `H(f)`.

    Equivalent to `CONSTANT` in a SAC PZ file (`A0 * sensitivity`, the
    analog stage's normalisation factor times the reference-frequency
    sensitivity), or FDSN StationXML's `NormalizationFactor *
    InstrumentSensitivity`. This is *not* the instrument's plain flat-band
    gain — see [`reference_sensitivity`][pysmo.Response.reference_sensitivity]
    for that — so dividing raw data by `overall_sensitivity` directly
    (rather than combining it with `poles`/`zeros`, or using
    `reference_sensitivity` instead) mis-scales the result by the `A0`
    factor, often several orders of magnitude.

    Negative values are permitted (but not zero): a negative `CONSTANT`/
    `NormalizationFactor` is how a reversed-polarity channel is recorded in
    the wild, not an error.
    """

    reference_sensitivity: NonZeroNumber | None
    """Total system sensitivity (counts per physical unit) at the response's
    own reference/normalisation frequency, with no `A0` normalisation folded
    in.

    Equivalent to SAC PZ's `SENSITIVITY` header value, or FDSN StationXML's
    `InstrumentSensitivity/Value`. This — not `overall_sensitivity`, which
    has `A0` folded in — is the correct divisor for a flat, zero-phase
    approximation of the response (e.g.
    [`remove_response`][pysmo.tools.signal.remove_response]'s
    sensitivity-only path). `None` if unavailable (e.g. a SAC PZ file
    without a `SENSITIVITY` header): callers needing it should raise rather
    than silently substituting `overall_sensitivity`. As with
    `overall_sensitivity`, a negative value indicates reversed polarity
    rather than an error; zero is not permitted.
    """

    input_units: str
    """Physical units produced by removing this response (e.g. `"M/S"`, `"M/S**2"`).

    Informational only: not validated against a fixed set of units, and not
    read by [`remove_response`][pysmo.tools.signal.remove_response] or
    [`integrate`][pysmo.tools.signal.integrate]/
    [`differentiate`][pysmo.tools.signal.differentiate] — callers are
    responsible for interpreting it themselves.
    """

@runtime_checkable
class StagedResponse(Response, Protocol):
    """Protocol class to define a `Response` with digital decimation stages.

    Extends `Response` with the digital FIR/IIR stages of the instrument's
    full signal chain, in stage (signal) order.
    """

    stages: list[ResponseStage]
    """Digital decimation stages, in signal order (stage 1 = closest to the
    analog sensor)."""

@runtime_checkable
class ResponseStage(Protocol):
    """Protocol class to define one digital (FIR/IIR) decimation stage."""

    input_sample_rate: PositiveNumber
    """Sample rate (Hz) this stage's filter coefficients operate at."""

    decimation_factor: int
    """Integer decimation factor applied by this stage."""

    numerator: list[float]
    """Feedforward ("b") filter coefficients."""

    denominator: list[float]
    """Feedback ("a") filter coefficients. `[1.0]` for a pure FIR stage."""

    correction: float
    """Time correction (seconds) already applied to the recorded data to
    cancel this stage's own filter delay.

    Equivalent to FDSN StationXML's `Decimation/Correction` (SEED Blockette
    57 field 8). Real digitisers timestamp their output as if this stage
    had zero delay, by shifting the data earlier by this amount; evaluating
    `numerator`/`denominator` alone (e.g. via `scipy.signal.freqz`) instead
    reproduces the filter's own, uncorrected delay. `0.0` (the default,
    correct for a stage with no such correction, e.g. a symmetric FIR
    stage where the coefficients themselves carry no net delay) leaves the
    coefficient-derived transfer function unchanged."""

@define(kw_only=True, slots=True)
class MiniSeismogram(SeismogramEndtimeMixin):
    """Minimal class for use with the [`Seismogram`][pysmo.Seismogram] type.

    Examples:
        ```python
        >>> from pysmo import MiniSeismogram, Seismogram
        >>> import pandas as pd
        >>> from datetime import timezone
        >>> import numpy as np
        >>> now = pd.Timestamp.now(timezone.utc)
        >>> delta = pd.Timedelta(seconds=0.1)
        >>> seismogram = MiniSeismogram(begin_time=now, delta=delta, data=np.random.rand(100))
        >>> isinstance(seismogram, Seismogram)
        True
        >>>
        ```
    """

    begin_time: UtcTimestamp = field(
        default=SeismogramDefaults.begin_time,
        converter=convert_to_utc_timestamp,
        on_setattr=setters.convert,
    )
    """Seismogram begin time."""

    delta: PositiveTimedelta = field(
        default=SeismogramDefaults.delta,
        converter=convert_to_timedelta,
        validator=[
            validators.instance_of(pd.Timedelta),
            validators.gt(pd.Timedelta(0)),
        ],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Seismogram sampling interval."""

    data: np.ndarray = field(
        factory=lambda: np.array([]),
        converter=convert_to_ndarray,
        validator=validators.instance_of(np.ndarray),
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Seismogram data."""

@define(kw_only=True, slots=True)
class MiniStation:
    """Minimal class for use with the [`Station`][pysmo.Station] type.

    Examples:
        ```python
        >>> from pysmo import MiniStation, Station, Location
        >>> station = MiniStation(latitude=-21.680301, longitude=-46.732601, name="CACB", network="BL", channel="BHZ", location="00")
        >>> isinstance(station, Station)
        True
        >>> isinstance(station, Location)
        True
        >>>
        ```
    """

    name: str = field(
        validator=[validators.min_len(1), validators.max_len(5)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Station name.

    See [`Station.name`][pysmo.Station.name] for more details.
    """

    network: str = field(
        validator=[validators.min_len(1), validators.max_len(2)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Network name.

    See [`Station.network`][pysmo.Station.network] for more details.
    """

    location: str = field(
        default="  ",
        validator=[validators.min_len(2), validators.max_len(2)],
        converter=_pad_string,
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Location ID.

    See [`Station.location`][pysmo.Station.location] for more details.
    """

    channel: str = field(
        validator=[validators.min_len(3), validators.max_len(3)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Channel code.

    See [`Station.channel`][pysmo.Station.channel] for more details.
    """

    latitude: float = field(
        converter=float,
        validator=[validators.ge(-90), validators.le(90)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Station latitude from -90 to 90 degrees."""

    longitude: float = field(
        converter=float,
        validator=[validators.gt(-180), validators.le(180)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Station longitude from -180 to 180 degrees."""

    elevation: float | None = field(
        default=None,
        converter=converters.optional(float),
        validator=validators.optional(validators.instance_of(float)),
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Station elevation."""

@define(kw_only=True, slots=True)
class MiniEvent:
    """Minimal class for use with the [`Event`][pysmo.Event] type.

    Examples:
        ```python
        >>> from pysmo import MiniEvent, Event, LocationWithDepth, Location
        >>> import pandas as pd
        >>> from datetime import timezone
        >>> now = pd.Timestamp.now(timezone.utc)
        >>> event = MiniEvent(latitude=-24.68, longitude=-26.73, depth=15234.0, time=now)
        >>> isinstance(event, Event)
        True
        >>> isinstance(event, Location)
        True
        >>> isinstance(event, LocationWithDepth)
        True
        >>>
        ```
    """

    time: UtcTimestamp = field(
        converter=convert_to_utc_timestamp,
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Event origin time."""

    latitude: float = field(
        converter=float,
        validator=[validators.ge(-90), validators.le(90)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Event latitude from -90 to 90 degrees."""

    longitude: float = field(
        converter=float,
        validator=[validators.gt(-180), validators.le(180)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Event longitude from -180 to 180 degrees."""

    depth: float = field(
        converter=float, on_setattr=setters.pipe(setters.convert, setters.validate)
    )
    """Event depth in metres (positive downward from the surface)."""

@define(kw_only=True, slots=True)
class MiniLocation:
    """Minimal class for use with the [`Location`][pysmo.Location] type.

    Examples:
        ```python
        >>> from pysmo import MiniLocation, Location
        >>> location = MiniLocation(latitude=41.8781, longitude=-87.6298)
        >>> isinstance(location, Location)
        True
        >>>
        ```
    """

    latitude: float = field(
        converter=float,
        validator=[validators.ge(-90), validators.le(90)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Latitude from -90 to 90 degrees."""

    longitude: float = field(
        converter=float,
        validator=[validators.gt(-180), validators.le(180)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Longitude from -180 to 180 degrees."""

@define(kw_only=True, slots=True)
class MiniLocationWithDepth:
    """Minimal class for use with the [`LocationWithDepth`][pysmo.LocationWithDepth] type.

    Examples:
        ```python
        >>> from pysmo import MiniLocationWithDepth, LocationWithDepth, Location
        >>> hypo = MiniLocationWithDepth(latitude=-24.68, longitude=-26.73, depth=15234.0)
        >>> isinstance(hypo, LocationWithDepth)
        True
        >>> isinstance(hypo, Location)
        True
        >>>
        ```
    """

    latitude: float = field(
        converter=float,
        validator=[validators.ge(-90), validators.le(90)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Location latitude from -90 to 90 degrees."""

    longitude: float = field(
        converter=float,
        validator=[validators.gt(-180), validators.le(180)],
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Location longitude from -180 to 180 degrees."""

    depth: float = field(
        converter=float, on_setattr=setters.pipe(setters.convert, setters.validate)
    )
    """Location depth in metres (positive downward from the surface)."""

@define(kw_only=True, slots=True)
class MiniResponse:
    """Minimal implementation of the [`Response`][pysmo.Response] type.

    Examples:
        ```python
        >>> from pysmo import MiniResponse, Response
        >>> response = MiniResponse(
        ...     poles=[-0.037 + 0.037j, -0.037 - 0.037j],
        ...     zeros=[0j, 0j],
        ...     overall_sensitivity=3.4e9,
        ...     input_units="M/S",
        ... )
        >>> isinstance(response, Response)
        True
        >>>
        ```
    """

    poles: list[complex] = field(
        converter=_convert_complex_list,
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Response poles.

    See [`Response.poles`][pysmo.Response.poles] for more details.
    """

    zeros: list[complex] = field(
        converter=_convert_complex_list,
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Response zeros.

    See [`Response.zeros`][pysmo.Response.zeros] for more details.
    """

    overall_sensitivity: NonZeroNumber = field(
        converter=float,
        validator=validate_nonzero,
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Scale factor combined with `poles`/`zeros` to reconstruct `H(f)`.

    See [`Response.overall_sensitivity`][pysmo.Response.overall_sensitivity]
    for more details.
    """

    reference_sensitivity: NonZeroNumber | None = field(
        default=None,
        converter=_convert_optional_float,
        validator=validators.optional(validate_nonzero),
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Total system sensitivity at the reference frequency, `A0` excluded.

    See
    [`Response.reference_sensitivity`][pysmo.Response.reference_sensitivity]
    for more details.
    """

    input_units: str = field(
        validator=validators.instance_of(str),
        on_setattr=setters.pipe(setters.convert, setters.validate),
    )
    """Physical units produced by removing this response.

    See [`Response.input_units`][pysmo.Response.input_units] for more details.
    """