//! The measuring half of Weld. Runs on the worker thread; deterministic, so tests drive
//! it directly.
//!
//! A Welch estimate of the cross-spectrum S = F_A·conj(F_B) and both auto-spectra runs
//! over Hann-windowed frames with 50 % overlap. Both layers share one complex FFT
//! (A in the real part, B in the imaginary part).
use crate::chain::Fit;
use crate::fit::{search, Point, Problem};
use crate::report::{band_of, curve_frequency, Report, Status, BANDS};
use crate::settings::{Mode, Settings};
use rustfft::{num_complex::Complex32, Fft, FftPlanner};
use std::sync::Arc;

/// Bins below this coherence do not steer the fit.
const MIN_COHERENCE: f32 = 0.3;
/// Below this mean coherence the layers count as unrelated.
const LOOSE_COHERENCE: f32 = 0.3;
/// Frames needed before the first fit.
const MIN_FRAMES: u32 = 8;
/// A new fit replaces the current one only when it gains this much (Follow mode).
const FOLLOW_GAIN_DB: f32 = 0.3;
const LEARN_GAIN_DB: f32 = 0.1;
/// Lag search range and the smallest lag reported as "late".
const LAG_RANGE_S: f32 = 0.005;
const LATE_S: f32 = 0.0001;
const LATE_STRENGTH: f32 = 0.3;
/// Report and fit intervals in seconds of analysed audio.
const REPORT_S: f32 = 0.25;
const FIT_S: f32 = 0.5;
/// −90 dBFS: anything quieter on the sidechain counts as no signal.
const SILENCE: f32 = 3.2e-5;

/// What a poll produced.
#[derive(Debug, Clone, Copy)]
pub struct Update {
    /// Set when a new alignment was accepted.
    pub fit: Option<Fit>,
    pub report: Report,
}

pub struct Analyzer {
    rate: f32,
    n: usize,
    hop: usize,
    window: Vec<f32>,
    ring_a: Vec<f32>,
    ring_b: Vec<f32>,
    write: usize,
    filled: usize,
    since_frame: usize,
    fft: Arc<dyn Fft<f32>>,
    ifft: Arc<dyn Fft<f32>>,
    buf: Vec<Complex32>,
    scratch: Vec<Complex32>,
    sab: Vec<Complex32>,
    saa: Vec<f32>,
    sbb: Vec<f32>,
    frames: u32,
    settings: Settings,
    current: Fit,
    since_report: usize,
    since_fit: usize,
    finish_learn: bool,
    side_level: f32,
    gated: bool,
}

