aimbat.models

class AimbatDataSource(SQLModel, table=True):
    """Class to store data source information."""

    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 data source ID.",
        schema_extra={"rich": RichColSpec(style="yellow", highlight=False)},
    )
    sourcename: str = Field(
        title="Source name",
        description="Path or name of the data source.",
    )
    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,
        allow_mutation=False,
        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.",
    )
    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,
    )

    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", cascade_delete=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. The attributes correspond exactly to the AIMBAT
    event parameters.
    """

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

    ramp_width: float = Field(
        default_factory=lambda: settings.ramp_width,
        title="Ramp width",
        description="Width of taper ramp up and down as a fraction of the window length.",
    )

    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).",
    )

    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("bandpass_fmax must be > bandpass_fmin")
        return self

    @model_validator(mode="after")
    def validate_iccs_context(self, info: ValidationInfo) -> Self:
        """Attempt ICCS construction with these parameters if requested.

        Requires `validate_iccs=True` and an `event` instance in the validation
        context.
        """
        context = info.context or {}
        if context.get("validate_iccs"):
            event = context.get("event")
            if event:
                from aimbat.core._iccs import validate_iccs_construction

                try:
                    validate_iccs_construction(event, parameters=self)
                except Exception as exc:
                    raise ValueError(f"ICCS validation failed: {exc}") from exc
        return self

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

    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 = Field(
        default=None,
        foreign_key="aimbateventparameters.id",
        ondelete="CASCADE",
        title="Event parameters ID",
        description="Foreign key referencing the source event parameters.",
    )
    snapshot: "AimbatSnapshot" = Relationship(
        back_populates="event_parameters_snapshot"
    )
    "The snapshot this record belongs to."
    parameters: AimbatEventParameters = Relationship(back_populates="snapshots")
    "The event parameters this snapshot was taken from."

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,
    )

    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", cascade_delete=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: PydanticPositiveTimedelta | 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."""

    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 = Field(
        default=None,
        foreign_key="aimbateventquality.id",
        ondelete="CASCADE",
        title="Event quality ID",
        description="Foreign key referencing the source event quality.",
    )
    quality: AimbatEventQuality = Relationship(back_populates="snapshots")
    "The event quality this snapshot was taken from."
    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",
    )
    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",
    )

    @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.

    At most one of the four FK fields is set per row. 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_at_most_one_parent",
        ),
    )

    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 _at_most_one_parent(self) -> "AimbatNote":
        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(
                "At most one of event_id, station_id, seismogram_id, snapshot_id"
                " may be set on AimbatNote."
            )
        return self

class AimbatSeismogram(SQLModel, table=True):
    """Class to store seismogram metadata."""

    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 seismogram ID.",
        schema_extra={"rich": RichColSpec(style="yellow", highlight=False)},
    )
    begin_time: PydanticTimestamp = Field(
        sa_type=SAPandasTimestamp,
        title="Begin time",
        description="Start time of the seismogram.",
    )
    delta: PydanticPositiveTimedelta = Field(
        sa_type=SAPandasTimedelta,
        title="Sampling interval",
        description="Sampling interval.",
    )
    t0: PydanticTimestamp = Field(
        sa_type=SAPandasTimestamp,
        title="Initial pick",
        description="Initial phase arrival pick.",
    )
    datasource: AimbatDataSource = Relationship(
        back_populates="seismogram", cascade_delete=True
    )
    "Data source for the seismogram waveform."
    station_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatstation.id",
        ondelete="CASCADE",
        title="Station ID",
        description="Foreign key referencing the recording station.",
    )
    station: "AimbatStation" = Relationship(back_populates="seismograms")
    "The station that recorded this seismogram."
    event_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatevent.id",
        ondelete="CASCADE",
        title="Event ID",
        description="Foreign key referencing the parent event.",
    )
    extra: dict[Hashable, Any] = Field(
        default_factory=dict,
        sa_column=Column(MutableDict.as_mutable(PickleType)),
        title="Extra metadata",
        description="Dictionary to store any additional metadata for the seismogram.",
    )
    event: "AimbatEvent" = Relationship(back_populates="seismograms")
    "The event this seismogram belongs to."
    parameters: "AimbatSeismogramParameters" = Relationship(
        back_populates="seismogram",
        cascade_delete=True,
    )
    "Processing parameters for this seismogram."
    quality: AimbatSeismogramQuality | None = Relationship(
        back_populates="seismogram",
        cascade_delete=True,
    )
    "Live quality metrics for this seismogram."

    if TYPE_CHECKING:
        # Add same default values for type checking purposes
        # as in AimbatSeismogramParametersBase
        flip: bool = False
        select: bool = True
        t1: Timestamp | None = None
        data: npt.NDArray[np.float64] = np.array([])

        @property
        def end_time(self) -> Timestamp: ...

    else:

        @computed_field
        def end_time(self) -> PydanticTimestamp:
            """End time of the seismogram, derived from begin_time, delta, and data length."""
            if len(self.data) == 0:
                return self.begin_time
            return self.begin_time + self.delta * (len(self.data) - 1)

        @property
        def flip(self) -> bool:
            """Whether the seismogram should be flipped."""
            return self.parameters.flip

        @flip.setter
        def flip(self, value: bool) -> None:
            self.parameters.flip = value

        @property
        def select(self) -> bool:
            """Whether this seismogram should be used for processing."""
            return self.parameters.select

        @select.setter
        def select(self, value: bool) -> None:
            self.parameters.select = value

        @property
        def t1(self) -> Timestamp | None:
            """Working phase arrival pick."""
            return self.parameters.t1

        @t1.setter
        def t1(self, value: Timestamp | None) -> None:
            self.parameters.t1 = value

        @property
        def data(self) -> npt.NDArray[np.float64]:
            """Seismogram waveform data array."""
            if self.datasource is None:
                raise ValueError("Expected a valid datasource name, got None.")
            return read_seismogram_data(
                self.datasource.sourcename, self.datasource.datatype
            )

        @data.setter
        def data(self, value: npt.NDArray[np.float64]) -> None:
            if self.datasource is None:
                raise ValueError("Expected a valid datasource name, got None.")
            write_seismogram_data(
                self.datasource.sourcename, self.datasource.datatype, value
            )