//! CSI Spectrogram Generation
//!
//! Constructs 2D time-frequency matrices via Short-Time Fourier Transform (STFT)
//! applied to temporal CSI amplitude streams. The resulting spectrograms are the
//! standard input format for CNN-based WiFi activity recognition.
//!
//! # References
//! - Used in virtually all CNN-based WiFi sensing papers since 2018

use ndarray::Array2;
use num_complex::Complex64;
use rustfft::{Fft, FftPlanner};
use ruvector_attn_mincut::attn_mincut;
use std::f64::consts::PI;
use std::sync::Arc;

/// Configuration for spectrogram generation.
#[derive(Debug, Clone)]
pub struct SpectrogramConfig {
    /// FFT window size (number of samples per frame)
    pub window_size: usize,
    /// Hop size (step between consecutive frames). Smaller = more overlap.
    pub hop_size: usize,
    /// Window function to apply
    pub window_fn: WindowFunction,
    /// Whether to compute power (magnitude squared) or magnitude
    pub power: bool,
}

impl Default for SpectrogramConfig {
    fn default() -> Self {
        Self {
            window_size: 256,
            hop_size: 64,
            window_fn: WindowFunction::Hann,
            power: true,
        }
    }
}

/// Window function types.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WindowFunction {
    /// Rectangular (no windowing)
    Rectangular,
    /// Hann window (cosine-squared taper)
    Hann,
    /// Hamming window
    Hamming,
    /// Blackman window (lower sidelobe level)
    Blackman,
}

/// Result of spectrogram computation.
#[derive(Debug, Clone)]
pub struct Spectrogram {
    /// Power/magnitude values: rows = frequency bins, columns = time frames.
    /// Only positive frequencies (0 to Nyquist), so rows = window_size/2 + 1.
    pub data: Array2<f64>,
    /// Number of frequency bins
    pub n_freq: usize,
    /// Number of time frames
    pub n_time: usize,
    /// Frequency resolution (Hz per bin)
    pub freq_resolution: f64,
    /// Time resolution (seconds per frame)
    pub time_resolution: f64,
}

/// Compute spectrogram of a 1D signal.
///
/// Returns a time-frequency matrix suitable as CNN input.
pub fn compute_spectrogram(
    signal: &[f64],
    sample_rate: f64,
    config: &SpectrogramConfig,
) -> Result<Spectrogram, SpectrogramError> {
    if signal.len() < config.window_size {
        return Err(SpectrogramError::SignalTooShort {
            signal_len: signal.len(),
            window_size: config.window_size,
        });
    }
    if config.hop_size == 0 {
        return Err(SpectrogramError::InvalidHopSize);
    }
    if config.window_size == 0 {
        return Err(SpectrogramError::InvalidWindowSize);
    }

    let mut planner = FftPlanner::new();
    let fft = planner.plan_fft_forward(config.window_size);
    let window = make_window(config.window_fn, config.window_size);
    Ok(compute_spectrogram_with_plan(
        signal,
        sample_rate,
        config,
        &fft,
        &window,
    ))
}

/// STFT core that runs against a **pre-planned** FFT and pre-built window.
///
/// ADR-154 §7.4 #20: `compute_spectrogram` re-plans the FFT on every call, so
/// `compute_multi_subcarrier_spectrogram` (which calls it once per subcarrier)
/// re-planned the same length-`window_size` FFT for *every* subcarrier. This
/// helper hoists the plan + window out of the per-subcarrier loop. The numeric
/// body is byte-for-byte the old loop — only the plan/window construction is
/// lifted — so the output is **bit-identical** to the per-call path (asserted by
/// `multi_subcarrier_hoisted_plan_bit_identical`). Callers must pass a plan
/// built for exactly `config.window_size` and a window of that length.
fn compute_spectrogram_with_plan(
    signal: &[f64],
    sample_rate: f64,
    config: &SpectrogramConfig,
    fft: &Arc<dyn Fft<f64>>,
    window: &[f64],
) -> Spectrogram {
    debug_assert_eq!(window.len(), config.window_size, "window/plan size mismatch");
    debug_assert_eq!(fft.len(), config.window_size, "FFT/window size mismatch");

    let n_frames = (signal.len() - config.window_size) / config.hop_size + 1;
    let n_freq = config.window_size / 2 + 1;

    let mut data = Array2::zeros((n_freq, n_frames));

    for frame in 0..n_frames {
        let start = frame * config.hop_size;
        let end = start + config.window_size;

        // Apply window and convert to complex
        let mut buffer: Vec<Complex64> = signal[start..end]
            .iter()
            .zip(window.iter())
            .map(|(&s, &w)| Complex64::new(s * w, 0.0))
            .collect();

        fft.process(&mut buffer);

        // Store positive frequencies
        for bin in 0..n_freq {
            let mag = buffer[bin].norm();
            data[[bin, frame]] = if config.power { mag * mag } else { mag };
        }
    }

    Spectrogram {
        data,
        n_freq,
        n_time: n_frames,
        freq_resolution: sample_rate / config.window_size as f64,
        time_resolution: config.hop_size as f64 / sample_rate,
    }
}

