pysmo.tools.project

@define(kw_only=True, frozen=True)
class FetchContext:
    """Context handed to `transform` alongside the freshly downloaded seismogram.

    Bundles the originating [`ProjectEntry`][pysmo.tools.project.ProjectEntry]
    with what this specific fetch resolved but that doesn't belong on
    `ProjectEntry` itself — recomputed fresh on every fetch, never
    persisted (unlike `entry.checksum`, which is deliberately pinned).

    Note the deliberate naming overlap with `entry.starttime`/`entry.endtime`:
    those are the entry's possibly-`None` *explicit override* (see
    [`ProjectEntry`][pysmo.tools.project.ProjectEntry]), while
    `starttime`/`endtime` here are always-populated and reflect the window
    that was *actually used* — identical to the entry's own when an explicit
    override was given, resolved from `predicted` otherwise. A `transform`
    wanting "the window this fetch actually covered" should read
    `context.starttime`/`context.endtime`, not
    `context.entry.starttime`/`context.entry.endtime`.
    """

    entry: ProjectEntry
    """The entry this seismogram was fetched for."""

    starttime: pd.Timestamp
    """Absolute start of the window actually used for this fetch."""

    endtime: pd.Timestamp
    """Absolute end of the window actually used for this fetch."""

    predicted: pd.Timestamp | None
    """Predicted phase arrival used to derive the window, or `None` if
    `entry.starttime`/`entry.endtime` were used directly."""

@define(kw_only=True)
class ProjectEntry:
    """One station/event selection within a [`PysmoProject`][pysmo.tools.project.PysmoProject].

    Pairs a station with either an event (for a phase-arrival-relative
    window, resolved at fetch time) or an explicit absolute time window
    (for event-less or continuous data), or both — an explicit window
    always takes precedence over one derived from `event`. See
    [`PysmoProject`][pysmo.tools.project.PysmoProject] for how the window is
    actually resolved.
    """

    station: Station
    """Station to fetch waveform data for."""

    event: Event | None = None
    """Event used to derive a phase-arrival-relative window, if `starttime`/`endtime` are not set."""

    starttime: pd.Timestamp | None = field(
        default=None,
        converter=converters.optional(convert_to_utc_timestamp),
        on_setattr=setters.convert,
    )
    """Explicit start of the fetch window (UTC). Overrides `event` when set together with `endtime`."""

    endtime: pd.Timestamp | None = field(
        default=None,
        converter=converters.optional(convert_to_utc_timestamp),
        on_setattr=setters.convert,
    )
    """Explicit end of the fetch window (UTC). Overrides `event` when set together with `starttime`."""

    checksum: str | None = field(default=None)
    """Checksum of the fetched seismogram, set on first fetch; `None` until then.

    A mismatch on a later fetch means the underlying archive data changed
    since this entry was first fetched — see
    [`PysmoProject.on_checksum_mismatch`][pysmo.tools.project.PysmoProject.on_checksum_mismatch]
    for how that is reported. Deliberately mutated by
    [`PysmoProject`][pysmo.tools.project.PysmoProject] as a side effect of
    fetching, so it is captured the next time the containing project is
    pickled — this is what makes it a durable reproducibility pin rather
    than a one-session-only check.

    Because this field is mutated in place, a `ProjectEntry` shared across
    two different `PysmoProject` instances (e.g. reused deliberately in an
    iterative workflow, or accidentally via a shared `entries` list) has its
    checksum set/checked by *whichever* project fetches it first — the
    entry doesn't belong to one project. Sharing entries across projects is
    fine; sharing them without being aware their checksum state is joint,
    not per-project, is the surprise to avoid.
    """

