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):
    """Class to store seismic event information."""

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

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    time: PydanticTimestamp = Field(
        unique=True, sa_type=SAPandasTimestamp, allow_mutation=False
    )
    latitude: float
    longitude: float
    depth: float | None = None
    last_modified: PydanticTimestamp | None = Field(
        default=None, sa_type=SAPandasTimestamp
    )
    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 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)
    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
            )