aimbat.models

class AimbatDataSource(SQLModel, table=True):
    """Logical source identifier and format of the waveform data for a seismogram."""

    model_config = SQLModelConfig(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    # Named explicitly so `create_all` (via `create_project`) and the Alembic
    # migration that added it (`72c6b97febca`) produce an identically-named
    # constraint, keeping the two schema-creation paths in sync.
    __table_args__ = (
        UniqueConstraint("sourcename", name="uq_aimbatdatasource_sourcename"),
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        title="ID",
        description="Unique data source ID.",
        schema_extra={"rich": RichColSpec(style="yellow", highlight=False)},
    )
    sourcename: str = Field(
        title="Source name",
        description=(
            "Opaque logical source identifier for the waveform data. "
            + "Uniqueness is enforced as plain string equality, not "
            + "canonicalised, so a file-based source reachable under two "
            + "different (but equivalent) path spellings is not recognised "
            + "as the same source - callers are responsible for "
            + "canonicalising before storing."
        ),
    )
    datatype: DataType = Field(
        default=DataType.SAC,
        title="Data type",
        description="Data type of the data source.",
    )
    seismogram_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatseismogram.id",
        ondelete="CASCADE",
        title="Seismogram ID",
        description="Foreign key referencing the parent seismogram.",
        schema_extra={"rich": RichColSpec(style="magenta", highlight=False)},
    )
    seismogram: "AimbatSeismogram" = Relationship(back_populates="datasource")
    "The seismogram this data source belongs to."

class AimbatEvent(SQLModel, table=True):
    """Seismic event with origin time, location, and depth."""

    model_config = SQLModelConfig(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        title="ID",
        description="Unique event ID.",
        schema_extra={"rich": RichColSpec(style="yellow", highlight=False)},
    )
    time: PydanticTimestamp = Field(
        unique=True,
        sa_type=SAPandasTimestamp,
        title="Time",
        description="Event origin time (UTC).",
    )
    latitude: float = Field(title="Latitude", description="Event latitude in degrees.")
    longitude: float = Field(
        title="Longitude", description="Event longitude in degrees."
    )
    depth: float | None = Field(
        default=None, title="Depth", description="Event depth in metres."
    )
    last_modified: PydanticTimestamp | None = Field(
        default=None,
        sa_type=SAPandasTimestamp,
        title="Last modified",
        description="Timestamp of the last parameter modification.",
    )
    stack_modified: PydanticTimestamp | None = Field(
        default=None,
        sa_type=SAPandasTimestamp,
        title="Stack modified",
        description=(
            "Timestamp of the last change to a parameter that requires the ICCS "
            + "instance to be rebuilt (event window, ramp, bandpass, corners; "
            + "per-seismogram t1, flip, select). Drives `BoundICCS.is_stale`. "
            + "Unlike `last_modified`, an MCCC-only or `min_cc` change does not "
            + "bump it."
        ),
    )
    seismograms: list[AimbatSeismogram] = Relationship(
        back_populates="event", cascade_delete=True
    )
    "List of seismograms of this event."

    parameters: AimbatEventParameters = Relationship(
        back_populates="event", cascade_delete=True
    )
    "Event parameters."

    quality: AimbatEventQuality | None = Relationship(
        back_populates="event", cascade_delete=True
    )
    "Live quality metrics for this event."

    snapshots: list[AimbatSnapshot] = Relationship(
        back_populates="event", cascade_delete=True
    )
    "List of snapshots."

    if TYPE_CHECKING:
        # Column properties defined below, but add same default values for type checking purposes
        seismogram_count: int = 0
        station_count: int = 0
        snapshot_count: int = 0

class AimbatEventParameters(AimbatEventParametersBase, table=True):
    """Processing parameters common to all seismograms of a particular event."""

    model_config = SQLModelConfig(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    __table_args__ = (
        Index("ix_aimbateventparameters_event_id", "event_id", unique=True),
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        title="ID",
        description="Unique ID.",
        schema_extra={
            "rich": RichColSpec(style="yellow", no_wrap=True, highlight=False)
        },
    )
    event_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatevent.id",
        ondelete="CASCADE",
        title="Event ID",
        description="Foreign key referencing the parent event.",
        schema_extra={
            "rich": RichColSpec(style="magenta", no_wrap=True, highlight=False)
        },
    )
    event: "AimbatEvent" = Relationship(back_populates="parameters")
    "The event these parameters belong to."
    snapshots: list["AimbatEventParametersSnapshot"] = Relationship(
        back_populates="parameters", passive_deletes=True
    )
    "Parameter snapshots for this event."