@define(kw_only=True)
class PysmoProject[T]:
    """Declares station/event data to fetch on demand and transform into `T`.

    A `PysmoProject` holds a flat list of
    [`ProjectEntry`][pysmo.tools.project.ProjectEntry] objects plus the
    parameters needed to resolve each entry's fetch window and the
    `transform` callable that turns a freshly downloaded
    [`Seismogram`][pysmo.Seismogram] into the caller's target type `T`. No
    waveform data is stored on the instance between calls beyond an
    in-memory cache of already-fetched-and-transformed results.

    See the [module documentation][pysmo.tools.project] for a worked
    example.

    Note: Thread-safety
        The in-memory fetch cache is safe to touch from multiple threads
        calling [`seismogram`][pysmo.tools.project.PysmoProject.seismogram],
        [`seismograms_for`][pysmo.tools.project.PysmoProject.seismograms_for],
        or [`fetch_all`][pysmo.tools.project.PysmoProject.fetch_all] on the
        same instance concurrently. This does not parallelise fetching
        itself, though: `fetch_seismogram` is called outside the lock, so
        two threads racing the same not-yet-cached entry both still fetch
        before one result wins and is cached.
    """

    entries: list[ProjectEntry] = field(
        factory=list, on_setattr=setters.pipe(setters.convert, _on_setattr_clear_cache)
    )
    """Station/event/window selections making up this project."""

    transform: Callable[[Seismogram, FetchContext], T] = field(
        on_setattr=setters.pipe(setters.convert, _on_setattr_clear_cache)
    )
    """Applied to a freshly downloaded seismogram; converts the result into the target type `T`.

    Called with the downloaded seismogram and a
    [`FetchContext`][pysmo.tools.project.FetchContext] carrying the
    originating entry and this fetch's resolved window/predicted arrival.
    This is also where ordinary data preparation belongs — detrending,
    resampling, removing the instrument response — not just the final
    conversion to `T`, since it is the one place every fetch already passes
    through. Free to do anything else it needs — including its own
    additional fetches (e.g. instrument response metadata via
    [`StationXML.fetch`][pysmo.classes.StationXML.fetch], demonstrated in
    the [module documentation][pysmo.tools.project]'s own example — this
    design doesn't fetch or know about response data itself, deliberately).
    Must be a top-level function in an importable module — not a lambda or
    closure — if the containing `PysmoProject` is to be pickled; a callable
    `attrs` class with only picklable fields (same example) is the
    alternative once `transform` needs its own configuration.
    """

    fetch_seismogram: Callable[[Station, pd.Timestamp, pd.Timestamp], Seismogram] = (
        field(
            default=_default_fetch_seismogram,
            on_setattr=setters.pipe(setters.convert, _on_setattr_clear_cache),
        )
    )
    """Downloads a seismogram for a station and absolute time window.

    Defaults to a private helper wrapping
    [`GeoCsvSeismogram.fetch`][pysmo.classes.GeoCsvSeismogram.fetch] — the
    explicit "always fresh, never cached" choice.

    For any project where reproducibility matters, substitute a
    [`SqliteArchiveFetcher`][pysmo.tools.archive.SqliteArchiveFetcher]
    instance instead — this is the *recommended* value for real analysis
    work, not a power-user option sitting alongside the default on equal
    footing; see its own docstring for why, and how it differs from
    `ProjectEntry.checksum`'s live-fetch drift detection. It only pins the
    waveform, though — see the [module documentation][pysmo.tools.project]'s
    second example for the gotcha this doesn't cover: `transform` making its
    own additional fetches. Any other
    callable of the right shape (e.g. one wrapping
    [`SAC.fetch`][pysmo.classes.SAC.fetch]) also works — this field changes
    the retrieval path without subclassing. Must be picklable by reference
    (a top-level function, or an attrs instance with only picklable
    fields — not a lambda or closure), same constraint as `transform`.
    """

    phase: str = field(
        default="P", on_setattr=setters.pipe(setters.convert, _on_setattr_clear_cache)
    )
    """Seismic phase used to derive a window from an entry's `event`."""

    pre_pick: NonPositiveTimedelta = field(
        default=pd.Timedelta(minutes=-2),
        converter=convert_to_timedelta,
        validator=[
            validators.instance_of(pd.Timedelta),
            validators.le(pd.Timedelta(0)),
        ],
        on_setattr=setters.pipe(
            setters.convert, setters.validate, _on_setattr_clear_cache
        ),
    )
    """Offset from the predicted arrival to the window start. Must be zero or negative."""

    post_pick: PositiveTimedelta = field(
        default=pd.Timedelta(minutes=8),
        converter=convert_to_timedelta,
        validator=[
            validators.instance_of(pd.Timedelta),
            validators.gt(pd.Timedelta(0)),
        ],
        on_setattr=setters.pipe(
            setters.convert, setters.validate, _on_setattr_clear_cache
        ),
    )
    """Offset from the predicted arrival to the window end. Must be positive."""

    travel_time_backend: TravelTimeBackend | None = field(
        default=None, on_setattr=setters.pipe(setters.convert, _on_setattr_clear_cache)
    )
    """Optional override for travel-time calculation. See [`pysmo.tools.web.TravelTimeBackend`][]."""

    on_checksum_mismatch: Literal["warn", "raise", "ignore"] = field(
        default="warn",
        validator=validators.in_(("warn", "raise", "ignore")),
    )
    """Behaviour when a fetched seismogram's checksum no longer matches the
    one recorded on `entry.checksum` from its first fetch.

    `"warn"` (default) emits a `UserWarning` and still returns the new data;
    `"raise"` raises `ValueError` instead of returning anything, for a
    pipeline that should hard-stop on detected drift; `"ignore"` returns the
    new data with no signal at all. In every case `entry.checksum` keeps the
    value from the *first* fetch — it is never overwritten by a mismatching
    value, so a mismatch is reported (or not) consistently on every
    subsequent fetch, not just the first time it's noticed.

    Deliberately not part of cache invalidation: changing this policy only
    affects how a *future* mismatch is handled, it doesn't change what data
    was fetched or would be re-fetched, so nothing about previously cached
    results becomes stale when it changes.
    """

    _cache: dict[_CacheKey, tuple[str, T]] = field(
        init=False, factory=dict, repr=False, eq=False
    )
    _lock: threading.Lock = field(
        init=False, factory=threading.Lock, repr=False, eq=False
    )
    """Guards reads/writes of `_cache` against concurrent access from more
    than one thread — see the class docstring's thread-safety note."""

    def __getstate__(self) -> dict:
        """Drop the fetch cache and lock, neither of which can survive pickling."""
        state = attrs_getstate(self, {"_cache": {}})
        del state["_lock"]
        return state

    def __setstate__(self, state: dict) -> None:
        """Restore state without triggering any `on_setattr` hooks, then create a fresh lock."""
        attrs_setstate(self, state)
        object.__setattr__(self, "_lock", threading.Lock())

    def clear_cache(self) -> None:
        """Clear the in-memory fetch cache.

        Cleared automatically whenever
        [`entries`][pysmo.tools.project.PysmoProject.entries],
        [`transform`][pysmo.tools.project.PysmoProject.transform],
        [`fetch_seismogram`][pysmo.tools.project.PysmoProject.fetch_seismogram],
        [`phase`][pysmo.tools.project.PysmoProject.phase],
        [`pre_pick`][pysmo.tools.project.PysmoProject.pre_pick],
        [`post_pick`][pysmo.tools.project.PysmoProject.post_pick], or
        [`travel_time_backend`][pysmo.tools.project.PysmoProject.travel_time_backend]
        is *reassigned*.

        Call this manually after any in-place mutation of
        [`entries`][pysmo.tools.project.PysmoProject.entries] (e.g. `append`,
        `remove`, or index assignment), which isn't observable by
        `on_setattr` and therefore doesn't clear the cache automatically.
        """
        self._cache.clear()

    def _resolve_window(
        self, entry: ProjectEntry
    ) -> tuple[pd.Timestamp, pd.Timestamp, pd.Timestamp | None]:
        """Resolve the absolute fetch window and predicted arrival for one entry.

        Returns:
            `(starttime, endtime, predicted_arrival)` — `predicted_arrival`
            is `None` when `entry.starttime`/`entry.endtime` were used
            directly rather than derived from `entry.event`.

        Raises:
            ValueError: If `entry` has neither a usable explicit window nor
                an event to derive one from, or if no `phase` arrival is
                predicted for this station/event geometry.
        """
        if entry.starttime is not None and entry.endtime is not None:
            return entry.starttime, entry.endtime, None
        if entry.event is not None:
            dist = haversine(entry.event, entry.station)
            tt = fetch_travel_times(
                entry.event.depth / 1000.0,
                dist,
                [self.phase],
                travel_time_backend=self.travel_time_backend,
            )
            if self.phase not in tt:
                raise ValueError(
                    f"No {self.phase!r} arrival predicted for "
                    f"{entry.station.network}.{entry.station.name} at this "
                    "distance/depth."
                )
            predicted = entry.event.time + pd.Timedelta(seconds=tt[self.phase])
            return predicted + self.pre_pick, predicted + self.post_pick, predicted
        raise ValueError("ProjectEntry needs either an explicit window or an event.")

    def _fetch(self, entry: ProjectEntry, *, _stacklevel: int = 3) -> T:
        """Fetch, transform, and cache the seismogram for one entry.

        Internal primitive — see
        [`seismogram`][pysmo.tools.project.PysmoProject.seismogram] for the
        public, station/event-based accessor built on top of this.

        Args:
            entry: The station/event/window selection to fetch.

        Returns:
            The transformed result for `entry`, from cache if this exact
            entry has been fetched before.

        Raises:
            ValueError: If `entry` has neither a usable window nor an event
                (via `_resolve_window`); if the underlying fetch raises
                (e.g. no waveform data for the resolved window); or if the
                checksum no longer matches and `on_checksum_mismatch="raise"`.
        """
        event_key: _EventKey = (
            (
                entry.event.latitude,
                entry.event.longitude,
                entry.event.depth,
                entry.event.time,
            )
            if entry.event is not None
            else None
        )
        key: _CacheKey = (
            entry.station.network,
            entry.station.name,
            entry.station.location,
            entry.station.channel,
            event_key,
            entry.starttime,
            entry.endtime,
        )
        with self._lock:
            cached = self._cache.get(key)
        if cached is None:
            # Deliberately outside the lock — see the class docstring's
            # thread-safety note: two threads racing the same not-yet-cached
            # key both fetch here (a stampede) before one result wins.
            starttime, endtime, predicted = self._resolve_window(entry)
            seismogram = self.fetch_seismogram(entry.station, starttime, endtime)
            checksum = _checksum(seismogram)
            context = FetchContext(
                entry=entry,
                starttime=starttime,
                endtime=endtime,
                predicted=predicted,
            )
            cached = (checksum, self.transform(seismogram, context))
            with self._lock:
                cached = self._cache.setdefault(key, cached)

        checksum, result = cached
        if entry.checksum is None:
            entry.checksum = checksum
        elif entry.checksum != checksum and self.on_checksum_mismatch != "ignore":
            message = (
                f"Fetched data for {entry.station.network}.{entry.station.name} "
                "no longer matches the checksum recorded when this entry was "
                "first fetched — the underlying archive may have been revised."
            )
            if self.on_checksum_mismatch == "raise":
                raise ValueError(message)
            # `_stacklevel` is threaded in from the public entry point
            # (`seismogram`/`fetch_all` pass the default; `seismograms_for`
            # passes one level deeper) so the warning always points at the
            # user's own call site, not an intermediate method.
            warnings.warn(message, stacklevel=_stacklevel)
        return result

    @property
    def stations(self) -> list[Station]:
        """Distinct stations across all entries, in first-seen order.

        A plain `@property`, not `@cached_property` — recomputed on each
        access, same as
        [`ICCS.cc_seismograms`][pysmo.tools.iccs.ICCS.cc_seismograms]'s
        precedent for a no-arg derived list view in this codebase. Compares
        with `==` (attrs-generated equality, not identity or hashing —
        `Station` is not hashable).
        """
        seen: list[Station] = []
        for entry in self.entries:
            if entry.station not in seen:
                seen.append(entry.station)
        return seen

    @property
    def events(self) -> list[Event]:
        """Distinct events across all entries, excluding event-less entries.

        In first-seen order; compares with `==`, same caveat as `stations`.
        """
        seen: list[Event] = []
        for entry in self.entries:
            if entry.event is not None and entry.event not in seen:
                seen.append(entry.event)
        return seen

    def events_for(self, station: Station) -> list[Event | None]:
        """Events available for one station, in first-seen order.

        `None` appears in the result if `station` has an event-less entry —
        an event-less selection is a first-class member of this list, not a
        special case to check for separately.
        """
        seen: list[Event | None] = []
        for entry in self.entries:
            if entry.station == station and entry.event not in seen:
                seen.append(entry.event)
        return seen

    def stations_for(self, event: Event | None) -> list[Station]:
        """Stations available for one event, in first-seen order.

        Pass `None` for stations with an event-less entry.
        """
        seen: list[Station] = []
        for entry in self.entries:
            if entry.event == event and entry.station not in seen:
                seen.append(entry.station)
        return seen

    def seismogram(
        self, station: Station, event: Event | None = None, *, _stacklevel: int = 3
    ) -> T:
        """Fetch (or return from cache) the result for one station/event combination.

        Args:
            station: Station to fetch.
            event: Event to fetch for, or `None` for an event-less entry.

        Returns:
            The transformed result for the matching entry.

        Raises:
            KeyError: If no entry matches this station/event combination.
            ValueError: If more than one entry matches — an authoring
                mistake (e.g. the same station/event added twice with
                different explicit windows), surfaced rather than silently
                resolved by picking one.
        """
        matches = [e for e in self.entries if e.station == station and e.event == event]
        if not matches:
            raise KeyError("No entry for this station/event combination.")
        if len(matches) > 1:
            raise ValueError(
                "More than one entry matches this station/event combination."
            )
        return self._fetch(matches[0], _stacklevel=_stacklevel)

    def seismograms_for(self, event: Event) -> list[T]:
        """All seismograms for one event — e.g. ready for `ICCS(seismograms=...)`.

        Built from
        [`stations_for`][pysmo.tools.project.PysmoProject.stations_for] and
        [`seismogram`][pysmo.tools.project.PysmoProject.seismogram], not an
        independent filter over `entries`.

        Typed to require an `Event`, unlike `stations_for`/`events_for`
        (which both treat `None` as first-class) — deliberately: this
        method exists for the event-based bulk-fetch use case (`ICCS`),
        which has no equivalent "all event-less entries" workflow to
        support. `[seismogram(s, None) for s in stations_for(None)]`
        already covers that case directly if it's ever needed.
        """
        return [
            self.seismogram(station, event, _stacklevel=4)
            for station in self.stations_for(event)
        ]

    def fetch_all(self) -> list[T]:
        """Fetch every entry in the project.

        With the default, always-fresh `fetch_seismogram`, this just warms
        `_cache` for the session. With an archive-backed `fetch_seismogram`
        (e.g.
        [`SqliteArchiveFetcher`][pysmo.tools.archive.SqliteArchiveFetcher]),
        this is what actually populates the archive — a single, explicit
        "get everything this project needs into the archive" call, rather
        than relying on incidental use of `seismogram`/`seismograms_for` to
        cover every entry eventually.

        Returns:
            One transformed result per entry, in `entries` order.
        """
        return [self._fetch(entry) for entry in self.entries]