/// Compute spectrogram for each subcarrier from a temporal CSI matrix.
///
/// Input: `csi_temporal` is (num_samples × num_subcarriers) amplitude matrix.
/// Returns one spectrogram per subcarrier.
pub fn compute_multi_subcarrier_spectrogram(
    csi_temporal: &Array2<f64>,
    sample_rate: f64,
    config: &SpectrogramConfig,
) -> Result<Vec<Spectrogram>, SpectrogramError> {
    let (n_samples, n_sc) = csi_temporal.dim();

    // ADR-154 §7.4 #20: validate *once* (same checks `compute_spectrogram`
    // makes), then plan the FFT + build the window *once* and reuse them across
    // every subcarrier instead of re-planning per column. The window length is
    // identical for all subcarriers, so this is pure hoisting — output stays
    // bit-identical to the per-call path.
    if n_samples < config.window_size {
        return Err(SpectrogramError::SignalTooShort {
            signal_len: n_samples,
            window_size: config.window_size,
        });
    }
    if config.hop_size == 0 {
        return Err(SpectrogramError::InvalidHopSize);
    }
    if config.window_size == 0 {
        return Err(SpectrogramError::InvalidWindowSize);
    }

    let mut planner = FftPlanner::new();
    let fft = planner.plan_fft_forward(config.window_size);
    let window = make_window(config.window_fn, config.window_size);

    let mut spectrograms = Vec::with_capacity(n_sc);
    for sc in 0..n_sc {
        let col: Vec<f64> = csi_temporal.column(sc).to_vec();
        spectrograms.push(compute_spectrogram_with_plan(
            &col,
            sample_rate,
            config,
            &fft,
            &window,
        ));
    }

    Ok(spectrograms)
}

/// Generate a window function.
///
/// ADR-154: the cosine windows divide by `(size - 1)`, which is zero for
/// `size == 1` (→ NaN samples) and underflows the empty-range maths for tiny
/// sizes. We short-circuit `size <= 1` to a safe constant window (empty for 0,
/// single unit sample for 1) before any `size - 1` arithmetic runs.
fn make_window(kind: WindowFunction, size: usize) -> Vec<f64> {
    if size <= 1 {
        return vec![1.0; size];
    }
    match kind {
        WindowFunction::Rectangular => vec![1.0; size],
        WindowFunction::Hann => (0..size)
            .map(|i| 0.5 * (1.0 - (2.0 * PI * i as f64 / (size - 1) as f64).cos()))
            .collect(),
        WindowFunction::Hamming => (0..size)
            .map(|i| 0.54 - 0.46 * (2.0 * PI * i as f64 / (size - 1) as f64).cos())
            .collect(),
        WindowFunction::Blackman => (0..size)
            .map(|i| {
                let n = (size - 1) as f64;
                0.42 - 0.5 * (2.0 * PI * i as f64 / n).cos()
                    + 0.08 * (4.0 * PI * i as f64 / n).cos()
            })
            .collect(),
    }
}

/// Apply attention-gating to a computed CSI spectrogram using ruvector-attn-mincut.
///
/// Treats each time frame as an attention token (d = n_freq_bins features,
/// seq_len = n_time_frames tokens). Self-attention (Q=K=V) gates coherent
/// body-motion frames and suppresses uncorrelated noise/interference frames.
///
/// # Arguments
/// * `spectrogram` - Row-major [n_freq_bins × n_time_frames] f32 slice
/// * `n_freq` - Number of frequency bins (feature dimension d)
/// * `n_time` - Number of time frames (sequence length)
/// * `lambda` - Gating strength: 0.1 = mild, 0.3 = moderate, 0.5 = aggressive
///
/// # Returns
/// Gated spectrogram as Vec<f32>, same shape as input
pub fn gate_spectrogram(
    spectrogram: &[f32],
    n_freq: usize,
    n_time: usize,
    lambda: f32,
) -> Vec<f32> {
    debug_assert_eq!(
        spectrogram.len(),
        n_freq * n_time,
        "spectrogram length must equal n_freq * n_time"
    );

    if n_freq == 0 || n_time == 0 {
        return spectrogram.to_vec();
    }

    // Q = K = V = spectrogram (self-attention over time frames)
    let result = attn_mincut(
        spectrogram,
        spectrogram,
        spectrogram,
        n_freq, // d = feature dimension
        n_time, // seq_len = time tokens
        lambda,
        /*tau=*/ 2,
        /*eps=*/ 1e-7_f32,
    );
    result.output
}