class AimbatEventParametersBase(SQLModel):
    """Base class defining event-level processing parameters for AIMBAT.

    This class serves as a base that is inherited by the actual classes that
    create the database tables. It covers the settings applied uniformly to
    all seismograms of an event: window bounds, taper ramp width, bandpass
    filter settings, the CC threshold for automatic de-selection, and MCCC
    settings.
    """

    completed: bool = Field(
        default=False,
        title="Completed",
        description="Mark an event as completed.",
    )

    ramp_width: PydanticNonNegativeFloat = Field(
        default_factory=lambda: settings.ramp_width,
        sa_type=Float,
        title="Ramp width",
        description=(
            "Width of taper ramp as a multiple of the window length. Values "
            + "greater than 1 are valid; the ramp extends outside the window."
        ),
    )

    window_pre: PydanticNegativeTimedelta = Field(
        sa_type=SAPandasTimedelta,
        default_factory=lambda: settings.window_pre,
        title="Window pre",
        description="Pre-pick window length in seconds.",
    )

    window_post: PydanticPositiveTimedelta = Field(
        sa_type=SAPandasTimedelta,
        default_factory=lambda: settings.window_post,
        title="Window post",
        description="Post-pick window length in seconds.",
    )

    bandpass_apply: bool = Field(
        default_factory=lambda: settings.bandpass_apply,
        title="Bandpass apply",
        description="Whether to apply bandpass filter to seismograms.",
    )

    bandpass_fmin: float = Field(
        default_factory=lambda: settings.bandpass_fmin,
        ge=0,
        title="Bandpass f min",
        description=(
            "Minimum frequency for bandpass filter in Hz (ignored if "
            + "`bandpass_apply` is False)."
        ),
    )

    bandpass_fmax: float = Field(
        default_factory=lambda: settings.bandpass_fmax,
        gt=0,
        title="Bandpass f max",
        description=(
            "Maximum frequency for bandpass filter in Hz (ignored if "
            + "`bandpass_apply` is False)."
        ),
    )

    corners: int = Field(
        default_factory=lambda: settings.corners,
        gt=0,
        title="Corners",
        description=(
            "Number of corners (poles) for the bandpass filter (ignored if "
            + "`bandpass_apply` is False)."
        ),
        sa_column_kwargs={"server_default": "2"},
    )

    min_cc: float = Field(
        ge=0.0,
        le=1.0,
        default_factory=lambda: settings.min_cc,
        title="Min CC",
        description=(
            "Minimum cross-correlation used when automatically de-selecting "
            + "seismograms."
        ),
    )

    mccc_damp: float = Field(
        default_factory=lambda: settings.mccc_damp,
        ge=0,
        title="MCCC damp",
        description="Damping factor for MCCC algorithm.",
    )

    mccc_min_cc: float = Field(
        default_factory=lambda: settings.mccc_min_cc,
        ge=0,
        le=1,
        title="MCCC min CC",
        description=(
            "Minimum correlation coefficient required to include a pair in the "
            + "MCCC inversion."
        ),
    )

    @model_validator(mode="after")
    def check_freq_range(self) -> Self:
        """Validate that `bandpass_fmax` is strictly greater than `bandpass_fmin`."""
        if self.bandpass_fmax <= self.bandpass_fmin:
            raise ValueError(
                f"bandpass_fmax ({self.bandpass_fmax:g} Hz) must be greater than "
                + f"bandpass_fmin ({self.bandpass_fmin:g} Hz). Setting one bound at "
                + "a time, widen the band first: raise bandpass_fmax before "
                + "bandpass_fmin, lower bandpass_fmin before bandpass_fmax."
            )
        return self

class AimbatEventParametersSnapshot(AimbatEventParametersBase, table=True):
    """Snapshot of processing parameters for a particular event.

    Frozen history: the source event is still reachable via
    `snapshot.event_id`, and `parameters_id` is set to `NULL` if the live
    parameter row is deleted.
    """

    model_config = SQLModelConfig(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    snapshot_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatsnapshot.id",
        ondelete="CASCADE",
        title="Snapshot ID",
        description="Foreign key referencing the parent snapshot.",
    )
    parameters_id: uuid.UUID | None = Field(
        default=None,
        foreign_key="aimbateventparameters.id",
        ondelete="SET NULL",
        title="Event parameters ID",
        description=(
            "Foreign key referencing the source event parameters, or NULL if it "
            + "has been deleted."
        ),
    )
    snapshot: "AimbatSnapshot" = Relationship(
        back_populates="event_parameters_snapshot"
    )
    "The snapshot this record belongs to."
    parameters: AimbatEventParameters | None = Relationship(back_populates="snapshots")
    "The event parameters this snapshot was taken from, if still present."

