"""Channel metadata, site position models, and channel-related enums."""

from __future__ import annotations

from enum import IntEnum
from typing import assert_never

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

from radiens_core.exceptions import DatasetError
from radiens_core.models.signals import SignalType


class SignalUnits(IntEnum):
    """Signal measurement units."""

    UNKNOWN = 0
    MICROVOLTS = 1
    VOLTS = 2
    BINARY = 3


class AreaUnits(IntEnum):
    """Area measurement units for electrode sites."""

    SQUARE_MICROMETERS = 0  # µm² (base unit)
    SQUARE_MILLIMETERS = 1  # mm²
    SQUARE_CENTIMETERS = 2  # cm²
    SQUARE_NANOMETERS = 3  # nm²
    SQUARE_METERS = 4  # m²


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

    Attributes:
        probe_x: X-coordinate in probe coordinate system.
        probe_y: Y-coordinate in probe coordinate system.
        probe_z: Z-coordinate in probe coordinate system.
        tissue_x: X-coordinate in tissue coordinate system.
        tissue_y: Y-coordinate in tissue coordinate system.
        tissue_z: Z-coordinate in tissue coordinate system.
        site_shape: Geometric shape of the electrode site.
        site_area_um2: Surface area of the electrode site in square micrometers.
        site_lim_x_min: Minimum X-limit of the site bounding box.
        site_lim_x_max: Maximum X-limit of the site bounding box.
        site_lim_y_min: Minimum Y-limit of the site bounding box.
        site_lim_y_max: Maximum Y-limit of the site bounding box.
        site_lim_z_min: Minimum Z-limit of the site bounding box.
        site_lim_z_max: Maximum Z-limit of the site bounding box.
        scs_to_tcs_x: X-offset for probe to tissue coordinate transformation.
        scs_to_tcs_y: Y-offset for probe to tissue coordinate transformation.
        scs_to_tcs_z: Z-offset for probe to tissue coordinate transformation.
    """

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

    # Geometry metadata
    site_shape: str
    site_area_um2: float

    # Bounding box limits (probe coordinates)
    site_lim_x_min: float
    site_lim_x_max: float
    site_lim_y_min: float
    site_lim_y_max: float
    site_lim_z_min: float
    site_lim_z_max: float

    # Coordinate system transformation offsets
    scs_to_tcs_x: float
    scs_to_tcs_y: float
    scs_to_tcs_z: float

    @property
    def bounding_box_size(self) -> tuple[float, float, float]:
        """Site bounding box size (width, height, depth) in micrometers."""
        return (
            self.site_lim_x_max - self.site_lim_x_min,
            self.site_lim_y_max - self.site_lim_y_min,
            self.site_lim_z_max - self.site_lim_z_min,
        )

    @property
    def bounding_box_center(self) -> tuple[float, float, float]:
        """Site bounding box center coordinates in probe coordinate system."""
        return (
            (self.site_lim_x_min + self.site_lim_x_max) / 2,
            (self.site_lim_y_min + self.site_lim_y_max) / 2,
            (self.site_lim_z_min + self.site_lim_z_max) / 2,
        )

    def get_area(self, unit: AreaUnits) -> float:
        """Get site area in the specified units.

        Args:
            unit: Target area unit for conversion

        Returns:
            Site area converted to the specified unit
        """
        # Base area is in µm²
        base_area_um2 = self.site_area_um2

        if unit == AreaUnits.SQUARE_MICROMETERS:
            return base_area_um2
        elif unit == AreaUnits.SQUARE_MILLIMETERS:
            # 1 mm² = 1,000,000 µm²
            return base_area_um2 / 1_000_000
        elif unit == AreaUnits.SQUARE_CENTIMETERS:
            # 1 cm² = 100,000,000 µm²
            return base_area_um2 / 100_000_000
        elif unit == AreaUnits.SQUARE_NANOMETERS:
            # 1 µm² = 1,000,000 nm²
            return base_area_um2 * 1_000_000
        elif unit == AreaUnits.SQUARE_METERS:
            # 1 m² = 1,000,000,000,000 µm²
            return base_area_um2 / 1_000_000_000_000
        else:
            assert_never(unit)


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

    Attributes:
        ntv: Native channel indices.
        dset: Dataset channel indices.
        sys: System channel indices.
    """

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