/// Errors from spectrogram computation.
#[derive(Debug, thiserror::Error)]
pub enum SpectrogramError {
    #[error("Signal too short ({signal_len} samples) for window size {window_size}")]
    SignalTooShort {
        signal_len: usize,
        window_size: usize,
    },

    #[error("Hop size must be > 0")]
    InvalidHopSize,

    #[error("Window size must be > 0")]
    InvalidWindowSize,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_spectrogram_dimensions() {
        let sample_rate = 100.0;
        let signal: Vec<f64> = (0..1000)
            .map(|i| (i as f64 / sample_rate * 2.0 * PI * 5.0).sin())
            .collect();

        let config = SpectrogramConfig {
            window_size: 128,
            hop_size: 32,
            window_fn: WindowFunction::Hann,
            power: true,
        };

        let spec = compute_spectrogram(&signal, sample_rate, &config).unwrap();
        assert_eq!(spec.n_freq, 65); // 128/2 + 1
        assert_eq!(spec.n_time, (1000 - 128) / 32 + 1); // 28 frames
        assert_eq!(spec.data.dim(), (65, 28));
    }

    #[test]
    fn test_single_frequency_peak() {
        // A pure 10 Hz tone at 100 Hz sampling → peak at bin 10/100*256 ≈ bin 26
        let sample_rate = 100.0;
        let freq = 10.0;
        let signal: Vec<f64> = (0..1024)
            .map(|i| (2.0 * PI * freq * i as f64 / sample_rate).sin())
            .collect();

        let config = SpectrogramConfig {
            window_size: 256,
            hop_size: 128,
            window_fn: WindowFunction::Hann,
            power: true,
        };

        let spec = compute_spectrogram(&signal, sample_rate, &config).unwrap();

        // Find peak frequency bin in the first frame
        let frame = spec.data.column(0);
        let peak_bin = frame
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(i, _)| i)
            .unwrap();

