"""Spike sorting and neural recording domain models."""

from __future__ import annotations

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


class SpikesSpec(BaseModel):
    """Spike interface metadata for a linked dataset.

    Attributes:
        biointerface_id: Unique identifier for the spike interface; used as
            ``spikeSorterID`` in RPC calls that require it.
        dataset_uid: Dataset UID this spike interface is associated with.
        num_neurons: Total number of sorted neurons across all channels.
        num_sites: Number of recording sites (channels) with spike data.
        enabled_channel_indices: Native channel indices of enabled recording sites,
            sorted in ascending order.
        time_range: Full time range of the recording in seconds ``(start, end)``.
        spike_time_range: Time range over which spike data is available in seconds
            ``(start, end)``.
    """

    model_config = ConfigDict(frozen=True)

    biointerface_id: str
    dataset_uid: str
    num_neurons: int
    num_sites: int
    enabled_channel_indices: list[int]
    time_range: tuple[float, float]
    spike_time_range: tuple[float, float]


class ChannelSpikeData(BaseModel):
    """Spike timestamps and neuron labels for a single recording channel.

    Attributes:
        ntv_chan_idx: Native channel index.
        timestamps: Spike timestamps in seconds, shape ``(N,)``.
        labels: Integer neuron label per spike, shape ``(N,)``. Matches the
            neuron label reported by ``get_neurons()``.
        waveforms: Spike waveforms in raw units, shape ``(N, waveform_n_points)``,
            or ``None`` if waveforms were not requested.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    ntv_chan_idx: int
    timestamps: npt.NDArray[np.float64]
    labels: npt.NDArray[np.int32]
    waveforms: npt.NDArray[np.float32] | None


class NeuronSpikeData(BaseModel):
    """Spike timestamps for a single sorted neuron.

    Attributes:
        label: Integer neuron label. Correlates with ``NeuronInfo.spike_label``
            (string) from ``get_neurons()``.
        ntv_chan_idx: Native channel index of the primary recording site for this
            neuron. ``-1`` if the channel could not be determined from the spike spec.
        timestamps: Spike timestamps in seconds, shape ``(N,)``.
        waveforms: Spike waveforms in raw units, shape ``(N, waveform_n_points)``,
            or ``None`` if waveforms were not requested.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    label: int
    ntv_chan_idx: int
    timestamps: npt.NDArray[np.float64]
    waveforms: npt.NDArray[np.float32] | None


class NeuronInfo(BaseModel):
    """Descriptive statistics for a single sorted neuron.

    Attributes:
        neuron_id: Unique neuron identifier string.
        site_ntv_chan_idx: Native channel index of the primary recording site.
        spike_label: Neuron label as a string (e.g. ``"1"``, ``"2"``).
        spike_count: Total number of spikes detected.
        spike_rate: Mean spike rate in Hz.
        mean_abs_peak_wfm: Mean absolute peak waveform amplitude.
        snr: Signal-to-noise ratio.
        pos_probe: Probe coordinate ``[x, y, z]`` in microns.
        pos_tissue: Tissue coordinate ``[x, y, z]`` in microns.
    """

    model_config = ConfigDict(frozen=True)

    neuron_id: str
    site_ntv_chan_idx: int
    spike_label: str
    spike_count: int
    spike_rate: float
    mean_abs_peak_wfm: float
    snr: float
    pos_probe: list[float]
    pos_tissue: list[float]


class NeuronsResult(BaseModel):
    """Neuron descriptors organized by recording channel.

    Attributes:
        neurons_by_channel: Map from native channel index to a list of neuron
            descriptors detected on that channel.
    """

    model_config = ConfigDict(frozen=True)

    neurons_by_channel: dict[int, list[NeuronInfo]]
