"""
SpikeCascade: three-stage fidelity cascade for real-time spike sorting.

Stage 1: amplitude threshold (cheap, ~1 us).
Stage 2: CNN template match (medium, ~50 us).
Stage 3: full KMeans clustering (expensive, ~1 ms, ~1% of spikes).

The CascadeReport returned by `infer` implements the calibration
contract from the paper: it reports activity_fraction, stage3_fraction,
latency_ms, and correlation (NaN if the baseline is empty).
"""

from __future__ import annotations

import math
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel

from configuration_spikecascade import SpikeCascadeConfig
from preprocessing import bandpass_filter, extract_snippets


# ─────────────────────────────────────────────────────────────────────
# Stage 2: CNN template matcher
# ─────────────────────────────────────────────────────────────────────

class TemplateCNN(nn.Module):
    """
    A small 1D CNN that maps a multi-channel spike snippet to a
    probability distribution over a bank of learned templates.

    Input:  (batch, n_channels, template_len)
    Output: (batch, template_bank_size) logits, (batch, hidden_dim)
    """

    def __init__(self, cfg: SpikeCascadeConfig):
        super().__init__()
        c1, c2, c3 = cfg.cnn_channels
        self.conv1 = nn.Conv1d(cfg.n_channels, c1, kernel_size=7, padding=3)
        self.conv2 = nn.Conv1d(c1, c2, kernel_size=5, padding=2)
        self.conv3 = nn.Conv1d(c2, c3, kernel_size=3, padding=1)
        self.pool = nn.AdaptiveAvgPool1d(1)
        self.fc = nn.Linear(c3, cfg.hidden_dim)
        self.head = nn.Linear(cfg.hidden_dim, cfg.template_bank_size)

    def forward(self, x: torch.Tensor):
        h = F.relu(self.conv1(x))
        h = F.relu(self.conv2(h))
        h = F.relu(self.conv3(h))
        h = self.pool(h).squeeze(-1)
        h = F.relu(self.fc(h))
        return self.head(h), h


# ─────────────────────────────────────────────────────────────────────
# Cascade report
# ─────────────────────────────────────────────────────────────────────

@dataclass
class CascadeReport:
    units: Dict[int, List[float]] = field(default_factory=dict)
    activity_fraction: float = 0.0
    stage3_fraction: float = 0.0
    latency_ms: float = 0.0
    correlation: float = float("nan")
    warning: Optional[str] = None


# ─────────────────────────────────────────────────────────────────────
# The model
# ─────────────────────────────────────────────────────────────────────