        let peak_freq = peak_bin as f64 * spec.freq_resolution;
        assert!(
            (peak_freq - freq).abs() < spec.freq_resolution * 2.0,
            "Peak at {:.1} Hz, expected {:.1} Hz",
            peak_freq,
            freq
        );
    }

    #[test]
    fn test_window_functions_symmetric() {
        for wf in [
            WindowFunction::Hann,
            WindowFunction::Hamming,
            WindowFunction::Blackman,
        ] {
            let w = make_window(wf, 64);
            for i in 0..32 {
                assert!(
                    (w[i] - w[63 - i]).abs() < 1e-10,
                    "{:?} not symmetric at {}",
                    wf,
                    i
                );
            }
        }
    }

    #[test]
    fn test_rectangular_window_all_ones() {
        let w = make_window(WindowFunction::Rectangular, 100);
        assert!(w.iter().all(|&v| (v - 1.0).abs() < 1e-10));
    }

    // ADR-154: degenerate window sizes must not divide by (n-1)==0 → NaN.
    #[test]
    fn make_window_size_0_and_1_are_safe() {
        for wf in [
            WindowFunction::Hann,
            WindowFunction::Hamming,
            WindowFunction::Blackman,
            WindowFunction::Rectangular,
        ] {
            assert!(make_window(wf, 0).is_empty(), "{wf:?} size-0 must be empty");
            let w1 = make_window(wf, 1);
            assert_eq!(w1.len(), 1, "{wf:?} size-1 must have one sample");
            assert!(
                w1[0].is_finite() && (w1[0] - 1.0).abs() < 1e-12,
                "{wf:?} size-1 must be a finite unit sample, got {}",
                w1[0]
            );
        }
    }

    #[test]
    fn test_signal_too_short() {
        let signal = vec![1.0; 10];
        let config = SpectrogramConfig {
            window_size: 256,
            ..Default::default()
        };
        assert!(matches!(
            compute_spectrogram(&signal, 100.0, &config),
            Err(SpectrogramError::SignalTooShort { .. })
        ));
    }

    #[test]
    fn test_multi_subcarrier() {
        let n_samples = 500;
        let n_sc = 8;
        let csi = Array2::from_shape_fn((n_samples, n_sc), |(t, sc)| {
            let freq = 1.0 + sc as f64 * 0.5;
            (2.0 * PI * freq * t as f64 / 100.0).sin()
        });

        let config = SpectrogramConfig {
            window_size: 128,
            hop_size: 64,
            ..Default::default()
        };

        let specs = compute_multi_subcarrier_spectrogram(&csi, 100.0, &config).unwrap();
        assert_eq!(specs.len(), n_sc);
        for spec in &specs {
            assert_eq!(spec.n_freq, 65);
        }
    }

    // ADR-154 §7.4 #20: the FFT-planner hoist in
    // `compute_multi_subcarrier_spectrogram` must produce **bit-identical**
    // output to calling `compute_spectrogram` (fresh planner) per subcarrier.
    // We compare `f64::to_bits` of every spectrogram value across several
    // window functions and a realistic 56-subcarrier CSI matrix — the planner
    // change only reorders *when* the (identical) plan is built, never the math.
    #[test]
    fn multi_subcarrier_hoisted_plan_bit_identical() {
        let n_samples = 600;
        let n_sc = 56; // canonical-56 grid — the production subcarrier count
        let sample_rate = 100.0;
        let csi = Array2::from_shape_fn((n_samples, n_sc), |(t, sc)| {
            // Deterministic, non-trivial per-subcarrier content.
            let freq = 0.7 + sc as f64 * 0.13;
            (2.0 * PI * freq * t as f64 / sample_rate).sin()
                + 0.3 * (2.0 * PI * (freq * 2.1) * t as f64 / sample_rate).cos()
        });

        for window_fn in [
            WindowFunction::Hann,
            WindowFunction::Hamming,
            WindowFunction::Blackman,
            WindowFunction::Rectangular,
        ] {
            for &power in &[true, false] {
                let config = SpectrogramConfig {
                    window_size: 128,
                    hop_size: 37, // non-divisor hop to exercise frame edges
                    window_fn,
                    power,
                };

                // AFTER: hoisted-plan path.
                let hoisted =
                    compute_multi_subcarrier_spectrogram(&csi, sample_rate, &config).unwrap();

                // BEFORE: independent per-subcarrier fresh-planner path.
                let reference: Vec<Spectrogram> = (0..n_sc)
                    .map(|sc| {
                        let col: Vec<f64> = csi.column(sc).to_vec();
                        compute_spectrogram(&col, sample_rate, &config).unwrap()
                    })
                    .collect();

                assert_eq!(hoisted.len(), reference.len());
                for (sc, (h, r)) in hoisted.iter().zip(reference.iter()).enumerate() {
                    assert_eq!(h.data.dim(), r.data.dim(), "dim sc={sc} {window_fn:?}");
                    for (a, b) in h.data.iter().zip(r.data.iter()) {
                        assert_eq!(
                            a.to_bits(),
                            b.to_bits(),
                            "bit mismatch sc={sc} {window_fn:?} power={power}: {a} vs {b}"
                        );
                    }
                    assert_eq!(h.freq_resolution.to_bits(), r.freq_resolution.to_bits());
                    assert_eq!(h.time_resolution.to_bits(), r.time_resolution.to_bits());
                }
            }
        }
    }
}

#[cfg(test)]
mod gate_tests {
    use super::*;

    #[test]
    fn gate_spectrogram_preserves_shape() {
        let n_freq = 16_usize;
        let n_time = 10_usize;
        let spectrogram: Vec<f32> = (0..n_freq * n_time).map(|i| i as f32 * 0.01).collect();
        let gated = gate_spectrogram(&spectrogram, n_freq, n_time, 0.3);
        assert_eq!(gated.len(), n_freq * n_time);
    }

    #[test]
    fn gate_spectrogram_zero_lambda_is_identity_ish() {
        let n_freq = 8_usize;
        let n_time = 4_usize;
        let spectrogram: Vec<f32> = vec![1.0; n_freq * n_time];
        // Uniform input — gated output should also be approximately uniform
        let gated = gate_spectrogram(&spectrogram, n_freq, n_time, 0.01);
        assert_eq!(gated.len(), n_freq * n_time);
        // All values should be finite
        assert!(gated.iter().all(|x| x.is_finite()));
    }
}