"""Spike data proto <-> domain converters."""

from __future__ import annotations

from collections import defaultdict
from typing import TYPE_CHECKING

import numpy as np
import numpy.typing as npt

from radiens_core.models.spikes import (
    ChannelSpikeData,
    NeuronInfo,
    NeuronSpikeData,
    NeuronsResult,
    SpikesSpec,
)

if TYPE_CHECKING:
    from radiens_core._generated import biointerface_pb2, spikesorter_pb2

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _timestamps_to_sec(
    ts_ints: list[int],
    time_range: list[float],
    ts_range: list[int],
) -> npt.NDArray[np.float64]:
    """Convert a list of hardware int64 timestamps to seconds.

    Uses linear interpolation anchored by the ``timeRange`` and
    ``timeStampRange`` fields from the raster reply.  Returns an array of
    zeros if the ranges are degenerate (empty or zero span).
    """
    n = len(ts_ints)
    if n == 0:
        return np.empty(0, dtype=np.float64)
    if len(time_range) < 2 or len(ts_range) < 2:
        return np.zeros(n, dtype=np.float64)
    span_ts = ts_range[1] - ts_range[0]
    if span_ts == 0:
        return np.full(n, time_range[0], dtype=np.float64)
    ts_arr = np.array(ts_ints, dtype=np.int64)
    scale = (time_range[1] - time_range[0]) / span_ts
    return np.asarray(time_range[0] + (ts_arr - ts_range[0]) * scale, dtype=np.float64)


def _collect_dataIdxs_by_label(
    raster: spikesorter_pb2.SpikeSorterRasterDataReply,
) -> dict[int, list[int]]:
    """Build a mapping from label → list of waveform row indices.

    Each ``SpikeTimestamp.dataIdxs`` contains the row index into the dense
    waveform matrix (``SpikesGetSpikesDense``) for that spike.  Takes the
    first element of ``dataIdxs`` per spike (the primary waveform row).
    """
    by_label: dict[int, list[int]] = defaultdict(list)
    for st in raster.spikeTimestamps:
        if len(st.dataIdxs) > 0:
            by_label[st.label].append(st.dataIdxs[0])
    return dict(by_label)


# ---------------------------------------------------------------------------
# SpikesSpec
# ---------------------------------------------------------------------------


def spikes_spec_from_proto(
    proto: biointerface_pb2.SpikesSpecReply,
) -> SpikesSpec:
    """Convert ``SpikesSpecReply`` proto to domain ``SpikesSpec``."""
    if proto.HasField("tR") and len(proto.tR.timeRangeSec) >= 2:
        time_range: tuple[float, float] = (
            proto.tR.timeRangeSec[0],
            proto.tR.timeRangeSec[1],
        )
    else:
        time_range = (0.0, 0.0)

    if proto.HasField("spikeTR") and len(proto.spikeTR.timeRangeSec) >= 2:
        spike_time_range: tuple[float, float] = (
            proto.spikeTR.timeRangeSec[0],
            proto.spikeTR.timeRangeSec[1],
        )
    else:
        spike_time_range = (0.0, 0.0)

    return SpikesSpec(
        biointerface_id=proto.biointerfaceID,
        dataset_uid=proto.datasetUID,
        num_neurons=proto.numNeurons,
        num_sites=proto.numSites,
        enabled_channel_indices=list(proto.enabledSiteNtvChanIdx),
        time_range=time_range,
        spike_time_range=spike_time_range,
    )


# ---------------------------------------------------------------------------
# Channel spike data (RasterMode.CHANNELS)
# ---------------------------------------------------------------------------