class SpikeCascadeForSorting(PreTrainedModel):
    config_class = SpikeCascadeConfig

    def __init__(self, config: SpikeCascadeConfig):
        super().__init__(config)
        self.config = config
        self.template_cnn = TemplateCNN(config)

        # Template bank: (template_bank_size, n_channels, template_len).
        # Populated by `set_template_bank` or loaded from the checkpoint.
        self.register_buffer(
            "template_bank",
            torch.zeros(
                config.template_bank_size,
                config.n_channels,
                config.template_len_samples,
            ),
        )
        # Per-template confidence calibration (set during training).
        self.register_buffer(
            "template_prior", torch.ones(config.template_bank_size) / config.template_bank_size
        )

        # Stage 3 fallback: KMeans centroids. Populated at inference time.
        self._stage3_centroids: Optional[np.ndarray] = None

    # ── template bank management ───────────────────────────────────
    def set_template_bank(self, templates: torch.Tensor, priors: Optional[torch.Tensor] = None):
        """
        templates: (template_bank_size, n_channels, template_len)
        priors:    (template_bank_size,) optional
        """
        assert templates.shape == self.template_bank.shape
        self.template_bank.copy_(templates)
        if priors is not None:
            self.template_prior.copy_(priors / priors.sum())

    # ── stage 1: threshold ─────────────────────────────────────────
    def _stage1_detect(self, spike_band: np.ndarray) -> np.ndarray:
        """
        Returns indices (sample positions) where any channel crosses
        |x| > threshold_sigma * robust_sigma.
        """
        sigma = np.median(np.abs(spike_band), axis=1, keepdims=True) / 0.6745
        sigma = np.maximum(sigma, 1e-6)
        crossings = (np.abs(spike_band) > self.config.threshold_sigma * sigma).any(axis=0)
        return np.flatnonzero(crossings)

    # ── stage 2: CNN template match ────────────────────────────────
    @torch.no_grad()
    def _stage2_match(self, snippets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """
        snippets: (n_spikes, n_channels, template_len)
        Returns: (unit_ids, confidences)
        """
        x = torch.from_numpy(snippets).float()
        logits, _ = self.template_cnn(x)
        probs = F.softmax(logits, dim=-1).numpy()
        unit_ids = probs.argmax(axis=1)
        confidences = probs.max(axis=1)
        return unit_ids, confidences

    # ── stage 3: KMeans fallback ───────────────────────────────────
    def _stage3_cluster(self, snippets: np.ndarray) -> np.ndarray:
        """
        Full clustering on the ambiguous snippets. Fits KMeans on the
        fly if no centroids are cached, then assigns.
        """
        from sklearn.cluster import KMeans

        flat = snippets.reshape(snippets.shape[0], -1)
        if self._stage3_centroids is None or self._stage3_centroids.shape[1] != flat.shape[1]:
            k = min(self.config.stage3_kmeans_k, max(2, flat.shape[0] // 4))
            km = KMeans(n_clusters=k, max_iter=self.config.stage3_max_iter, n_init=1)
            labels = km.fit_predict(flat)
            self._stage3_centroids = km.cluster_centers_
            return labels + self.config.template_bank_size  # offset so IDs do not collide
        # Assign to nearest centroid
        d = ((flat[:, None, :] - self._stage3_centroids[None, :, :]) ** 2).sum(axis=2)
        return d.argmin(axis=1) + self.config.template_bank_size

    # ── the cascade ────────────────────────────────────────────────
    def infer(
        self,
        voltage: np.ndarray,
        sample_rate: Optional[int] = None,
        baseline_spike_counts: Optional[Dict[int, int]] = None,
    ) -> CascadeReport:
        """
        voltage: (n_channels, n_samples), float32, raw or spike-band.
        sample_rate: defaults to config.sample_rate.
        baseline_spike_counts: optional {unit_id: count} from an offline
            sort, used to compute the correlation. If None, correlation
            is left as NaN and a warning is issued.
        """
        t0 = time.perf_counter()
        sample_rate = sample_rate or self.config.sample_rate
        n_channels, n_samples = voltage.shape
        if n_channels != self.config.n_channels:
            raise ValueError(
                f"expected {self.config.n_channels} channels, got {n_channels}"
            )

        # Preprocess: bandpass to the spike band.
        spike_band = bandpass_filter(
            voltage, sample_rate, self.config.spike_band_hz[0], self.config.spike_band_hz[1]
        )

        # Stage 1.
        spike_positions = self._stage1_detect(spike_band)
        n_windows = max(1, n_samples - self.config.template_len_samples)
        activity_fraction = len(spike_positions) / n_windows

        if len(spike_positions) == 0:
            # Zero-spike branch: report NaN, not 0.0.
            latency_ms = (time.perf_counter() - t0) * 1000
            return CascadeReport(
                units={},
                activity_fraction=0.0,
                stage3_fraction=0.0,
                latency_ms=latency_ms,
                correlation=float("nan"),
                warning=(
                    "zero spikes detected; correlation is NaN, not 0. "
                    "The speedup over an empty baseline is undefined."
                ),
            )

        # Extract snippets centred on the crossings.
        snippets = extract_snippets(
            spike_band, spike_positions, self.config.template_len_samples
        )  # (n_spikes, n_channels, template_len)

        # Stage 2.
        unit_ids_s2, conf_s2 = self._stage2_match(snippets)

        # Stage 3: only the low-confidence spikes.
        ambiguous = conf_s2 < self.config.template_confidence_threshold
        stage3_fraction = ambiguous.mean()
        if ambiguous.any():
            unit_ids_s3 = self._stage3_cluster(snippets[ambiguous])
            unit_ids = unit_ids_s2.copy()
            unit_ids[ambiguous] = unit_ids_s3
        else:
            unit_ids = unit_ids_s2

        # Convert positions to milliseconds.
        spike_times_ms = spike_positions / sample_rate * 1000.0

        # Build the per-unit spike-time dictionary.
        units: Dict[int, List[float]] = {}
        for uid, t in zip(unit_ids.tolist(), spike_times_ms.tolist()):
            units.setdefault(int(uid), []).append(float(t))

        # Correlation against the offline baseline (if provided).
        correlation = float("nan")
        if baseline_spike_counts is not None:
            correlation = self._correlation(units, baseline_spike_counts)

        latency_ms = (time.perf_counter() - t0) * 1000
        return CascadeReport(
            units=units,
            activity_fraction=activity_fraction,
            stage3_fraction=stage3_fraction,
            latency_ms=latency_ms,
            correlation=correlation,
        )

    @staticmethod
    def _correlation(
        units: Dict[int, List[float]], baseline: Dict[int, int]
    ) -> float:
        all_ids = sorted(set(units) | set(baseline))
        if not all_ids:
            return float("nan")
        a = np.array([len(units.get(u, [])) for u in all_ids], dtype=float)
        b = np.array([baseline.get(u, 0) for u in all_ids], dtype=float)
        if a.std() == 0.0 or b.std() == 0.0:
            return float("nan")
        return float(np.corrcoef(a, b)[0, 1])