class AimbatEventQuality(AimbatEventQualityBase, table=True):
    """Live quality metrics for a seismic event.

    One row per event. Updated in place whenever MCCC runs.
    Fields are `None` until MCCC has been executed.
    """

    model_config = SQLModelConfig(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    __table_args__ = (Index("ix_aimbateventquality_event_id", "event_id", unique=True),)

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    event_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatevent.id",
        ondelete="CASCADE",
        title="Event ID",
        description="Foreign key referencing the parent event.",
    )
    event: "AimbatEvent" = Relationship(back_populates="quality")
    "The event these quality metrics belong to."
    snapshots: list["AimbatEventQualitySnapshot"] = Relationship(
        back_populates="quality", passive_deletes=True
    )
    "Quality snapshots taken from this live record."

class AimbatEventQualityBase(SQLModel):
    """Base class defining event-level quality metrics.

    Fields are `None` when the corresponding algorithm has not been run yet.
    """

    mccc_rmse: PydanticNonNegativeTimedelta | None = Field(
        default=None,
        sa_type=SAPandasTimedelta,
        title="MCCC RMSE",
        description="Root-mean-square error of the MCCC inversion fit across the whole array.",
    )

class AimbatEventQualitySnapshot(AimbatEventQualityBase, table=True):
    """Snapshot of quality metrics for a seismic event.

    Frozen history: the source event is still reachable via
    `snapshot.event_id`, and `event_quality_id` is set to `NULL` if the live
    quality row is deleted.
    """

    model_config = SQLModelConfig(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    event_quality_id: uuid.UUID | None = Field(
        default=None,
        foreign_key="aimbateventquality.id",
        ondelete="SET NULL",
        title="Event quality ID",
        description=(
            "Foreign key referencing the source event quality, or NULL if it has "
            + "been deleted."
        ),
    )
    quality: AimbatEventQuality | None = Relationship(back_populates="snapshots")
    "The event quality this snapshot was taken from, if still present."
    snapshot_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatsnapshot.id",
        ondelete="CASCADE",
        title="Snapshot ID",
        description="Foreign key referencing the parent snapshot.",
    )
    snapshot: "AimbatSnapshot" = Relationship(back_populates="event_quality_snapshot")
    "The snapshot this record belongs to."

class AimbatEventRead(BaseModel):
    """Read model for AimbatEvent including computed counts."""

    model_config = ConfigDict(
        frozen=True,
        alias_generator=to_camel,
        populate_by_name=True,
    )

    id: UUID = Field(
        title="ID",
        description="Unique identifier for the event",
        json_schema_extra={
            "rich": RichColSpec(style="yellow"),  # type: ignore[dict-item]
        },
    )
    short_id: str | None = Field(
        default=None,
        title="Short ID",
        description="Shortened unique identifier",
        json_schema_extra={
            "tui": TuiColSpec(display_title="ID"),  # type: ignore[dict-item]
            "rich": RichColSpec(display_title="ID", style="yellow", highlight=False),  # type: ignore[dict-item]
        },
    )
    completed: bool = Field(
        title="Completed",
        description="Indicates if the event's parameters are marked as completed",
        json_schema_extra={"tui": TuiColSpec(text_align="center")},  # type: ignore[dict-item]
    )
    time: PydanticTimestamp = Field(
        title="Event time",
        description="Origin time of the event",
        json_schema_extra={
            "tui": TuiColSpec(formatter=fmt_timestamp),  # type: ignore[dict-item]
            "rich": RichColSpec(formatter=fmt_timestamp),  # type: ignore[dict-item]
        },
    )
    latitude: float = Field(
        title="Latitude",
        description="Latitude of the event",
    )
    longitude: float = Field(
        title="Longitude",
        description="Longitude of the event",
    )
    depth: float | None = Field(
        title="Depth",
        description="Depth of the event",
        json_schema_extra={
            "tui": TuiColSpec(display_title="Depth km", formatter=fmt_depth_km),  # type: ignore[dict-item]
            "rich": RichColSpec(
                display_title=r"Depth \[km]", justify="right", formatter=fmt_depth_km
            ),  # type: ignore[dict-item]
        },
    )
    seismogram_count: int = Field(
        title="Seismograms",
        description="Number of seismograms associated with this event",
    )
    station_count: int = Field(
        title="Stations",
        description="Number of stations associated with this event",
    )
    snapshot_count: int = Field(
        title="Snapshots",
        description="Number of snapshots associated with this event",
    )
    last_modified: PydanticTimestamp | None = Field(
        default=None,
        title="Last modified",
        description="Timestamp of the last modification of this event's parameters",
        json_schema_extra={
            "tui": TuiColSpec(formatter=fmt_timestamp),  # type: ignore[dict-item]
            "rich": RichColSpec(formatter=fmt_timestamp),  # type: ignore[dict-item]
        },
    )

    @classmethod
    def from_event(cls, event: AimbatEvent, session: Session | None = None) -> Self:
        """Create an AimbatEventRead from an AimbatEvent ORM instance."""
        data = event.model_dump()

        if session is not None:
            from aimbat.utils import uuid_shortener

            data["short_id"] = uuid_shortener(session, event)

        data.update(
            {
                "completed": event.parameters.completed if event.parameters else False,
                "seismogram_count": event.seismogram_count or 0,
                "station_count": event.station_count or 0,
                "snapshot_count": event.snapshot_count or 0,
            }
        )
        return cls(**data)