def channel_spike_data_from_raster(
    proto: spikesorter_pb2.SpikeSorterRasterDataReply,
) -> dict[int, ChannelSpikeData]:
    """Convert a CHANNELS-mode raster reply to per-channel spike data.

    Groups ``SpikeTimestamp`` objects by ``label`` (= native channel index in
    CHANNELS mode).  Converts hardware int64 timestamps to seconds using the
    linear interpolation range embedded in the reply.  Waveforms are always
    ``None``; call :func:`merge_waveforms_channel` to attach them.

    Args:
        proto: ``SpikeSorterRasterDataReply`` from a CHANNELS-mode request.

    Returns:
        Dict mapping native channel index to :class:`ChannelSpikeData`.
    """
    time_range = list(proto.timeRange)
    ts_range = list(proto.timeStampRange)

    # Group spike timestamps and labels by channel label
    chan_ts: dict[int, list[int]] = defaultdict(list)
    chan_labels: dict[int, list[int]] = defaultdict(list)
    for st in proto.spikeTimestamps:
        chan_ts[st.label].append(st.spikeTimestamp)
        chan_labels[st.label].append(st.label)

    result: dict[int, ChannelSpikeData] = {}
    for chan_idx, ts_ints in chan_ts.items():
        timestamps = _timestamps_to_sec(ts_ints, time_range, ts_range)
        labels = np.array(chan_labels[chan_idx], dtype=np.int32)
        result[chan_idx] = ChannelSpikeData(
            ntv_chan_idx=chan_idx,
            timestamps=timestamps,
            labels=labels,
            waveforms=None,
        )
    return result


def merge_waveforms_channel(
    result: dict[int, ChannelSpikeData],
    raster: spikesorter_pb2.SpikeSorterRasterDataReply,
    dense: biointerface_pb2.SpikesSpikeDataDenseReply,
) -> dict[int, ChannelSpikeData]:
    """Attach waveforms to channel spike data using the dense waveform matrix.

    Uses ``SpikeTimestamp.dataIdxs`` from the raster reply to index into the
    float32 dense matrix (shape ``[totalNSpikes, waveformNPoints]``).

    Args:
        result: Channel spike data returned by :func:`channel_spike_data_from_raster`.
        raster: The same ``SpikeSorterRasterDataReply`` used to build *result*.
        dense: ``SpikesSpikeDataDenseReply`` from ``SpikesGetSpikesDense``.

    Returns:
        New dict with waveforms filled in for each channel present in both
        *result* and the dense reply.  Channels with no waveform rows are
        returned unchanged (waveforms remain ``None``).
    """
    total = dense.totalNSpikes
    n_pts = dense.waveformNPoints
    if total == 0 or n_pts == 0 or len(dense.data) == 0:
        return result

    wfm = np.frombuffer(dense.data, dtype=np.float32).reshape(total, n_pts)
    by_label = _collect_dataIdxs_by_label(raster)

    updated: dict[int, ChannelSpikeData] = {}
    for chan_idx, csd in result.items():
        row_indices = by_label.get(chan_idx)
        if row_indices and len(row_indices) > 0:
            waveforms: npt.NDArray[np.float32] | None = wfm[np.array(row_indices, dtype=np.intp)]
        else:
            waveforms = None
        updated[chan_idx] = ChannelSpikeData(
            ntv_chan_idx=csd.ntv_chan_idx,
            timestamps=csd.timestamps,
            labels=csd.labels,
            waveforms=waveforms,
        )
    return updated


# ---------------------------------------------------------------------------
# Neuron spike data (RasterMode.NEURONS)
# ---------------------------------------------------------------------------


def neuron_spike_data_from_raster(
    proto: spikesorter_pb2.SpikeSorterRasterDataReply,
    label_to_chan: dict[int, int],
) -> dict[int, NeuronSpikeData]:
    """Convert a NEURONS-mode raster reply to per-neuron spike data.

    Groups ``SpikeTimestamp`` objects by ``label`` (= integer neuron label in
    NEURONS mode).  Converts hardware int64 timestamps to seconds.  The
    parent channel for each neuron is looked up in *label_to_chan*; if the
    label is not found, ``ntv_chan_idx`` is set to ``-1``.

    Args:
        proto: ``SpikeSorterRasterDataReply`` from a NEURONS-mode request.
        label_to_chan: Mapping of integer neuron label → native channel index,
            built from ``SpikesSpecReply.site`` by the caller.

    Returns:
        Dict mapping integer neuron label to :class:`NeuronSpikeData`.
    """
    time_range = list(proto.timeRange)
    ts_range = list(proto.timeStampRange)

    neuron_ts: dict[int, list[int]] = defaultdict(list)
    for st in proto.spikeTimestamps:
        neuron_ts[st.label].append(st.spikeTimestamp)

    result: dict[int, NeuronSpikeData] = {}
    for label, ts_ints in neuron_ts.items():
        timestamps = _timestamps_to_sec(ts_ints, time_range, ts_range)
        result[label] = NeuronSpikeData(
            label=label,
            ntv_chan_idx=label_to_chan.get(label, -1),
            timestamps=timestamps,
            waveforms=None,
        )
    return result


