//! An in-process acoustic echo canceller built on a **Normalized Least Mean
//! Squares (NLMS)** adaptive FIR filter.
//!
//! The idea: we know the far-end signal exactly (it's what we're about to play
//! out), and the mic picks up an echo of it shaped by the unknown room. An
//! adaptive filter continuously estimates that room response from the data:
//! each sample it predicts the echo as a weighted sum of recent far-end samples,
//! subtracts that prediction from the mic, and nudges its weights to shrink the
//! leftover. With only far-end present, the leftover (the "error" signal) drives
//! the weights toward the true echo path; once converged, the error *is* the
//! cleaned near-end audio.
//!
//! NLMS over plain LMS divides the update by the reference's instantaneous energy,
//! which makes the convergence speed independent of how loud the far-end is — the
//! reason it's the workhorse of practical echo cancellation.
//!
//! This is the pure, sample-at-a-time core. Double-talk handling and a residual
//! suppressor build on top of it; integration into the live audio path wraps it.

use std::collections::VecDeque;

/// A single-channel NLMS echo canceller.
pub struct Nlms {
    /// Adaptive filter coefficients (the estimated echo-path impulse response).
    /// `weights[k]` multiplies the far-end sample `k` steps in the past.
    weights: Vec<f32>,
    /// Ring buffer of recent far-end (reference) samples, same length as weights.
    history: Vec<f32>,
    /// Index in `history` where the newest sample currently lives.
    pos: usize,
    /// Step size (0 < mu < 2). Larger = faster adaptation but more misadjustment.
    mu: f32,
    /// Regularization added to the energy normalizer to avoid divide-by-zero and
    /// runaway updates when the far-end is near silent.
    eps: f32,
    /// When false, weights are held fixed (used during double-talk).
    adapting: bool,
}

impl Nlms {
    /// Creates an NLMS canceller with `taps` filter coefficients (should cover the
    /// echo path's length: delay + reverb tail), step size `mu` in `(0, 2)`, and
    /// regularization `eps`. Sensible defaults: `mu = 0.5`, `eps = 1e-6`.
    pub fn new(taps: usize, mu: f32, eps: f32) -> Self {
        let taps = taps.max(1);
        Nlms {
            weights: vec![0.0; taps],
            history: vec![0.0; taps],
            pos: 0,
            mu,
            eps,
            adapting: true,
        }
    }

    /// Number of filter taps.
    pub fn taps(&self) -> usize {
        self.weights.len()
    }

    /// Enables or disables weight adaptation. Disable during double-talk so the
    /// near-end voice doesn't corrupt the learned echo path.
    pub fn set_adapting(&mut self, on: bool) {
        self.adapting = on;
    }

    /// A copy of the current estimated impulse response (the filter weights),
    /// newest-tap-first. Useful for tests/inspection.
    pub fn weights(&self) -> &[f32] {
        &self.weights
    }

    /// Processes one sample pair and returns the echo-cancelled output.
    ///
    /// `reference` is the far-end sample about to be (or being) played out;
    /// `mic` is the simultaneously-captured near-end mic sample (near-end voice +
    /// echo of the far-end + noise). The returned value is `mic` minus the
    /// filter's echo estimate — i.e. the cleaned near-end.
    pub fn process_sample(&mut self, reference: f32, mic: f32) -> f32 {
        let l = self.weights.len();
        // Store newest reference at the current ring position.
        self.history[self.pos] = reference;

        // Predict the echo: weights[k] pairs with the sample k steps back, which
        // sits at ring index (pos - k). Accumulate the reference energy too.
        let mut echo_est = 0.0f32;
        let mut energy = 0.0f32;
        for k in 0..l {
            let idx = (self.pos + l - k) % l;
            let s = self.history[idx];
            echo_est += self.weights[k] * s;
            energy += s * s;
        }

        let error = mic - echo_est;

        // NLMS weight update: w += mu * error * x / (||x||^2 + eps).
        if self.adapting {
            let g = self.mu * error / (energy + self.eps);
            for k in 0..l {
                let idx = (self.pos + l - k) % l;
                self.weights[k] += g * self.history[idx];
            }
        }

        // Advance the ring so the next sample overwrites the oldest entry.
        self.pos = (self.pos + 1) % l;
        error
    }

