"""Tests for spike data converters."""

from __future__ import annotations

import numpy as np
import pytest

from radiens_core._converters._spikes import (
    channel_spike_data_from_raster,
    merge_waveforms_channel,
    merge_waveforms_neuron,
    neuron_spike_data_from_raster,
    neurons_result_from_proto,
    spikes_spec_from_proto,
)
from radiens_core._generated import biointerface_pb2, common_pb2, spikesorter_pb2
from radiens_core.models.spikes import NeuronInfo, SpikesSpec

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_spike_timestamp(
    ts: int,
    label: int,
    data_idxs: list[int] | None = None,
) -> biointerface_pb2.SpikeTimestamp:
    return biointerface_pb2.SpikeTimestamp(
        spikeTimestamp=ts,
        label=label,
        dataIdxs=data_idxs if data_idxs is not None else [],
    )


def _make_raster_reply(
    spike_timestamps: list[biointerface_pb2.SpikeTimestamp],
    time_range: list[float],
    ts_range: list[int],
) -> spikesorter_pb2.SpikeSorterRasterDataReply:
    return spikesorter_pb2.SpikeSorterRasterDataReply(
        spikeTimestamps=spike_timestamps,
        timeRange=time_range,
        timeStampRange=ts_range,
    )


# ---------------------------------------------------------------------------
# SpikesSpec converter
# ---------------------------------------------------------------------------


class TestSpikesSpecConverter:
    def test_basic_fields(self) -> None:
        tr = common_pb2.TimeRange(timeRangeSec=[1.0, 5.0])
        spike_tr = common_pb2.TimeRange(timeRangeSec=[1.5, 4.5])
        proto = biointerface_pb2.SpikesSpecReply(
            biointerfaceID="bio-123",
            datasetUID="ds-uid",
            numNeurons=4,
            numSites=2,
            enabledSiteNtvChanIdx=[0, 1],
            tR=tr,
            spikeTR=spike_tr,
        )
        spec = spikes_spec_from_proto(proto)
        assert isinstance(spec, SpikesSpec)
        assert spec.biointerface_id == "bio-123"
        assert spec.dataset_uid == "ds-uid"
        assert spec.num_neurons == 4
        assert spec.num_sites == 2
        assert spec.enabled_channel_indices == [0, 1]
        assert spec.time_range == (1.0, 5.0)
        assert spec.spike_time_range == (1.5, 4.5)

    def test_missing_time_range_defaults_to_zero(self) -> None:
        proto = biointerface_pb2.SpikesSpecReply(
            biointerfaceID="bio-456",
            datasetUID="ds-uid2",
        )
        spec = spikes_spec_from_proto(proto)
        assert spec.time_range == (0.0, 0.0)
        assert spec.spike_time_range == (0.0, 0.0)


# ---------------------------------------------------------------------------
# channel_spike_data_from_raster
# ---------------------------------------------------------------------------


class TestChannelSpikeDataConverter:
    def test_groups_by_channel_label(self) -> None:
        """Spikes are grouped by their label (= channel index in CHANNELS mode)."""
        sts = [
            _make_spike_timestamp(100, label=0),
            _make_spike_timestamp(200, label=1),
            _make_spike_timestamp(300, label=0),
        ]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = channel_spike_data_from_raster(raster)

        assert set(result.keys()) == {0, 1}
        assert len(result[0].timestamps) == 2
        assert len(result[1].timestamps) == 1

    def test_timestamp_conversion(self) -> None:
        """Int64 timestamps are converted linearly to seconds."""
        # ts_range [0, 1000] maps to time_range [0.0, 1.0]
        # ts=500 → 0.5 s
        sts = [_make_spike_timestamp(500, label=0)]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = channel_spike_data_from_raster(raster)
        assert pytest.approx(result[0].timestamps[0]) == 0.5

    def test_labels_array_matches_timestamps(self) -> None:
        sts = [
            _make_spike_timestamp(100, label=0),
            _make_spike_timestamp(200, label=0),
        ]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = channel_spike_data_from_raster(raster)
        # label field for CHANNELS mode = channel index, so both spikes have label 0
        assert list(result[0].labels) == [0, 0]

    def test_waveforms_none_by_default(self) -> None:
        sts = [_make_spike_timestamp(100, label=0)]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = channel_spike_data_from_raster(raster)
        assert result[0].waveforms is None

    def test_empty_raster_returns_empty_dict(self) -> None:
        raster = _make_raster_reply([], time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = channel_spike_data_from_raster(raster)
        assert result == {}

    def test_degenerate_ts_range_returns_zeros(self) -> None:
        """Zero-span timestamp range produces zero timestamps."""
        sts = [_make_spike_timestamp(500, label=0)]
        raster = _make_raster_reply(sts, time_range=[1.0, 1.0], ts_range=[500, 500])
        result = channel_spike_data_from_raster(raster)
        assert result[0].timestamps[0] == pytest.approx(1.0)


# ---------------------------------------------------------------------------
# merge_waveforms_channel
# ---------------------------------------------------------------------------


class TestMergeWaveformsChannel:
    def _make_dense_reply(
        self,
        wfm_matrix: list[list[float]],
    ) -> biointerface_pb2.SpikesSpikeDataDenseReply:
        arr = np.array(wfm_matrix, dtype=np.float32)
        total, n_pts = arr.shape
        return biointerface_pb2.SpikesSpikeDataDenseReply(
            data=arr.tobytes(),
            totalNSpikes=total,
            waveformNPoints=n_pts,
        )

    def test_waveforms_attached_via_data_idxs(self) -> None:
        """dataIdxs links each spike to a row in the dense matrix."""
        # 3 waveform rows, 4 samples each
        wfm = [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]]
        dense = self._make_dense_reply(wfm)

        # chan 0 → spike at row 0, chan 1 → spike at row 2
        sts = [
            _make_spike_timestamp(100, label=0, data_idxs=[0]),
            _make_spike_timestamp(200, label=1, data_idxs=[2]),
        ]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        base_result = channel_spike_data_from_raster(raster)
        merged = merge_waveforms_channel(base_result, raster, dense)

        assert merged[0].waveforms is not None
        np.testing.assert_array_almost_equal(merged[0].waveforms[0], [1.0, 2.0, 3.0, 4.0])
        assert merged[1].waveforms is not None
        np.testing.assert_array_almost_equal(merged[1].waveforms[0], [9.0, 10.0, 11.0, 12.0])

    def test_empty_dense_leaves_waveforms_none(self) -> None:
        sts = [_make_spike_timestamp(100, label=0, data_idxs=[0])]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        base_result = channel_spike_data_from_raster(raster)
        dense = biointerface_pb2.SpikesSpikeDataDenseReply(
            data=b"",
            totalNSpikes=0,
            waveformNPoints=0,
        )
        merged = merge_waveforms_channel(base_result, raster, dense)
        assert merged[0].waveforms is None