def merge_waveforms_neuron(
    result: dict[int, NeuronSpikeData],
    raster: spikesorter_pb2.SpikeSorterRasterDataReply,
    dense: biointerface_pb2.SpikesSpikeDataDenseReply,
) -> dict[int, NeuronSpikeData]:
    """Attach waveforms to neuron spike data using the dense waveform matrix.

    Uses ``SpikeTimestamp.dataIdxs`` from the raster reply to index into the
    float32 dense matrix (shape ``[totalNSpikes, waveformNPoints]``).

    Args:
        result: Neuron spike data returned by :func:`neuron_spike_data_from_raster`.
        raster: The same ``SpikeSorterRasterDataReply`` used to build *result*.
        dense: ``SpikesSpikeDataDenseReply`` from ``SpikesGetSpikesDense``.

    Returns:
        New dict with waveforms filled in for each neuron present in both
        *result* and the dense reply.  Neurons with no waveform rows are
        returned unchanged (waveforms remain ``None``).
    """
    total = dense.totalNSpikes
    n_pts = dense.waveformNPoints
    if total == 0 or n_pts == 0 or len(dense.data) == 0:
        return result

    wfm = np.frombuffer(dense.data, dtype=np.float32).reshape(total, n_pts)
    by_label = _collect_dataIdxs_by_label(raster)

    updated: dict[int, NeuronSpikeData] = {}
    for label, nsd in result.items():
        row_indices = by_label.get(label)
        if row_indices and len(row_indices) > 0:
            waveforms: npt.NDArray[np.float32] | None = wfm[np.array(row_indices, dtype=np.intp)]
        else:
            waveforms = None
        updated[label] = NeuronSpikeData(
            label=nsd.label,
            ntv_chan_idx=nsd.ntv_chan_idx,
            timestamps=nsd.timestamps,
            waveforms=waveforms,
        )
    return updated


# ---------------------------------------------------------------------------
# NeuronInfo / NeuronsResult
# ---------------------------------------------------------------------------


def neuron_info_from_proto(
    pb: biointerface_pb2.NeuronDescriptor,
) -> NeuronInfo:
    """Convert a ``NeuronDescriptor`` proto to domain ``NeuronInfo``."""
    return NeuronInfo(
        neuron_id=pb.neuronID,
        site_ntv_chan_idx=pb.siteNtvChanIdx,
        spike_label=pb.spikeLabel,
        spike_count=pb.spikeCount,
        spike_rate=pb.spikeRate,
        mean_abs_peak_wfm=pb.meanAbsPeakWfm,
        snr=pb.snr,
        pos_probe=list(pb.posProbe),
        pos_tissue=list(pb.posTissue),
    )


def neurons_result_from_proto(
    pb: biointerface_pb2.BiointerfaceGetNeuronsReply,
) -> NeuronsResult:
    """Convert ``BiointerfaceGetNeuronsReply`` proto to domain ``NeuronsResult``.

    The proto ``site`` field is a ``MessageMap[int, NeuronL]`` keyed by
    native channel index.  Each ``NeuronL.neuron`` list contains the neuron
    descriptors for that channel.
    """
    neurons_by_channel: dict[int, list[NeuronInfo]] = {}
    for ntv_chan_idx, neuron_l in pb.site.items():
        neurons_by_channel[ntv_chan_idx] = [neuron_info_from_proto(nd) for nd in neuron_l.neuron]
    return NeuronsResult(neurons_by_channel=neurons_by_channel)