    /// Convenience: cancels a whole buffer, returning the cleaned output. `mic`
    /// and `reference` are processed pairwise over their shared length.
    pub fn process(&mut self, reference: &[f32], mic: &[f32]) -> Vec<f32> {
        let n = reference.len().min(mic.len());
        let mut out = Vec::with_capacity(n);
        for i in 0..n {
            out.push(self.process_sample(reference[i], mic[i]));
        }
        out
    }
}

/// Geigel double-talk detector.
///
/// The danger with an adaptive echo canceller is *double-talk*: when the
/// near-end person speaks at the same time as the far-end, the mic contains
/// near-end voice the filter has no reference for. If it keeps adapting it will
/// (wrongly) try to model that voice as echo and diverge, wrecking the painstakingly
/// learned echo path. The fix is to detect double-talk and freeze adaptation
/// while it lasts (the filter coasts on its last good estimate).
///
/// Geigel's classic test: declare near-end speech present when the mic level
/// rises above a fraction `threshold` of the recent peak far-end level. The echo
/// is always an attenuated, smeared copy of the far-end, so a mic louder than the
/// far-end's recent peak (scaled) can only be near-end energy. A `hangover` keeps
/// the freeze latched for a while after the trigger so brief dips mid-utterance
/// don't immediately re-enable adaptation.
pub struct DoubleTalkDetector {
    /// Sliding window (samples) over which the peak far-end level is tracked.
    window: usize,
    /// Detection threshold: `|mic| ≥ threshold · max|far|` over the window ⇒ talk.
    threshold: f32,
    /// Samples to keep adaptation frozen after the last trigger.
    hangover: usize,
    /// Monotonic-decreasing deque of `(sample_index, |far|)` for O(1) sliding max.
    mono: VecDeque<(usize, f32)>,
    /// Running sample counter.
    i: usize,
    /// Samples of freeze remaining.
    hangover_left: usize,
}

impl DoubleTalkDetector {
    /// `window` should roughly cover the echo path length (e.g. the filter taps),
    /// `threshold` is the Geigel fraction (≈0.5–0.7 for ~6 dB echo coupling), and
    /// `hangover` is how long to stay frozen after a trigger (e.g. ~50 ms).
    pub fn new(window: usize, threshold: f32, hangover: usize) -> Self {
        DoubleTalkDetector {
            window: window.max(1),
            threshold,
            hangover,
            mono: VecDeque::new(),
            i: 0,
            hangover_left: 0,
        }
    }

    /// Feeds one sample pair; returns `true` if double-talk is currently declared
    /// (i.e. adaptation should be frozen this sample).
    pub fn update(&mut self, reference: f32, mic: f32) -> bool {
        let ax = reference.abs();
        // Maintain the monotonic-decreasing deque: drop smaller tail entries.
        while matches!(self.mono.back(), Some(&(_, v)) if v <= ax) {
            self.mono.pop_back();
        }
        self.mono.push_back((self.i, ax));
        // Evict entries that have fallen out of the sliding window.
        let window_start = (self.i + 1).saturating_sub(self.window);
        while matches!(self.mono.front(), Some(&(idx, _)) if idx < window_start) {
            self.mono.pop_front();
        }
        let far_max = self.mono.front().map(|&(_, v)| v).unwrap_or(0.0);
        self.i += 1;

        // Geigel decision (only meaningful when the far-end is actually active).
        if far_max > 1e-6 && mic.abs() >= self.threshold * far_max {
            self.hangover_left = self.hangover;
        }
        if self.hangover_left > 0 {
            self.hangover_left -= 1;
            true
        } else {
            false
        }
    }
}