class ChannelInfo(BaseModel):
    """Complete channel information for a specific channel.

    Attributes:
        ntv_idx: Native channel index.
        dset_idx: Dataset channel index.
        sys_idx: System channel index.
        signal_type: Type of the signal. References SignalType.
        site_position: Site position and geometry. Defaults to None.
    """

    model_config = ConfigDict(frozen=True)
    ntv_idx: int
    dset_idx: int
    sys_idx: int
    signal_type: SignalType
    site_position: SitePosition | None = None


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

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

    Attributes:
        dataset_uid: Unique identifier for the dataset.
        fs: Sampling frequency in Hertz.
        source_label: Label for the data source.
        neighbor_max_radius_um: Maximum radius for neighbor search in micrometers.
        signal_units: Mapping of signal types to their units. References SignalType and SignalUnits.
        has_discrete: Whether the dataset contains discrete signals.
        has_continuous_amp: Whether the dataset contains continuous amplifier signals.
        sig_map: Mapping of signal types to their indices. References SignalType and KeyIdxs.
        selected_sig_map: Mapping of signal types to their selected indices. References SignalType and KeyIdxs.
        num_channels: Total number of channels per signal type. References SignalType.
        selected_num_channels: Number of selected channels per signal type. References SignalType.
        site_positions: List of site positions for amplifier 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)

    def get_channel_info(
        self,
        ntv_idx: int,
        sig_type: SignalType,
    ) -> ChannelInfo:
        """Get complete channel information from native index.

        Args:
            ntv_idx: Native channel index
            sig_type: Signal type to look up

        Returns:
            Complete channel information including indices and site position

        Raises:
            DatasetError: If ntv_idx not found for the signal type
        """
        key_idxs = self.sig_map[sig_type]

        try:
            position = key_idxs.ntv.index(ntv_idx)
        except ValueError:
            msg = f"NTV index {ntv_idx} not found for signal type {sig_type}"
            raise DatasetError(msg) from None

        site_position = None
        if sig_type == SignalType.AMP and position < len(self.site_positions):
            site_position = self.site_positions[position]

        return ChannelInfo(
            ntv_idx=ntv_idx,
            dset_idx=key_idxs.dset[position],
            sys_idx=key_idxs.sys[position],
            signal_type=sig_type,
            site_position=site_position,
        )

    def ntv_from_dset(self, dset_idx: int, sig_type: SignalType) -> int:
        """Convert dataset index to native index.

        Args:
            dset_idx: Dataset index
            sig_type: Signal type

        Returns:
            Native index for the channel

        Raises:
            DatasetError: If dset_idx not found for the signal type
        """
        key_idxs = self.sig_map[sig_type]
        try:
            position = key_idxs.dset.index(dset_idx)
            return key_idxs.ntv[position]
        except ValueError:
            msg = f"Dataset index {dset_idx} not found for signal type {sig_type}"
            raise DatasetError(msg) from None

    def ntv_from_sys(self, sys_idx: int, sig_type: SignalType) -> int:
        """Convert system index to native index.

        Args:
            sys_idx: System index
            sig_type: Signal type

        Returns:
            Native index for the channel

        Raises:
            DatasetError: If sys_idx not found for the signal type
        """
        key_idxs = self.sig_map[sig_type]
        try:
            position = key_idxs.sys.index(sys_idx)
            return key_idxs.ntv[position]
        except ValueError:
            msg = f"System index {sys_idx} not found for signal type {sig_type}"
            raise DatasetError(msg) from None

    @property
    def amp_site_count(self) -> int:
        """Number of amplifier channel sites."""
        return len(self.site_positions)

    @property
    def total_selected_channels(self) -> int:
        """Total number of selected channels across all signal types."""
        return sum(self.selected_num_channels.values())
