"""Channel metadata and site position models."""

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel, ConfigDict

from radiens_core.models.enums import SignalType, SignalUnits


class SitePosition(BaseModel):
    """Electrode site position in probe and tissue coordinates."""

    model_config = ConfigDict(frozen=True)
    probe_x: float
    probe_y: float
    probe_z: float
    tissue_x: float
    tissue_y: float
    tissue_z: float


class KeyIdxs(BaseModel):
    """Channel indices in all three key index spaces."""

    model_config = ConfigDict(frozen=True)
    ntv: list[int]
    dset: list[int]
    sys: list[int]


class ChannelMetadata(BaseModel):
    """Channel metadata for a dataset.

    Contains channel configuration, probe geometry, signal types,
    and indices for all channels.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)
    dataset_uid: str
    fs: float
    source_label: str
    neighbor_max_radius_um: float
    signal_units: dict[SignalType, SignalUnits]
    has_discrete: bool
    has_continuous_amp: bool

    # Channel indices per signal type
    sig_map: dict[SignalType, KeyIdxs]
    selected_sig_map: dict[SignalType, KeyIdxs]
    num_channels: dict[SignalType, int]
    selected_num_channels: dict[SignalType, int]

    # Site positions (one per amp channel, native index order)
    site_positions: list[SitePosition]

    def index(self, sig_type: SignalType) -> KeyIdxs:
        """Get channel indices for a signal type."""
        return self.sig_map[sig_type]

    def selected_index(self, sig_type: SignalType) -> KeyIdxs:
        """Get selected channel indices for a signal type."""
        return self.selected_sig_map[sig_type]

    def get_num_channels(self, sig_type: SignalType) -> int:
        """Get total number of channels for a signal type."""
        return self.num_channels[sig_type]

    def all_amp_indices(self) -> npt.NDArray[np.int64]:
        """Get all amplifier channel native indices."""
        return np.array(self.sig_map[SignalType.AMP].ntv, dtype=np.int64)