impl Analyzer {
    pub fn new(rate: f32, fit: Fit) -> Self {
        // About 11.7 Hz per bin at any rate.
        let n = ((4096. * rate / 48000.) as usize)
            .next_power_of_two()
            .clamp(1024, 65536);
        let mut planner = FftPlanner::new();
        let fft = planner.plan_fft_forward(n);
        let ifft = planner.plan_fft_inverse(n);
        let scratch_len = fft
            .get_inplace_scratch_len()
            .max(ifft.get_inplace_scratch_len());
        let window = (0..n)
            .map(|i| 0.5 - 0.5 * (std::f32::consts::TAU * i as f32 / n as f32).cos())
            .collect();
        let bins = n / 2 + 1;
        Self {
            rate,
            n,
            hop: n / 2,
            window,
            ring_a: vec![0.; n],
            ring_b: vec![0.; n],
            write: 0,
            filled: 0,
            since_frame: 0,
            fft,
            ifft,
            buf: vec![Complex32::default(); n],
            scratch: vec![Complex32::default(); scratch_len],
            sab: vec![Complex32::default(); bins],
            saa: vec![0.; bins],
            sbb: vec![0.; bins],
            frames: 0,
            settings: Settings::default(),
            current: fit.sanitized(rate),
            since_report: 0,
            since_fit: 0,
            finish_learn: false,
            side_level: 0.,
            gated: false,
        }
    }
    pub fn fft_size(&self) -> usize {
        self.n
    }
    pub fn current(&self) -> &Fit {
        &self.current
    }
    pub fn frames(&self) -> u32 {
        self.frames
    }
    pub fn set_settings(&mut self, settings: Settings) {
        let old = self.settings.mode;
        if settings.mode == Mode::Learn && old != Mode::Learn {
            self.clear();
        }
        if old == Mode::Learn && settings.mode != Mode::Learn {
            self.finish_learn = true;
        }
        self.settings = settings;
    }
    /// Forgets the measurement.
    pub fn clear(&mut self) {
        self.sab.fill(Complex32::default());
        self.saa.fill(0.);
        self.sbb.fill(0.);
        self.frames = 0;
    }
    /// Feeds mono samples of the main layer `a` and the sidechain layer `b`.
    pub fn push(&mut self, a: &[f32], b: &[f32]) {
        for (&x, &y) in a.iter().zip(b) {
            self.ring_a[self.write] = x;
            self.ring_b[self.write] = y;
            self.write = (self.write + 1) % self.n;
            self.filled = (self.filled + 1).min(self.n);
            self.since_frame += 1;
            if self.filled == self.n && self.since_frame >= self.hop {
                self.since_frame = 0;
                self.frame();
            }
        }
        self.since_report += a.len().min(b.len());
        self.since_fit += a.len().min(b.len());
    }
    fn frame(&mut self) {
        let n = self.n;
        let (mut ea, mut eb) = (0f32, 0f32);
        for i in 0..n {
            let j = (self.write + i) % n;
            let (a, b) = (self.ring_a[j], self.ring_b[j]);
            ea += a * a;
            eb += b * b;
            self.buf[i] = Complex32::new(a * self.window[i], b * self.window[i]);
        }
        let (ra, rb) = ((ea / n as f32).sqrt(), (eb / n as f32).sqrt());
        self.side_level = rb;
        let gate = 10f32.powf(self.settings.gate / 20.);
        self.gated = ra < gate || rb < gate;
        if self.gated {
            return;
        }
        self.fft
            .process_with_scratch(&mut self.buf, &mut self.scratch);
        self.frames = self.frames.saturating_add(1);
        let count = 1. / self.frames as f32;
        let alpha = match self.settings.mode {
            Mode::Learn => count,
            _ => {
                let follow = 1. - (-(self.hop as f32) / (self.settings.response * self.rate)).exp();
                count.max(follow)
            }
        };
        for k in 0..=n / 2 {
            let x = self.buf[k];
            let xc = self.buf[(n - k) % n].conj();
            let fa = (x + xc) * 0.5;
            let fb = (x - xc) * Complex32::new(0., -0.5);
            let old = self.sab[k];
            self.sab[k] = old + (fa * fb.conj() - old) * alpha;
            self.saa[k] += (fa.norm_sqr() - self.saa[k]) * alpha;
            self.sbb[k] += (fb.norm_sqr() - self.sbb[k]) * alpha;
        }
    }
    fn frequency(&self, k: usize) -> f32 {
        k as f32 * self.rate / self.n as f32
    }
    fn coherence(&self, k: usize) -> f32 {
        let d = self.saa[k] * self.sbb[k];
        if d > 0. {
            (self.sab[k].norm_sqr() / d).min(1.)
        } else {
            0.
        }
    }
    fn range(&self) -> (f32, f32) {
        let high = self.settings.high.min(0.45 * self.rate);
        (self.settings.low.min(high * 0.5), high)
    }
    /// Fit points between Low and High: one per bin at the bottom, merged to 1/32 octave
    /// higher up.
    fn points(&self) -> Vec<Point> {
        let (low, high) = self.range();
        let mut points: Vec<Point> = Vec::new();
        let mut last = usize::MAX;
        for k in 1..=self.n / 2 {
            let f = self.frequency(k);
            if f < low || f > high {
                continue;
            }
            let slot = (32. * (f / low).log2()) as usize;
            if slot != last {
                points.push(Point::new(f, self.rate));
                last = slot;
            }
            let p = points.last_mut().expect("pushed above");
            if self.coherence(k) >= MIN_COHERENCE {
                p.fit += self.sab[k];
            }
            p.all += self.sab[k];
            p.power += self.saa[k] + self.sbb[k];
        }
        points
    }
    /// Produces an update once enough audio has been analysed since the last one.
    pub fn poll(&mut self) -> Option<Update> {
        let report_due = self.since_report as f32 >= REPORT_S * self.rate;
        if !report_due && !self.finish_learn {
            return None;
        }
        self.since_report = 0;
        let fit_due = self.since_fit as f32 >= FIT_S * self.rate;
        let mode = self.settings.mode;
        let threshold = if self.finish_learn {
            Some(LEARN_GAIN_DB)
        } else if fit_due && mode == Mode::Follow {
            Some(FOLLOW_GAIN_DB)
        } else if fit_due && mode == Mode::Learn {
            Some(LEARN_GAIN_DB)
        } else {
            None
        };
        if fit_due {
            self.since_fit = 0;
        }
        self.finish_learn = false;
        Some(self.analyse(threshold))
    }
    /// Measures, fits when `threshold` is given (minimum predicted gain in dB over the
    /// current fit), and reports.
    pub fn analyse(&mut self, threshold: Option<f32>) -> Update {
        let points = self.points();
        let problem = Problem::new(points, self.rate);
        let (weight, coherent) = problem.points.iter().fold((0., 0.), |(w, c), p| {
            let x = p.all.norm();
            (w + p.w * x, c + p.w * x * coherence_of(p))
        });
        let mean_coherence = if weight > 0. { coherent / weight } else { 0. };
        let (lag_ms, lag_strength) = self.lag();
        let late = lag_ms > LATE_S * 1000. && lag_strength > LATE_STRENGTH;
        let status = if self.side_level < SILENCE && self.frames == 0 {
            Status::NoSidechain
        } else if self.frames < MIN_FRAMES {
            if self.gated {
                Status::Quiet
            } else {
                Status::Waiting
            }
        } else if mean_coherence < LOOSE_COHERENCE || problem.is_empty() {
            Status::Loose
        } else if late {
            Status::Late
        } else {
            Status::Aligned
        };
        let mut changed = None;
        if let Some(threshold) = threshold {
            if matches!(status, Status::Aligned | Status::Late) {
                let candidate = search(&problem, self.settings.sections, &self.current);
                let gain = problem.mean_gain(
                    &problem.bands(&self.current),
                    &problem.bands(&candidate.fit),
                );
                if gain.is_finite() && gain >= threshold && candidate.fit != self.current {
                    self.current = candidate.fit;
                    changed = Some(candidate.fit);
                }
            }
        }
        let mut report = self.report(status, &problem);
        report.lag_ms = lag_ms;
        report.lag_strength = lag_strength;
        report.changed = changed.is_some();
        Update {
            fit: changed,
            report,
        }
    }
    fn report(&self, status: Status, problem: &Problem) -> Report {
        let fit = self.current;
        let mut report = Report {
            status,
            fit,
            frames: self.frames,
            low: self.range().0,
            high: self.range().1,
            smear_ms: fit.smear_ms(self.rate),
            ..Report::default()
        };
        // Display bands over all bins.
        // Sums per band: before, after, ideal, coherence·power, power.
        let mut sums = [[0f32; 5]; BANDS];
        for k in 1..=self.n / 2 {
            let f = self.frequency(k);
            if !(20. ..=0.45 * self.rate).contains(&f) {
                continue;
            }
            let (hr, hi) = fit.response(f, self.rate);
            let s = self.sab[k];
            let p = self.saa[k] + self.sbb[k];
            let b = &mut sums[band_of(f)];
            b[0] += p + 2. * s.re;
            b[1] += p + 2. * (s * Complex32::new(hr, hi)).re;
            b[2] += p + 2. * s.norm();
            b[3] += self.coherence(k) * p;
            b[4] += p;
        }
        for (band, [before, after, ideal, coherent, power]) in report.bands.iter_mut().zip(sums) {
            if ideal > 1e-12 && power > 0. {
                band.before = db_ratio(before, ideal);
                band.after = db_ratio(after, ideal);
                band.coherence = (coherent / power).clamp(0., 1.);
                band.valid = true;
            }
        }
        // Phase curve: bins within ±1/24 octave of each point.
        for (i, point) in report.curve.iter_mut().enumerate() {
            let f = curve_frequency(i);
            if f > 0.45 * self.rate {
                continue;
            }
            let lo = ((f * 2f32.powf(-1. / 24.)) * self.n as f32 / self.rate).floor() as usize;
            let hi = ((f * 2f32.powf(1. / 24.)) * self.n as f32 / self.rate).ceil() as usize;
            let centre = (f * self.n as f32 / self.rate).round() as usize;
            let (lo, hi) = (
                lo.max(1).min(centre.max(1)),
                hi.min(self.n / 2).max(centre.max(1)),
            );
            let (mut s, mut pa, mut pb) = (Complex32::default(), 0f32, 0f32);
            for k in lo..=hi {
                s += self.sab[k];
                pa += self.saa[k];
                pb += self.sbb[k];
            }
            if pa * pb <= 0. {
                continue;
            }
            let (hr, hi) = fit.response(f, self.rate);
            let h = Complex32::new(hr, hi);
            point.before = s.arg().to_degrees();
            point.after = (s * h).arg().to_degrees();
            point.chain = h.arg().to_degrees();
            point.coherence = (s.norm_sqr() / (pa * pb)).min(1.);
            point.valid = true;
        }
        // Headline numbers over the fit range, every band weighted the same.
        if !problem.is_empty() {
            let before = problem.before_bands();
            report.cost_db = problem.mean_gain(&before, &problem.ideal_bands());
            report.recovered_db = problem.mean_gain(&before, &problem.bands(&fit));
        }
        report
    }
    /// Time offset from the phase transform of the cross-spectrum (GCC-PHAT) between
    /// Low and High. Returns (ms, strength); positive means this layer is late.
    fn lag(&mut self) -> (f32, f32) {
        let n = self.n;
        let (low, high) = self.range();
        self.buf.fill(Complex32::default());
        let mut used = 0usize;
        for k in 1..n / 2 {
            let f = self.frequency(k);
            if f < low || f > high || self.coherence(k) < MIN_COHERENCE {
                continue;
            }
            let s = self.sab[k];
            let m = s.norm();
            if m <= 0. {
                continue;
            }
            let u = s / m;
            self.buf[k] = u;
            self.buf[n - k] = u.conj();
            used += 1;
        }
        if used < 4 {
            return (0., 0.);
        }
        self.ifft
            .process_with_scratch(&mut self.buf, &mut self.scratch);
        let range = (LAG_RANGE_S * self.rate) as isize;
        let (mut best, mut at) = (0f32, 0isize);
        for tau in -range..=range {
            let v = self.buf[tau.rem_euclid(n as isize) as usize].re.abs();
            if v > best {
                best = v;
                at = tau;
            }
        }
        (at as f32 / self.rate * 1000., best / (2 * used) as f32)
    }
}

fn coherence_of(p: &Point) -> f32 {
    if p.all.norm() > 0. {
        p.fit.norm() / p.all.norm()
    } else {
        0.
    }
}
fn db_ratio(a: f32, b: f32) -> f32 {
    (10. * (a.max(1e-20) / b.max(1e-20)).log10()).clamp(-40., 40.)
}