aimbat.core

def add_data_to_project(
    session: Session,
    data_sources: Sequence[os.PathLike[str] | str],
    data_type: DataType,
    station_id: UUID | None = None,
    event_id: UUID | None = None,
    dry_run: bool = False,
    on_progress: Callable[[int, int], None] | None = None,
) -> tuple[list[AimbatDataSource], set[UUID], set[UUID], set[UUID], list[str]]:
    """Add data sources to the AIMBAT database.

    What gets created depends on which capabilities `data_type` supports:

    - Station + event + seismogram: all three records are created and linked,
      and an `AimbatDataSource` entry is stored.
    - Station or event only (e.g. `JSON_STATION`, `JSON_EVENT`): only the
      relevant metadata records are created; no seismogram or data source entry
      is stored.

    Use `station_id` or `event_id` to skip extracting station or event metadata
    from the data source and link to a pre-existing record instead.

    A new event whose origin time nearly, but not exactly, matches a
    pre-existing event's time is handled per
    `settings.event_duplicate_tolerance`, `event_duplicate_raise_tolerance`,
    and `event_duplicate_strict`. Within the tighter tolerance the existing
    event is reused, with a log warning (a dry run additionally returns the
    message in `duplicate_warnings`). Within the wider "ambiguous gap" band,
    a `ValueError` is raised unconditionally, even during a dry run.

    Args:
        session: The SQLModel database session.
        data_sources: List of data sources to add.
        data_type: Type of data.
        station_id: UUID of an existing station to use instead of extracting
            one from each data source.
        event_id: UUID of an existing event to use instead of extracting one
            from each data source.
        dry_run: If True, do not commit changes to the database.
        on_progress: Optional callback invoked as `on_progress(done, total)`
            after each data source is processed, for callers that want to
            display progress.

    Returns:
        A 5-tuple of `(added_datasources, existing_station_ids,
        existing_event_ids, existing_seismogram_ids, duplicate_warnings)`.
        `added_datasources` is every `AimbatDataSource` touched by this call,
        new or reused. `existing_*_ids` are the sets of station/event/
        seismogram IDs that already existed in the database *before* this
        call, so callers can tell which entries in `added_datasources` are
        newly created versus reused by comparing IDs against these sets.
        `duplicate_warnings` lists near-duplicate-event reuse messages
        collected during a dry run (always empty on a real-add call, which
        logs the same message as a warning instead).

    Raises:
        ValueError: If a near-duplicate event falls in the "ambiguous gap"
            band (see above).

    Anything raised while a data source is being processed carries a
    `Data source: <name>` note identifying it, which
    `utils.exception_message` renders alongside the message.
    """

    logger.info(f"Adding {len(data_sources)} {data_type} data sources to project.")

    missing_sources = [str(ds) for ds in data_sources if not Path(ds).exists()]
    if missing_sources:
        raise FileNotFoundError(
            "Data source(s) not found: " + ", ".join(missing_sources)
        )

    if station_id is not None and session.get(AimbatStation, station_id) is None:
        raise NoResultFound(f"No station found with ID {station_id}.")
    if event_id is not None and session.get(AimbatEvent, event_id) is None:
        raise NoResultFound(f"No event found with ID {event_id}.")

    existing_station_ids, existing_event_ids, existing_seismogram_ids = (
        _existing_entity_ids(session)
    )

    # Mutated in place as new events are created, so near-duplicate
    # detection also covers events created earlier in this same batch,
    # unlike existing_event_ids, which stays a frozen pre-batch snapshot
    # for the return value's new-vs-reused semantics.
    known_event_ids = set(existing_event_ids)

    try:
        added_datasources: list[AimbatDataSource] = []
        duplicate_warnings: list[str] = []
        total = len(data_sources)
        with session.begin_nested() as nested:
            for done, datasource in enumerate(data_sources, start=1):
                try:
                    result, duplicate_warning = _process_datasource(
                        session,
                        datasource,
                        data_type,
                        station_id,
                        event_id,
                        dry_run,
                        known_event_ids,
                    )
                except Exception as e:
                    # A note rather than a wrapping exception, so callers
                    # catching the documented ValidationError/ValueError/
                    # NoResultFound still do. Displayed by
                    # `utils.exception_message`.
                    e.add_note(f"Data source: {datasource}")
                    raise
                if result is not None:
                    added_datasources.append(result)
                if duplicate_warning is not None:
                    duplicate_warnings.append(duplicate_warning)
                if on_progress is not None:
                    on_progress(done, total)

            if dry_run:
                _roll_back_dry_run(session, nested, added_datasources)
                return (
                    added_datasources,
                    existing_station_ids,
                    existing_event_ids,
                    existing_seismogram_ids,
                    duplicate_warnings,
                )

        session.commit()
        logger.info("Data added successfully.")
        return (
            added_datasources,
            existing_station_ids,
            existing_event_ids,
            existing_seismogram_ids,
            duplicate_warnings,
        )

    except Exception as e:
        logger.error(
            "Failed to add data. Rolling back changes. Error: " + exception_message(e)
        )
        raise