class AimbatNote(SQLModel, table=True):
    """Free-text Markdown note attached to an event, station, seismogram, or snapshot.

    Exactly one of the four FK fields is set per row, and at most one note
    exists per parent record (enforced by the per-FK partial unique indexes).
    Deletion of the parent record cascades to delete the note via the DB-level
    foreign key constraint.
    """

    model_config = SQLModelConfig(
        alias_generator=to_camel,
        populate_by_name=True,
    )

    __table_args__ = (
        CheckConstraint(
            (
                "(CASE WHEN event_id IS NOT NULL THEN 1 ELSE 0 END + CASE WHEN "
                + "station_id IS NOT NULL THEN 1 ELSE 0 END + CASE WHEN seismogram_id "
                + "IS NOT NULL THEN 1 ELSE 0 END + CASE WHEN snapshot_id IS NOT NULL "
                + "THEN 1 ELSE 0 END) = 1"
            ),
            name="aimbat_note_exactly_one_parent",
        ),
        Index(
            "ix_aimbatnote_event_id",
            "event_id",
            unique=True,
            sqlite_where=text("event_id IS NOT NULL"),
        ),
        Index(
            "ix_aimbatnote_station_id",
            "station_id",
            unique=True,
            sqlite_where=text("station_id IS NOT NULL"),
        ),
        Index(
            "ix_aimbatnote_seismogram_id",
            "seismogram_id",
            unique=True,
            sqlite_where=text("seismogram_id IS NOT NULL"),
        ),
        Index(
            "ix_aimbatnote_snapshot_id",
            "snapshot_id",
            unique=True,
            sqlite_where=text("snapshot_id IS NOT NULL"),
        ),
    )

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4,
        primary_key=True,
        description="Unique note ID.",
    )
    content: str = Field(
        default="",
        description="Note content in Markdown format.",
    )
    event_id: uuid.UUID | None = Field(
        default=None,
        foreign_key="aimbatevent.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent event.",
    )
    station_id: uuid.UUID | None = Field(
        default=None,
        foreign_key="aimbatstation.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent station.",
    )
    seismogram_id: uuid.UUID | None = Field(
        default=None,
        foreign_key="aimbatseismogram.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent seismogram.",
    )
    snapshot_id: uuid.UUID | None = Field(
        default=None,
        foreign_key="aimbatsnapshot.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent snapshot.",
    )

    @model_validator(mode="after")
    def _exactly_one_parent(self) -> "AimbatNote":
        # Table-model __init__ skips Pydantic validation, so this never runs
        # on the normal construction path; the DB CheckConstraint above is
        # the real guard. Kept for the direct-Pydantic-construction case.
        set_count = sum(
            fk is not None
            for fk in (
                self.event_id,
                self.station_id,
                self.seismogram_id,
                self.snapshot_id,
            )
        )
        if set_count != 1:
            raise ValueError(
                "Exactly one of event_id, station_id, seismogram_id, snapshot_id must"
                + " be set on AimbatNote."
            )
        return self