/// A complete in-process echo canceller: an [`Nlms`] adaptive filter guarded by a
/// [`DoubleTalkDetector`]. This is the unit the live audio path will wrap — feed
/// it the far-end reference and the mic, get back cleaned near-end audio, with
/// adaptation automatically frozen during double-talk.
pub struct EchoCanceller {
    nlms: Nlms,
    dtd: DoubleTalkDetector,
    /// Whether double-talk protection is active (off = raw NLMS, for A/B testing).
    dtd_enabled: bool,
    /// Count of samples the most recent run flagged as double-talk.
    double_talk_samples: usize,
    /// Total samples processed (for the double-talk rate).
    total_samples: usize,
}

impl EchoCanceller {
    /// Builds a canceller with `taps` filter length, NLMS `mu`/`eps`, and DTD
    /// `threshold`/`hangover`. The DTD window is the filter length.
    pub fn new(taps: usize, mu: f32, eps: f32, dtd_threshold: f32, hangover: usize) -> Self {
        EchoCanceller {
            nlms: Nlms::new(taps, mu, eps),
            dtd: DoubleTalkDetector::new(taps, dtd_threshold, hangover),
            dtd_enabled: true,
            double_talk_samples: 0,
            total_samples: 0,
        }
    }

    /// Enables/disables the double-talk freeze (disable to compare against raw NLMS).
    pub fn set_dtd_enabled(&mut self, on: bool) {
        self.dtd_enabled = on;
    }

    /// Fraction of processed samples that were flagged as double-talk (0.0–1.0).
    pub fn double_talk_rate(&self) -> f32 {
        if self.total_samples == 0 {
            0.0
        } else {
            self.double_talk_samples as f32 / self.total_samples as f32
        }
    }

    /// Processes one sample pair: runs the detector, freezes adaptation if it
    /// fires, and returns the cleaned output.
    pub fn process_sample(&mut self, reference: f32, mic: f32) -> f32 {
        let double_talk = self.dtd_enabled && self.dtd.update(reference, mic);
        self.nlms.set_adapting(!double_talk);
        self.total_samples += 1;
        if double_talk {
            self.double_talk_samples += 1;
        }
        self.nlms.process_sample(reference, mic)
    }