def add_seismograms_to_project(
    session: Session,
    items: Sequence[tuple[IccsSeismogram, Station, Event]],
    *,
    data_dir: str | os.PathLike[str],
    dry_run: bool = False,
    on_progress: Callable[[int, int], None] | None = None,
) -> tuple[list[AimbatDataSource], set[UUID], set[UUID], set[UUID], list[str]]:
    """Ingest already-built (seismogram, station, event) triples directly.

    The general-purpose counterpart to `add_data_to_project` for a caller
    that already has pysmo `Seismogram`/`Station`/`Event`-shaped objects in
    memory - fetched from a web service, built in a notebook, or produced by
    a library such as `pysmo.tools.project.PysmoProject` - with no file
    round-trip through a `DataType` AIMBAT already understands. Each
    seismogram's waveform is persisted as a miniSEED file under `data_dir`
    so it can be read back later like any other data source; station and
    event metadata are taken directly from the objects given.

    Every triple must carry a real `Event`: `AimbatSeismogram.event_id` is a
    required foreign key, and AIMBAT's processing model (ICCS, MCCC,
    snapshots, the active-event mechanism) is per-event throughout, so an
    event-less seismogram would have nothing to attach to. Reuse the same
    `Event` object across several triples to link them to one event -
    `_link_event`'s near-duplicate detection still applies, so a second,
    slightly different `Event` object (e.g. from a re-run) resolves to the
    same row rather than duplicating it.

    Args:
        session: The SQLModel database session.
        items: Sequence of `(seismogram, station, event)` triples to add.
        data_dir: Directory the persisted miniSEED files are written to.
            Created (with any missing parents) if it does not exist yet.
        dry_run: If True, do not commit changes to the database.
        on_progress: Optional callback invoked as `on_progress(done, total)`
            after each triple is processed, for callers that want to
            display progress.

    Returns:
        Same 5-tuple shape as `add_data_to_project`:
        `(added_datasources, existing_station_ids, existing_event_ids,
        existing_seismogram_ids, duplicate_warnings)`.

    Raises:
        ValueError: If a near-duplicate event falls in the "ambiguous gap"
            band (see `add_data_to_project`).

    Note:
        Miniseed files are written to `data_dir` before the final commit. If
        that commit fails (disk full, a constraint violation surfaced only
        at commit time, a concurrent writer), the files already written
        survive on disk with no corresponding database row - an orphan
        `aimbat data prune` cannot detect, since prune only looks for
        database rows whose file has vanished, never the reverse. Low
        severity (disk clutter, not data loss) and the failure window is
        narrow, but worth knowing about before relying on `data_dir` being
        an exact mirror of the database.
    """

    logger.info(f"Adding {len(items)} seismogram(s) to project directly.")

    data_dir = Path(data_dir)

    existing_station_ids, existing_event_ids, existing_seismogram_ids = (
        _existing_entity_ids(session)
    )

    known_event_ids = set(existing_event_ids)

    try:
        added_datasources: list[AimbatDataSource] = []
        duplicate_warnings: list[str] = []
        pending_writes: list[tuple[IccsSeismogram, Station, Path]] = []
        total = len(items)
        with session.begin_nested() as nested:
            for done, (seismogram, station, event) in enumerate(items, start=1):
                label = f"{station.network}.{station.name}"
                aimbat_station = _link_station(
                    session, AimbatStation.model_validate(station), label
                )

                new_aimbat_event = AimbatEvent.model_validate(
                    event, update={"parameters": AimbatEventParameters()}
                )
                aimbat_event, duplicate_warning = _link_event(
                    session, new_aimbat_event, dry_run, known_event_ids, label
                )
                if duplicate_warning is not None:
                    duplicate_warnings.append(duplicate_warning)

                path = _seismogram_path(seismogram, station, data_dir)
                sourcename, datatype = str(path), DataType.MSEED
                already_persisted = (
                    session.exec(
                        select(AimbatDataSource.id).where(
                            AimbatDataSource.sourcename == sourcename
                        )
                    ).first()
                    is not None
                )
                if already_persisted:
                    message = _format_seismogram_collision_message(label, sourcename)
                    logger.warning(message)
                    if dry_run:
                        duplicate_warnings.append(message)
                else:
                    # Writing is deferred until the whole batch has been
                    # validated and linked, so a later triple's failure (e.g.
                    # an ambiguous-gap event) never leaves earlier triples'
                    # files orphaned on disk with no committed DB row.
                    pending_writes.append((seismogram, station, path))

                new_aimbat_seismogram = AimbatSeismogram.model_validate(
                    seismogram,
                    update={
                        "t0": seismogram.t0,
                        "parameters": AimbatSeismogramParameters(),
                    },
                )
                aimbat_seismogram = _link_seismogram(
                    session, new_aimbat_seismogram, sourcename
                )
                aimbat_seismogram.station = aimbat_station
                aimbat_seismogram.event = aimbat_event

                logger.debug(
                    f"Linking seismogram from {sourcename} to Station="
                    + f"{aimbat_station.name} and EventTime={aimbat_event.time}."
                )

                aimbat_data_source = _link_datasource(
                    session, sourcename, datatype, aimbat_seismogram
                )
                added_datasources.append(aimbat_data_source)

                if on_progress is not None:
                    on_progress(done, total)

            if dry_run:
                _roll_back_dry_run(session, nested, added_datasources)
                return (
                    added_datasources,
                    existing_station_ids,
                    existing_event_ids,
                    existing_seismogram_ids,
                    duplicate_warnings,
                )

            if pending_writes:
                data_dir.mkdir(parents=True, exist_ok=True)
                written: list[Path] = []
                try:
                    for (
                        pending_seismogram,
                        pending_station,
                        pending_path,
                    ) in pending_writes:
                        _persist(pending_seismogram, pending_station, pending_path)
                        written.append(pending_path)
                except Exception:
                    for written_path in written:
                        written_path.unlink(missing_ok=True)
                    raise

        session.commit()
        logger.info("Data added successfully.")
        return (
            added_datasources,
            existing_station_ids,
            existing_event_ids,
            existing_seismogram_ids,
            duplicate_warnings,
        )

    except Exception as e:
        logger.error(f"Failed to add data. Rolling back changes. Error: {e}")
        raise

