aimbat.models

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

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    sourcename: str = Field(description="Path or name of the data source.")
    datatype: DataType = Field(description="Data type of the data source.")
    seismogram_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatseismogram.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent seismogram.",
    )
    seismogram: "AimbatSeismogram" = Relationship(back_populates="datasource")
    "The seismogram this data source belongs to."

class AimbatEvent(SQLModel, table=True):
    """Class to store seismic event information."""

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )

    is_default: bool | None = Field(
        default=None,
        unique=True,
        description="Indicates if an event is the default event.",
    )

    time: PydanticTimestamp = Field(
        unique=True,
        sa_type=SAPandasTimestamp,
        allow_mutation=False,
        description="Event time.",
    )

    latitude: float = Field(description="Event latitude.")
    longitude: float = Field(description="Event longitude.")
    depth: float | None = Field(default=None, description="Event depth.")
    last_modified: PydanticTimestamp | None = Field(
        default=None,
        sa_type=SAPandasTimestamp,
        description="Timestamp of the last modification to this event's parameters.",
    )
    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."

    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

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

    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",
        description="Foreign key referencing the parent event.",
    )
    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.",
    )

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

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

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

    mccc_min_ccnorm: float = Field(
        default_factory=lambda: settings.mccc_min_ccnorm,
        ge=0,
        le=1,
        title="MCCC min CC norm",
        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

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

    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",
        description="Foreign key referencing the parent snapshot.",
    )
    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."
    parameters_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbateventparameters.id",
        ondelete="CASCADE",
        description="Foreign key referencing the source event parameters.",
    )

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

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    begin_time: PydanticTimestamp = Field(
        sa_type=SAPandasTimestamp, description="Start time of the seismogram."
    )
    delta: PydanticPositiveTimedelta = Field(
        sa_type=SAPandasTimedelta, description="Sampling interval."
    )
    t0: PydanticTimestamp = Field(
        sa_type=SAPandasTimestamp, 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",
        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",
        description="Foreign key referencing the parent event.",
    )
    extra: dict[Hashable, Any] = Field(
        default_factory=dict,
        sa_column=Column(MutableDict.as_mutable(PickleType)),
        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."

    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
            )

class AimbatSeismogramParameters(AimbatSeismogramParametersBase, table=True):
    """Processing parameters for a single seismogram."""

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    seismogram_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatseismogram.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent seismogram.",
    )
    seismogram: "AimbatSeismogram" = Relationship(back_populates="parameters")
    "The seismogram these parameters belong to."
    snapshots: list["AimbatSeismogramParametersSnapshot"] = Relationship(
        back_populates="parameters", cascade_delete=True
    )
    "Parameter snapshots for this seismogram."

class AimbatSeismogramParametersBase(SQLModel):
    """Base class defining seismogram-level processing parameters for AIMBAT."""

    flip: bool = Field(
        default=False,
        description="Whether or not the seismogram should be flipped.",
    )

    select: bool = Field(
        default=True,
        description="Whether or not this seismogram should be used for processing.",
    )

    t1: PydanticTimestamp | None = Field(
        default=None,
        sa_type=SAPandasTimestamp,
        description=(
            "Working pick. This pick serves as working as well as output pick."
            " It is changed by: 1. Picking the phase arrival in the stack,"
            " 2. Running ICCS, 3. Running MCCC."
        ),
    )

class AimbatSeismogramParametersSnapshot(AimbatSeismogramParametersBase, table=True):
    """Snapshot of processing parameters for a single seismogram."""

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    seismogram_parameters_id: uuid.UUID = Field(
        foreign_key="aimbatseismogramparameters.id",
        ondelete="CASCADE",
        description="Foreign key referencing the source seismogram parameters.",
    )
    parameters: AimbatSeismogramParameters = Relationship(back_populates="snapshots")
    "The seismogram parameters this snapshot was taken from."
    snapshot_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatsnapshot.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent snapshot.",
    )
    snapshot: "AimbatSnapshot" = Relationship(
        back_populates="seismogram_parameters_snapshots"
    )
    "The snapshot this record belongs to."

class AimbatSnapshot(SQLModel, table=True):
    """Container for a point-in-time snapshot of event and seismogram parameters.

    The AimbatSnapshot class does not actually save any parameter data.
    It is used to keep track of the AimbatEventParametersSnapshot and
    AimbatSeismogramParametersSnapshot instances.
    """

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    date: PydanticTimestamp = Field(
        default_factory=lambda: Timestamp.now(tz=timezone.utc),
        unique=True,
        allow_mutation=False,
        sa_type=SAPandasTimestamp,
        description="Timestamp when the snapshot was created.",
    )
    comment: str | None = Field(
        default=None, description="Optional comment for the snapshot."
    )
    event_parameters_snapshot: AimbatEventParametersSnapshot = Relationship(
        back_populates="snapshot", cascade_delete=True
    )
    "Event parameter snapshot associated with this snapshot."
    seismogram_parameters_snapshots: list[AimbatSeismogramParametersSnapshot] = (
        Relationship(back_populates="snapshot", cascade_delete=True)
    )
    "Seismogram parameter snapshots associated with this snapshot."
    event_id: uuid.UUID = Field(
        default=None,
        foreign_key="aimbatevent.id",
        ondelete="CASCADE",
        description="Foreign key referencing the parent event.",
    )
    event: "AimbatEvent" = Relationship(back_populates="snapshots")
    "The event this snapshot belongs to."

    if TYPE_CHECKING:
        # defined as column properties below, but add same default values for type checking purposes
        seismogram_count: int = 0
        selected_seismogram_count: int = 0
        flipped_seismogram_count: int = 0

class AimbatStation(SQLModel, table=True):
    """Class to store station information."""

    id: uuid.UUID = Field(
        default_factory=uuid.uuid4, primary_key=True, description="Unique ID."
    )
    name: str = Field(allow_mutation=False, description="Station name.")
    network: str = Field(allow_mutation=False, description="Network name.")
    location: str = Field(allow_mutation=False, description="Location ID.")
    channel: str = Field(allow_mutation=False, description="Channel code.")
    latitude: float = Field(description="Station latitude.")
    longitude: float = Field(description="Station longitude.")
    elevation: float | None = Field(default=None, description="Station elevation.")
    seismograms: list[AimbatSeismogram] = Relationship(
        back_populates="station", cascade_delete=True
    )
    "Seismograms recorded at this station."