    /// Cancels a whole buffer pairwise over the shared length.
    pub fn process(&mut self, reference: &[f32], mic: &[f32]) -> Vec<f32> {
        let n = reference.len().min(mic.len());
        let mut out = Vec::with_capacity(n);
        for i in 0..n {
            out.push(self.process_sample(reference[i], mic[i]));
        }
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dsp::echo_path::EchoPath;
    use crate::dsp::generators::white_noise;
    use crate::dsp::metrics::erle;

    /// With only far-end present (no near-end talk), the filter should identify
    /// the echo path and drive ERLE high — and the late part of the call should
    /// be far better cancelled than the early part (it learns over time).
    #[test]
    fn converges_on_a_known_echo_path() {
        let far = white_noise(0.5, 48_000, 1);
        let path = EchoPath::synthetic(64, 192, 0.5, 7); // 256-tap path
        let echo = path.apply(&far);
        // Mic is echo-only (single-talk): the ideal case for adaptation.
        let mut aec = Nlms::new(512, 0.5, 1e-6);
        let cleaned = aec.process(&far, &echo);

        // Compare ERLE on the first vs last quarter.
        let q = echo.len() / 4;
        let early = erle(&echo[..q], &cleaned[..q]);
        let late = erle(&echo[3 * q..], &cleaned[3 * q..]);
        assert!(
            late > early + 10.0,
            "should improve markedly: early {early:.1} late {late:.1}"
        );
        assert!(
            late > 20.0,
            "converged ERLE should exceed 20 dB, got {late:.1}"
        );
    }

    #[test]
    fn frozen_filter_does_not_adapt() {
        let far = white_noise(0.5, 4000, 2);
        let echo = EchoPath::synthetic(16, 48, 0.5, 3).apply(&far);
        let mut aec = Nlms::new(128, 0.5, 1e-6);
        aec.set_adapting(false);
        aec.process(&far, &echo);
        // Weights must remain exactly zero with adaptation off.
        assert!(aec.weights().iter().all(|&w| w == 0.0));
    }

    #[test]
    fn passes_near_end_through_when_no_echo() {
        // No far-end (reference silent) -> filter predicts ~0 -> near-end passes
        // through essentially unchanged.
        let near = white_noise(0.3, 2000, 9);
        let silent_ref = vec![0.0f32; near.len()];
        let mut aec = Nlms::new(128, 0.5, 1e-6);
        let out = aec.process(&silent_ref, &near);
        for (a, b) in near.iter().zip(&out) {
            assert!(
                (a - b).abs() < 1e-6,
                "near-end should pass through: {a} vs {b}"
            );
        }
    }

    /// Builds a near-end signal that is silent until `onset`, then a loud tone —
    /// the canonical "far-end converges, then both talk" double-talk scenario.
    fn near_end_with_onset(len: usize, onset: usize) -> Vec<f32> {
        use std::f64::consts::PI;
        (0..len)
            .map(|i| {
                if i < onset {
                    0.0
                } else {
                    0.5 * (2.0 * PI * 300.0 * i as f64 / 48_000.0).sin() as f32
                }
            })
            .collect()
    }

    #[test]
    fn detector_fires_on_near_end_not_on_echo() {
        // Realistic ~-12 dB echo coupling (atten 0.25); louder coupling makes
        // Geigel false-positive on white-noise echo peaks (a real limitation).
        let far = white_noise(0.5, 20_000, 4);
        let echo = EchoPath::synthetic(64, 192, 0.25, 5).apply(&far);
        let onset = 10_000;
        let near = near_end_with_onset(far.len(), onset);
        let mic: Vec<f32> = echo.iter().zip(&near).map(|(&e, &n)| e + n).collect();

        let mut dtd = DoubleTalkDetector::new(256, 0.7, 200);
        let mut early_hits = 0;
        let mut late_hits = 0;
        for i in 0..far.len() {
            if dtd.update(far[i], mic[i]) {
                if i < onset {
                    early_hits += 1
                } else {
                    late_hits += 1
                }
            }
        }
        // Echo-only stretch should rarely trip; near-end stretch should trip a lot.
        let early_rate = early_hits as f32 / onset as f32;
        let late_rate = late_hits as f32 / (far.len() - onset) as f32;
        assert!(
            early_rate < 0.10,
            "false-positive rate {early_rate:.2} too high"
        );
        assert!(
            late_rate > 0.50,
            "missed double-talk, rate only {late_rate:.2}"
        );
    }

    #[test]
    fn dtd_protects_convergence_under_double_talk() {
        // Far-end runs the whole time; near-end starts at 60%. The filter should
        // converge during the single-talk opening, then the DTD freezes it so the
        // late (double-talk) echo stays cancelled. Without DTD it diverges.
        let far = white_noise(0.5, 60_000, 6);
        let path = EchoPath::synthetic(64, 192, 0.25, 7);
        let echo = path.apply(&far);
        let onset = (far.len() as f32 * 0.6) as usize;
        let near = near_end_with_onset(far.len(), onset);
        let mic: Vec<f32> = echo.iter().zip(&near).map(|(&e, &n)| e + n).collect();

        let late = onset..far.len();

        let mut with_dtd = EchoCanceller::new(512, 0.5, 1e-6, 0.7, 1200);
        let cleaned_d = with_dtd.process(&far, &mic);
        let res_d: Vec<f32> = cleaned_d.iter().zip(&near).map(|(&c, &n)| c - n).collect();
        let erle_dtd = erle(&echo[late.clone()], &res_d[late.clone()]);

        let mut no_dtd = EchoCanceller::new(512, 0.5, 1e-6, 0.7, 1200);
        no_dtd.set_dtd_enabled(false);
        let cleaned_n = no_dtd.process(&far, &mic);
        let res_n: Vec<f32> = cleaned_n.iter().zip(&near).map(|(&c, &n)| c - n).collect();
        let erle_no = erle(&echo[late.clone()], &res_n[late]);

        assert!(
            erle_dtd > erle_no + 15.0,
            "DTD should hold the echo path: with {erle_dtd:.1} dB vs without {erle_no:.1} dB"
        );
        assert!(
            erle_dtd > 15.0,
            "held filter should still cancel echo: {erle_dtd:.1} dB"
        );
    }
}