def build_iccs_from_snapshot(session: Session, snapshot_id: UUID) -> BoundICCS:
    """Build a read-only BoundICCS from a snapshot's parameters and live waveform data.

    Uses the snapshot's event and seismogram parameters (window, t1, flip, select,
    bandpass, etc.) but reads waveform data from the live datasources. Seismograms
    added after the snapshot was taken are not included in the snapshot; their live
    parameters are used instead. No DB writes occur at any point.

    Args:
        session: Database session.
        snapshot_id: ID of the AimbatSnapshot to load.

    Returns:
        BoundICCS instance built from the snapshot parameters.

    Raises:
        ValueError: If no snapshot with the given ID is found.
    """
    logger.info(f"Building ICCS from snapshot {snapshot_id}.")

    statement = (
        select(AimbatSnapshot)
        .where(AimbatSnapshot.id == snapshot_id)
        .options(
            selectinload(rel(AimbatSnapshot.event))
            .selectinload(rel(AimbatEvent.seismograms))
            .selectinload(rel(AimbatSeismogram.parameters)),
            selectinload(rel(AimbatSnapshot.event_parameters_snapshot)),
            selectinload(rel(AimbatSnapshot.seismogram_parameters_snapshots)),
        )
    )
    snapshot = session.exec(statement).one_or_none()

    if snapshot is None:
        raise ValueError(f"Snapshot {snapshot_id} not found.")

    ep = snapshot.event_parameters_snapshot
    snap_params = AimbatEventParametersBase.model_validate(ep)

    # Build a map from seismogram_id → snapshot parameters
    snap_seis_map = {
        sp.seismogram_id: sp for sp in snapshot.seismogram_parameters_snapshots
    }

    seismograms = []
    for seis in snapshot.event.seismograms:
        snap_sp = snap_seis_map.get(seis.id)
        if snap_sp is None:
            # Seismogram was added after the snapshot: use live parameters
            seis_params = AimbatSeismogramParametersBase.model_validate(seis.parameters)
        else:
            seis_params = AimbatSeismogramParametersBase.model_validate(snap_sp)
        seismograms.append(
            MiniIccsSeismogram(
                begin_time=seis.begin_time,
                delta=seis.delta,
                data=seis.data,
                t0=seis.t0,
                t1=seis_params.t1,
                flip=seis_params.flip,
                select=seis_params.select,
                extra={"id": seis.id},
            )
        )

    iccs = ICCS(
        seismograms=seismograms,
        window_pre=snap_params.window_pre,
        window_post=snap_params.window_post,
        ramp_width=snap_params.ramp_width,
        bandpass_apply=snap_params.bandpass_apply,
        bandpass_fmin=snap_params.bandpass_fmin,
        bandpass_fmax=snap_params.bandpass_fmax,
        corners=snap_params.corners,
        min_cc=snap_params.min_cc,
        context_width=settings.context_width,
    )
    return BoundICCS(
        iccs=iccs,
        event_id=snapshot.event_id,
        created_at=Timestamp.now("UTC"),
    )