"""
Online wrapper around SpikeCascadeForSorting. Consumes 50 ms chunks
from a SpikeGLX-like stream and returns per-chunk CascadeReports.

This is the class a closed-loop experiment would use: the loop runs
at ~1 kHz (one chunk per millisecond of wall clock), and the latency
of each `process_chunk` call is the closed-loop latency.
"""

import time
from collections import deque
from typing import Deque, Optional

import numpy as np

from modeling_spikecascade import SpikeCascadeForSorting, CascadeReport


class StreamingSorter:
    def __init__(
        self,
        model: SpikeCascadeForSorting,
        chunk_ms: float = 50.0,
        overlap_ms: float = 2.0,
    ):
        self.model = model
        self.sample_rate = model.config.sample_rate
        self.chunk_samples = int(chunk_ms / 1000 * self.sample_rate)
        self.overlap_samples = int(overlap_ms / 1000 * self.sample_rate)
        self.buffer: Deque[np.ndarray] = deque()
        self.reports: list[CascadeReport] = []

    def push(self, chunk: np.ndarray):
        """chunk: (n_channels, n_samples)."""
        assert chunk.shape[1] == self.chunk_samples
        self.buffer.append(chunk)

    def step(self) -> Optional[CascadeReport]:
        """
        Pop one chunk, run the cascade, and return the report.
        Returns None if the buffer is empty.
        """
        if not self.buffer:
            return None
        chunk = self.buffer.popleft()
        report = self.model.infer(chunk, sample_rate=self.sample_rate)
        self.reports.append(report)
        return report

    def run(self, n_steps: int) -> list[CascadeReport]:
        out = []
        for _ in range(n_steps):
            r = self.step()
            if r is None:
                break
            out.append(r)
        return out

    def summary(self) -> dict:
        if not self.reports:
            return {}
        return {
            "n_chunks": len(self.reports),
            "mean_activity_fraction": float(np.mean([r.activity_fraction for r in self.reports])),
            "mean_stage3_fraction": float(np.mean([r.stage3_fraction for r in self.reports])),
            "mean_latency_ms": float(np.mean([r.latency_ms for r in self.reports])),
            "mean_correlation": float(
                np.nanmean([r.correlation for r in self.reports])
            ),
            "zero_spike_chunks": sum(1 for r in self.reports if r.activity_fraction == 0.0),
        }