# ---------------------------------------------------------------------------
# neuron_spike_data_from_raster
# ---------------------------------------------------------------------------


class TestNeuronSpikeDataConverter:
    def test_groups_by_neuron_label(self) -> None:
        """Spikes are grouped by neuron label (NEURONS mode)."""
        sts = [
            _make_spike_timestamp(100, label=1),
            _make_spike_timestamp(200, label=2),
            _make_spike_timestamp(300, label=1),
        ]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = neuron_spike_data_from_raster(raster, label_to_chan={1: 0, 2: 1})

        assert set(result.keys()) == {1, 2}
        assert len(result[1].timestamps) == 2
        assert len(result[2].timestamps) == 1

    def test_ntv_chan_idx_populated_from_map(self) -> None:
        sts = [_make_spike_timestamp(100, label=3)]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = neuron_spike_data_from_raster(raster, label_to_chan={3: 7})
        assert result[3].ntv_chan_idx == 7

    def test_ntv_chan_idx_defaults_to_minus_one_if_not_found(self) -> None:
        sts = [_make_spike_timestamp(100, label=99)]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = neuron_spike_data_from_raster(raster, label_to_chan={})
        assert result[99].ntv_chan_idx == -1

    def test_waveforms_none_by_default(self) -> None:
        sts = [_make_spike_timestamp(100, label=1)]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        result = neuron_spike_data_from_raster(raster, label_to_chan={1: 0})
        assert result[1].waveforms is None


class TestMergeWaveformsNeuron:
    def _make_dense_reply(
        self,
        wfm_matrix: list[list[float]],
    ) -> biointerface_pb2.SpikesSpikeDataDenseReply:
        arr = np.array(wfm_matrix, dtype=np.float32)
        total, n_pts = arr.shape
        return biointerface_pb2.SpikesSpikeDataDenseReply(
            data=arr.tobytes(),
            totalNSpikes=total,
            waveformNPoints=n_pts,
        )

    def test_waveforms_attached_by_neuron_label(self) -> None:
        wfm = [[1.0, 2.0], [3.0, 4.0]]
        dense = self._make_dense_reply(wfm)

        sts = [
            _make_spike_timestamp(100, label=1, data_idxs=[0]),
            _make_spike_timestamp(200, label=2, data_idxs=[1]),
        ]
        raster = _make_raster_reply(sts, time_range=[0.0, 1.0], ts_range=[0, 1000])
        base_result = neuron_spike_data_from_raster(raster, label_to_chan={1: 0, 2: 0})
        merged = merge_waveforms_neuron(base_result, raster, dense)

        assert merged[1].waveforms is not None
        np.testing.assert_array_almost_equal(merged[1].waveforms[0], [1.0, 2.0])
        assert merged[2].waveforms is not None
        np.testing.assert_array_almost_equal(merged[2].waveforms[0], [3.0, 4.0])


# ---------------------------------------------------------------------------
# NeuronsResult converter
# ---------------------------------------------------------------------------


class TestNeuronsResultConverter:
    def test_neurons_by_channel(self) -> None:
        nd = biointerface_pb2.NeuronDescriptor(
            neuronID="n1",
            siteNtvChanIdx=3,
            spikeLabel="1",
            spikeCount=100,
            spikeRate=10.0,
            meanAbsPeakWfm=50.0,
            snr=3.5,
            posProbe=[1.0, 2.0, 3.0],
            posTissue=[4.0, 5.0, 6.0],
        )
        neuron_l = biointerface_pb2.BiointerfaceGetNeuronsReply.NeuronL(neuron=[nd])
        reply = biointerface_pb2.BiointerfaceGetNeuronsReply(site={3: neuron_l})

        result = neurons_result_from_proto(reply)

        assert 3 in result.neurons_by_channel
        infos = result.neurons_by_channel[3]
        assert len(infos) == 1
        info = infos[0]
        assert isinstance(info, NeuronInfo)
        assert info.neuron_id == "n1"
        assert info.site_ntv_chan_idx == 3
        assert info.spike_label == "1"
        assert info.spike_count == 100
        assert info.spike_rate == pytest.approx(10.0)
        assert info.snr == pytest.approx(3.5)
        assert info.pos_probe == [1.0, 2.0, 3.0]
        assert info.pos_tissue == [4.0, 5.0, 6.0]

    def test_empty_reply_returns_empty_dict(self) -> None:
        reply = biointerface_pb2.BiointerfaceGetNeuronsReply()
        result = neurons_result_from_proto(reply)
        assert result.neurons_by_channel == {}
