"""
Frequency Domain Spike Detection Module
======================================

This module implements frequency-domain spike detection using Short-Time Fourier Transform (STFT)
and Smoothed Nonlinear Energy Operator (SNEO) for detecting epileptic spikes in EEG recordings.

The algorithm combines spectral analysis with multi-band consensus voting to provide robust
spike detection that is less susceptible to artifacts compared to traditional amplitude-based methods.
"""

import logging
import warnings
from typing import Literal

try:
    import dask
except Exception:  # pragma: no cover
    dask = None

import numpy as np
import mne
from scipy.signal import iirfilter, iirnotch, sosfiltfilt, filtfilt, stft, windows

try:
    import spikeinterface.core as si

    SPIKEINTERFACE_AVAILABLE = True
except ImportError:  # pragma: no cover
    si = None
    SPIKEINTERFACE_AVAILABLE = False


class FrequencyDomainSpikeDetector:
    """
    Static class for frequency-domain spike detection using STFT and SNEO.

    This detector implements a multi-stage pipeline:
    1. Preprocessing (bandpass + notch filtering)
    2. STFT analysis at specific frequency bands
    3. SNEO detection with multi-band consensus voting
    4. Spike refinement and morphological validation
    5. Temporal clustering and deduplication
    """

    # Default detection parameters
    DEFAULT_PARAMS = {
        "bp": [3.0, 40.0],  # bandpass filter Hz
        "notch": 60.0,  # notch filter Hz
        "notch_q": 30.0,  # notch filter quality factor
        "freq_slices": [10.0, 20.0],  # STFT energy slices Hz
        "window_s": 0.125,  # STFT window size s
        "sneo_percentile": 99.99,  # SNEO threshold percentile
        "cluster_gap_ms": 80.0,  # clustering gap ms
        "search_ms": 160.0,  # spike refinement search window ms
        "baseline_ms": 500.0,  # baseline analysis window ms
        "k_sigma": 3.0,  # statistical significance threshold
        "smooth_window": 7,  # smoothing kernel size
        "vote_k": 2,  # minimum votes required for detection
        "smooth_len": 5,  # SNEO smoothing length
    }

    @staticmethod
    def detect_spikes_recording(
        recording: "si.BaseRecording",
        detection_params: dict = None,
        max_length: int = None,
        multiprocess_mode: Literal["dask", "serial"] = "serial",
    ) -> tuple[list[np.ndarray], "mne.io.RawArray"]:
        """
        Detect spikes in a recording using frequency-domain analysis.

        Args:
            recording (si.BaseRecording): The recording to analyze
            detection_params (dict, optional): Detection parameters. Uses DEFAULT_PARAMS if None
            max_length (int, optional): Maximum length in samples to analyze
            multiprocess_mode (Literal["dask", "serial"]): Processing mode

        Returns:
            tuple: (spike_indices_per_channel, mne_raw_with_annotations)
                - spike_indices_per_channel: List of arrays with spike sample indices per channel
                - mne_raw_with_annotations: MNE RawArray with spike annotations
        """
        if not SPIKEINTERFACE_AVAILABLE:
            raise ImportError(
                "SpikeInterface is required for frequency domain spike detection"
            )

        # Merge with default parameters
        params = FrequencyDomainSpikeDetector.DEFAULT_PARAMS.copy()
        if detection_params:
            params.update(detection_params)

        logging.info(f"Detecting spikes with parameters: {params}")
        logging.debug(f"Recording info: {recording}")
        logging.debug(f"Recording channels: {recording.get_channel_ids()}")

        # Get preprocessed recording
        rec_preprocessed = FrequencyDomainSpikeDetector._apply_preprocessing(
            recording, params
        )

        # Extract data for analysis
        raw_data = rec_preprocessed.get_traces(return_in_uV=True)  # (samples, channels)
        raw_data = raw_data.T  # (channels, samples)

        sampling_freq = rec_preprocessed.get_sampling_frequency()
        channel_names = [str(ch_id) for ch_id in rec_preprocessed.get_channel_ids()]
        n_channels = len(channel_names)

        # Apply max_length if specified
        if max_length and raw_data.shape[1] > max_length:
            raw_data = raw_data[:, :max_length]

        # Create MNE RawArray for consistency
        info = mne.create_info(
            ch_names=channel_names, sfreq=sampling_freq, ch_types="eeg"
        )
        mne_raw = mne.io.RawArray(data=raw_data, info=info)

        # Run spike detection
        match multiprocess_mode:
            case "dask":
                if dask is None:
                    raise ImportError("dask is required for multiprocess_mode='dask'")
                spike_tasks = [
                    dask.delayed(FrequencyDomainSpikeDetector._detect_spikes_channel)(
                        raw_data[ch, :], sampling_freq, params
                    )
                    for ch in range(n_channels)
                ]
                # dask.compute returns tuple, convert to list for consistency with serial mode
                spike_indices_per_channel = list(dask.compute(*spike_tasks))
            case _:
                spike_indices_per_channel = [
                    FrequencyDomainSpikeDetector._detect_spikes_channel(
                        raw_data[ch, :], sampling_freq, params
                    )
                    for ch in range(n_channels)
                ]

        # Add spike annotations to MNE object
        mne_raw_with_annotations = FrequencyDomainSpikeDetector._add_spike_annotations(
            mne_raw, spike_indices_per_channel, sampling_freq
        )

        return spike_indices_per_channel, mne_raw_with_annotations

    @staticmethod
    def _apply_preprocessing(
        recording: "si.BaseRecording", params: dict
    ) -> "si.BaseRecording":
        """Apply bandpass and notch filtering to the recording."""
        # Get raw data for scipy filtering (SpikeInterface preprocessing can be complex)
        raw_data = recording.get_traces(return_in_uV=True)  # (samples, channels)
        raw_data = raw_data.T  # (channels, samples)

        sampling_freq = recording.get_sampling_frequency()

        # Apply bandpass filter
        bp_lo, bp_hi = params["bp"]
        sos_bp = iirfilter(
            N=4,
            Wn=[bp_lo, bp_hi],
            btype="bandpass",
            ftype="butter",
            output="sos",
            fs=sampling_freq,
        )
        raw_filtered = sosfiltfilt(sos_bp, raw_data, axis=-1)

        # Apply notch filter
        notch_freq = params["notch"]
        notch_q = params["notch_q"]
        b_notch, a_notch = iirnotch(w0=notch_freq, Q=notch_q, fs=sampling_freq)
        raw_notch = filtfilt(b_notch, a_notch, raw_filtered, axis=-1)

        # Create new recording with filtered data
        # info = mne.create_info(
        #     ch_names=[str(ch_id) for ch_id in rec.get_channel_ids()],
        #     sfreq=sampling_freq,
        #     ch_types='eeg'
        # )
        # mne_raw = mne.io.RawArray(data=raw_notch, info=info)

        # Convert back to SpikeInterface format
        # Note: This is a simplified approach. In production, might want to use SpikeInterface preprocessing
        # SpikeInterface expects (n_times, n_channels) while our data is (n_channels, n_times)
        filtered_rec = si.NumpyRecording(
            raw_notch.T,
            sampling_frequency=sampling_freq,
            channel_ids=recording.get_channel_ids(),  # Use original, not clone (clone resets string IDs to integers)
        )

        return filtered_rec

    @staticmethod
    def _detect_spikes_channel(
        signal: np.ndarray, fs: float, params: dict
    ) -> np.ndarray:
        """
        Detect spikes in a single channel using the STFT+SNEO algorithm.

        Args:
            signal: 1D signal array
            fs: Sampling frequency
            params: Detection parameters dictionary

        Returns:
            spike_indices: Array of spike sample indices
        """
        logging.debug(
            f"Starting spike detection for channel - signal length: {len(signal)} samples, fs: {fs} Hz"
        )

        # Step 1: STFT analysis
        logging.debug(
            f"Step 1/4: Computing STFT slices at frequencies {params['freq_slices']} Hz with window size {params['window_s']} s"
        )
        slices_dict = FrequencyDomainSpikeDetector._compute_stft_slices(
            signal, fs, freqs=params["freq_slices"], window_s=params["window_s"]
        )
        logging.debug(
            f"Step 1/4: Completed STFT analysis - extracted {len(slices_dict)} frequency bands"
        )

        # Step 2: SNEO detection with multi-band voting
        logging.debug(
            f"Step 2/4: Applying SNEO detection with threshold percentile {params['sneo_percentile']} and vote_k={params['vote_k']}"
        )
        sneo_spikes, _ = FrequencyDomainSpikeDetector._apply_sneo_on_slices(
            slices_dict,
            fs,
            threshold_percentile=params["sneo_percentile"],
            smooth_len=params["smooth_len"],
            vote_k=params["vote_k"],
        )
        logging.info(
            f"Step 2/4: SNEO detection found {len(sneo_spikes)} candidate spikes"
        )

        # Step 3: Spike refinement and morphological validation
        logging.debug(
            f"Step 3/4: Refining spikes with morphological validation (search_ms={params['search_ms']}, k_sigma={params['k_sigma']})"
        )
        neg_spikes = FrequencyDomainSpikeDetector._enforce_downward_and_refine_minimal(
            signal,
            fs,
            sneo_spikes,
            search_ms=params["search_ms"],
            baseline_ms=params["baseline_ms"],
            k_sigma=params["k_sigma"],
            smooth_window=params["smooth_window"],
        )
        logging.info(
            f"Step 3/4: After refinement and validation: {len(neg_spikes)} spikes retained"
        )

        # Step 4: Temporal clustering and deduplication
        logging.debug(
            f"Step 4/4: Clustering and deduplicating spikes (min_gap_ms={params['cluster_gap_ms']})"
        )
        final_spikes = FrequencyDomainSpikeDetector._filter_close_spikes_by_min_local(
            signal,
            fs,
            neg_spikes,
            min_gap_ms=params["cluster_gap_ms"],
            window_ms=60.0,
            smooth_window=5,
        )
        logging.info(
            f"Step 4/4: Final spike count after clustering: {len(final_spikes)} spikes"
        )

        return final_spikes.astype(int)

    @staticmethod
    def _compute_stft_slices(
        signal, fs, freqs=(40, 60), window="hann", window_s=0.125, noverlap=None
    ):
        """
        Compute STFT and extract power at specific frequency bands.

        Args:
            signal: Input signal
            fs: Sampling frequency
            freqs: Frequency points to extract
            window: Window function for STFT
            nperseg: Length of each segment (defaults to ~125ms)
            noverlap: Number of points to overlap

        Returns:
            dict: Frequency -> energy time series mapping
        """
        N = len(signal)
        logging.debug(
            f"Computing STFT slices - signal length: {N} samples, fs: {fs} Hz"
        )

        # Default window length informed by expected spike width (~125 ms)
        nperseg = int(round(fs * window_s))
        logging.debug(
            f"STFT parameters: window={window}, nperseg={nperseg} samples ({window_s} s), noverlap={noverlap}"
        )

        f, t, Zxx = stft(
            signal, fs=fs, window=window, nperseg=nperseg, noverlap=noverlap
        )
        t_samp = (t * fs).astype(float)

        logging.debug(
            f"STFT output dimensions: frequency bins={len(f)}, time bins={len(t)}, Zxx shape={Zxx.shape}"
        )
        logging.debug(
            f"Frequency range: {f[0]:.2f} - {f[-1]:.2f} Hz, frequency resolution: {f[1] - f[0]:.2f} Hz"
        )
        logging.debug(
            f"Time range: {t[0]:.4f} - {t[-1]:.4f} s, STFT time samples: {len(t_samp)}, Original signal samples: {N}"
        )

        slices_dict = {}
        pow_spec = np.abs(Zxx) ** 2

        for f0 in freqs:
            idx = int(np.argmin(np.abs(f - f0)))
            actual_freq = f[idx]
            energy = pow_spec[idx]
            logging.debug(
                f"Extracting frequency slice: target={f0} Hz, actual={actual_freq:.2f} Hz (index={idx}), energy length={len(energy)}"
            )

            # REVIEW the output of stft should already be the correct length, so unsure why interpolation is required
            # Interpolate back to original sampling rate
            energy_resampled = np.interp(
                np.arange(N), t_samp, energy, left=energy[0], right=energy[-1]
            )
            logging.debug(
                f"Interpolated energy from {len(energy)} to {len(energy_resampled)} samples"
            )
            logging.debug(
                f"Energy stats - original: min={energy.min():.2e}, max={energy.max():.2e}, mean={energy.mean():.2e}"
            )
            logging.debug(
                f"Energy stats - resampled: min={energy_resampled.min():.2e}, max={energy_resampled.max():.2e}, mean={energy_resampled.mean():.2e}"
            )

            slices_dict[float(f0)] = energy_resampled

        logging.debug(f"STFT slices computed for {len(slices_dict)} frequency bands")
        return slices_dict

    @staticmethod
    def _sneo(x):
        """Smoothed Nonlinear Energy Operator (SNEO)."""
        return x[1:-1] ** 2 - x[2:] * x[:-2]

    @staticmethod
    def _apply_sneo_on_slices(
        slice_dict, fs, threshold_percentile=99.9, smooth_len=5, vote_k=2
    ):
        """
        Apply SNEO on frequency slices with multi-band consensus voting.

        Args:
            slice_dict: Dictionary of frequency -> energy arrays
            fs: Sampling frequency
            threshold_percentile: Percentile threshold for detection
            smooth_len: Smoothing window length
            vote_k: Minimum number of bands that must agree

        Returns:
            tuple: (final_spikes, sneo_combined)
        """
        sneo_len = (
            len(next(iter(slice_dict.values()))) - 2
        )  # REVIEW fragile hardcoded -2, what if the energy operation doesn't reduce the length by 2?
        sneo_combined = np.zeros(sneo_len)
        detections = []

        # Bartlett window for smoothing
        win = (
            windows.bartlett(smooth_len) if smooth_len > 1 else None
        )  # REVIEW I wonder if this can be replaced with a filter function -- becomes independent of sample frequency

        for _, slice_ in slice_dict.items():
            s = FrequencyDomainSpikeDetector._sneo(slice_)

            # Apply smoothing if specified
            if win is not None:
                if len(s) >= len(win):
                    s = np.convolve(s, win, mode="same")
                else:
                    warnings.warn(
                        f"Skipping smoothing, smoothing window length {len(win)} is greater than the length of the SNEO signal {len(s)}"
                    )

            sneo_combined += s

            # Threshold detection
            thr = np.percentile(s, threshold_percentile)
            # REVIEW +1 seems hardcoded -- better alternative would be to extrapolate the SNEO signal to the original length
            # or an implementation of the SNEO algorithm that preserves signal length
            cand = np.where(s > thr)[0] + 1  # +1 to account for SNEO shift
            detections.append(cand)

        if not detections:
            return np.array([]), sneo_combined

        # Multi-band consensus voting
        all_cand = np.concatenate(detections)
        unique, counts = np.unique(all_cand, return_counts=True)
        final_spikes = unique[counts >= vote_k]

        return final_spikes, sneo_combined

    @staticmethod
    def _enforce_downward_and_refine_minimal(
        signal,
        fs,
        candidates,
        search_ms=160,
        baseline_ms=500,
        k_sigma=3.0,
        smooth_window=7,
    ):
        """
        Refine spike candidates by enforcing downward deflection and statistical significance.

        Args:
            signal: Input signal
            fs: Sampling frequency
            candidates: Initial spike candidates
            search_ms: Search window around candidate (ms)
            baseline_ms: Baseline analysis window (ms)
            k_sigma: Statistical significance threshold
            smooth_window: Smoothing window size

        Returns:
            refined_spikes: Array of refined spike indices
        """
        if len(candidates) == 0:
            return np.array([])

        N = len(signal)
        half = int(
            round(fs * (search_ms / 1e3) / 2)
        )  # REVIEW rename these variables -- spike search window
        base_half = int(
            round(fs * (baseline_ms / 1e3) / 2)
        )  # REVIEW rename these variables -- window size to use as a baseline comparison
        out = []

        for c in np.asarray(candidates, dtype=int):
            # Extract search window
            L = max(0, c - half)
            R = min(N, c + half + 1)
            w = signal[L:R]

            # Apply smoothing
            if smooth_window and smooth_window > 1 and len(w) >= smooth_window:
                pad = smooth_window // 2
                wpad = np.pad(w, (pad, pad), mode="edge")
                kern = (
                    np.ones(smooth_window) / smooth_window
                )  # REVIEW I wonder if this can be replaced with a filter function -- becomes independent of sample frequency
                w = np.convolve(wpad, kern, mode="valid")

            # Find minimum (most negative deflection)
            rel_min = np.argmin(w)
            idx = L + rel_min

            # Extract baseline for statistical analysis
            B0 = max(0, idx - base_half)
            B1 = min(N, idx + base_half)
            base = signal[B0:B1]

            if len(base) < 10:
                warnings.warn(
                    f"Skipping baseline analysis for spike at index {idx}, baseline window length {len(base)} is less than 10"
                )
                continue

            # Exclude spike region from baseline
            exc0 = max(0, idx - half) - B0
            exc1 = min(N, idx + half + 1) - B0
            base_mask = np.ones(len(base), dtype=bool)
            base_mask[exc0:exc1] = False
            baseline_without_spike = base[base_mask] if base_mask.any() else base

            # Statistical significance test using robust statistics
            baseline_median = np.median(baseline_without_spike)
            baseline_mad = (
                np.median(np.abs(baseline_without_spike - baseline_median)) + 1e-12
            )  # REVIEW should this epsilon be a parameter?

            if signal[idx] <= baseline_median - k_sigma * 1.4826 * baseline_mad:
                # Morphological validation: check derivative pattern
                # REVIEW this appears hardcoded. Checks for a negative then positive derivative but
                # what if the before/after is noisy but there's still a trend?
                # maybe better to fit a line to the before/after and check slopes
                d = np.diff(signal[max(0, idx - 3) : min(N, idx + 4)])
                if d.size >= 2 and (
                    np.any(d[: max(1, len(d) // 2)] < 0)
                    and np.any(d[max(1, len(d) // 2) :] > 0)
                ):
                    out.append(idx)

        return np.array(sorted(set(out)))

    @staticmethod
    def _filter_close_spikes_by_min_local(
        signal, fs, spike_indices, min_gap_ms=20, window_ms=60, smooth_window=5
    ):
        """
        Filter out closely spaced spikes by selecting the most negative in each cluster.

        Args:
            signal: Input signal
            fs: Sampling frequency
            spike_indices: Array of spike indices
            min_gap_ms: Minimum gap between spikes (ms)
            window_ms: Window for local optimization (ms)
            smooth_window: Smoothing window size

        Returns:
            filtered_spikes: Array of filtered spike indices
        """
        if len(spike_indices) == 0:
            return np.array([])

        spike_indices = np.sort(np.asarray(spike_indices, dtype=int))
        min_gap = int(round(fs * (min_gap_ms / 1e3)))
        half_w = int(round(fs * (window_ms / 1e3) / 2))
        N = signal.shape[-1]

        # Group nearby spikes
        groups, current_cluster = [], [spike_indices[0]]
        for idx in spike_indices[1:]:
            if idx - current_cluster[-1] <= min_gap:
                current_cluster.append(idx)
            else:
                groups.append(current_cluster)
                current_cluster = [idx]
        groups.append(current_cluster)

        # Select best spike from each group
        chosen = []
        for g in groups:
            g = np.asarray(g)
            L = max(0, g.min() - half_w)
            R = min(N, g.max() + half_w + 1)
            w = signal[L:R]

            # Apply smoothing
            if smooth_window and smooth_window > 1 and len(w) >= smooth_window:
                pad = smooth_window // 2
                wpad = np.pad(w, (pad, pad), mode="edge")
                kern = (
                    np.ones(smooth_window, dtype=float) / smooth_window
                )  # REVIEW I wonder if this can be replaced with a filter function -- becomes independent of sample frequency
                w = np.convolve(wpad, kern, mode="valid")

            rel_min = int(np.argmin(w))
            best_idx = L + rel_min
            chosen.append(best_idx)

        return np.array(sorted(set(chosen)), dtype=int)

    @staticmethod
    def _add_spike_annotations(mne_raw, spike_indices_per_channel, sampling_freq):
        """
        Add spike annotations to MNE RawArray object.

        Args:
            mne_raw: MNE RawArray object
            spike_indices_per_channel: List of spike indices per channel
            sampling_freq: Sampling frequency

        Returns:
            mne_raw_with_annotations: MNE RawArray with spike annotations added
        """
        spike_annotations = []

        for ch_idx, channel_spikes in enumerate(spike_indices_per_channel):
            if len(channel_spikes) > 0:
                # Convert sample indices to seconds
                spike_times_seconds = (
                    channel_spikes + 1
                ) / sampling_freq  # REVIEW why is +1 here?
                for spike_time in spike_times_seconds:
                    spike_annotations.append(
                        {
                            "onset": float(spike_time),
                            "duration": 0.0,
                            "description": f"Spike_Ch{ch_idx}",
                        }
                    )

        if spike_annotations:
            # Sort by onset time
            spike_annotations.sort(key=lambda x: x["onset"])
            annotations = mne.Annotations(
                onset=[a["onset"] for a in spike_annotations],
                duration=[a["duration"] for a in spike_annotations],
                description=[a["description"] for a in spike_annotations],
            )
            mne_raw = mne_raw.copy().set_annotations(annotations)

        return mne_raw
