//! Fit quality on a synthetic saw bass, with targets taken from the prototype.
use rustfft::{num_complex::Complex64, FftPlanner};
use weld_core::{Analyzer, Chain, Fit, Mode, Settings, Status};

const RATE: f32 = 48000.;

/// Two seconds of an F1 saw (60 harmonics, fixed random phases) plus a little noise.
fn bass() -> Vec<f32> {
    let n = 2 * RATE as usize;
    let mut seed = 1u64;
    let mut rand = move || {
        seed = seed
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        (seed >> 11) as f64 / (1u64 << 53) as f64
    };
    let phases: Vec<f64> = (0..60).map(|_| rand() * std::f64::consts::TAU).collect();
    (0..n)
        .map(|i| {
            let t = i as f64 / RATE as f64;
            let saw: f64 = (1..60)
                .map(|k| (std::f64::consts::TAU * k as f64 * 87.3 * t + phases[k]).sin() / k as f64)
                .sum();
            (0.3 * saw + 0.015 * (rand() - 0.5)) as f32
        })
        .collect()
}
fn noise(n: usize, seed: u64) -> Vec<f32> {
    let mut s = seed;
    (0..n)
        .map(|_| {
            s = s
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            ((s >> 11) as f64 / (1u64 << 53) as f64 - 0.5) as f32 * 0.5
        })
        .collect()
}
fn delay(x: &[f32], samples: isize) -> Vec<f32> {
    (0..x.len() as isize)
        .map(|i| {
            let j = i - samples;
            if j >= 0 && (j as usize) < x.len() {
                x[j as usize]
            } else {
                0.
            }
        })
        .collect()
}
/// Direct-form biquad, b and a normalised by a0.
fn biquad(x: &[f32], b: [f64; 3], a: [f64; 3]) -> Vec<f32> {
    let (mut x1, mut x2, mut y1, mut y2) = (0f64, 0f64, 0f64, 0f64);
    x.iter()
        .map(|&v| {
            let v = v as f64;
            let y = (b[0] * v + b[1] * x1 + b[2] * x2 - a[1] * y1 - a[2] * y2) / a[0];
            (x2, x1, y2, y1) = (x1, v, y1, y);
            y as f32
        })
        .collect()
}
fn allpass(x: &[f32], fc: f64, q: f64) -> Vec<f32> {
    let w = std::f64::consts::TAU * fc / RATE as f64;
    let alpha = w.sin() / (2. * q);
    let a = [1. + alpha, -2. * w.cos(), 1. - alpha];
    biquad(x, [a[2], a[1], a[0]], a)
}
fn highpass(x: &[f32], fc: f64) -> Vec<f32> {
    let w = std::f64::consts::TAU * fc / RATE as f64;
    let alpha = w.sin() / (2. * std::f64::consts::FRAC_1_SQRT_2);
    let c = w.cos();
    biquad(
        x,
        [(1. + c) / 2., -(1. + c), (1. + c) / 2.],
        [1. + alpha, -2. * c, 1. - alpha],
    )
}
fn power_db(x: &[f32]) -> f64 {
    let skip = x.len() / 4;
    let p: f64 =
        x[skip..].iter().map(|&v| (v as f64).powi(2)).sum::<f64>() / (x.len() - skip) as f64;
    10. * p.max(1e-30).log10()
}
fn sum(a: &[f32], b: &[f32]) -> Vec<f32> {
    a.iter().zip(b).map(|(a, b)| a + b).collect()
}
fn apply(fit: Fit, x: &[f32]) -> Vec<f32> {
    let mut chain = Chain::new(fit, RATE);
    x.iter().map(|&v| chain.process(0, v)).collect()
}
/// Level of A + B when every bin of A is rotated onto B: the best any phase-only
/// processing could reach.
fn ideal_db(a: &[f32], b: &[f32]) -> f64 {
    let n = a.len();
    let mut planner = FftPlanner::<f64>::new();
    let fft = planner.plan_fft_forward(n);
    let ifft = planner.plan_fft_inverse(n);
    let mut fa: Vec<Complex64> = a.iter().map(|&v| Complex64::new(v as f64, 0.)).collect();
    let mut fb: Vec<Complex64> = b.iter().map(|&v| Complex64::new(v as f64, 0.)).collect();
    fft.process(&mut fa);
    fft.process(&mut fb);
    let mut y: Vec<Complex64> = fa
        .iter()
        .zip(&fb)
        .map(|(a, b)| {
            let unit = if b.norm() > 1e-12 {
                b / b.norm()
            } else {
                Complex64::new(1., 0.)
            };
            unit * a.norm() + b
        })
        .collect();
    ifft.process(&mut y);
    let out: Vec<f32> = y.iter().map(|v| (v.re / n as f64) as f32).collect();
    power_db(&out)
}
/// Power of A+B in octave bands from 60 Hz to 7.7 kHz, after minus before, in dB.
fn octave_changes(a: &[f32], aligned: &[f32], b: &[f32]) -> Vec<f64> {
    let n = a.len();
    let mut planner = FftPlanner::<f64>::new();
    let fft = planner.plan_fft_forward(n);
    let spectrum = |x: Vec<f32>| {
        let mut v: Vec<Complex64> = x.iter().map(|&s| Complex64::new(s as f64, 0.)).collect();
        fft.process(&mut v);
        v
    };
    let before = spectrum(sum(a, b));
    let after = spectrum(sum(aligned, b));
    (0..7)
        .map(|o| {
            let lo = 60. * 2f64.powi(o);
            let range = |s: &Vec<Complex64>| -> f64 {
                (0..n / 2)
                    .filter(|&k| {
                        let f = k as f64 * RATE as f64 / n as f64;
                        f >= lo && f < 2. * lo
                    })
                    .map(|k| s[k].norm_sqr())
                    .sum()
            };
            10. * (range(&after) / range(&before)).log10()
        })
        .collect()
}

struct Outcome {
    accepted: Option<Fit>,
    status: Status,
    lag_ms: f32,
    before: f64,
    after: f64,
    ideal: f64,
    octaves: Vec<f64>,
}
fn run(a: &[f32], b: &[f32]) -> Outcome {
    let mut analyzer = Analyzer::new(RATE, Fit::NONE);
    analyzer.set_settings(Settings {
        mode: Mode::Learn,
        low: 40.,
        high: 8000.,
        ..Settings::default()
    });
    for (x, y) in a.chunks(256).zip(b.chunks(256)) {
        analyzer.push(x, y);
    }
    let update = analyzer.analyse(Some(0.1));
    let fit = *analyzer.current();
    let aligned = apply(fit, a);
    Outcome {
        accepted: update.fit,
        status: update.report.status,
        lag_ms: update.report.lag_ms,
        before: power_db(&sum(a, b)),
        after: power_db(&sum(&aligned, b)),
        ideal: ideal_db(a, b),
        octaves: octave_changes(a, &aligned, b),
    }
}
fn show(name: &str, o: &Outcome) {
    eprintln!(
        "{name}: before {:.2} after {:.2} ideal {:.2} fit {:?} status {:?} lag {:.3} octaves {:?}",
        o.before,
        o.after,
        o.ideal,
        o.accepted.map(|f| (f.flip, f.active().to_vec())),
        o.status,
        o.lag_ms,
        o.octaves
            .iter()
            .map(|v| (v * 10.).round() / 10.)
            .collect::<Vec<_>>()
    );
}

#[test]
fn polarity_flip_needs_no_sections() {
    let a = bass();
    let b: Vec<f32> = a.iter().map(|v| -v).collect();
    let o = run(&a, &b);
    show("flip", &o);
    let fit = o.accepted.expect("a fit");
    assert!(fit.flip);
    assert_eq!(fit.count, 0);
    assert!((o.after - o.ideal).abs() < 0.05);
}

#[test]
fn an_allpass_is_found_again() {
    let a = bass();
    let b = allpass(&a, 400., 0.7);
    let o = run(&a, &b);
    show("allpass400", &o);
    let fit = o.accepted.expect("a fit");
    assert!(!fit.flip);
    assert_eq!(fit.count, 1, "{fit:?}");
    let s = fit.sections[0];
    assert!((s.fc / 400. - 1.).abs() < 0.05, "{s:?}");
    assert!((s.q / 0.7 - 1.).abs() < 0.1, "{s:?}");
    assert!(o.ideal - o.after < 0.1);
}

#[test]
fn a_late_sidechain_layer_is_caught_up() {
    let a = bass();
    for ms in [0.3, 1.0] {
        let b = delay(&a, (ms * 48.) as isize);
        let o = run(&a, &b);
        show(&format!("delay {ms}"), &o);
        assert!(o.ideal - o.after < 0.3, "{ms}");
        assert!(o.after > o.before);
        assert_ne!(o.status, Status::Late);
        // A delay leaves a comb in the sum: notches at odd multiples of 1/(2·delay).
        // A few allpass sections cannot follow a delay's phase all the way up, so they
        // move the notches rather than remove them; one octave may lose some level
        // while its neighbours gain more.
        let mean = o.octaves.iter().sum::<f64>() / o.octaves.len() as f64;
        assert!(mean > 0.4, "{ms} ms: mean octave change {mean}");
        for (i, v) in o.octaves.iter().enumerate() {
            assert!(*v > -2., "{ms} ms, octave {i}: {v}");
        }
    }
}

#[test]
fn an_inverted_highpassed_layer_is_aligned() {
    let a = bass();
    let b: Vec<f32> = highpass(&a, 300.).iter().map(|v| -v).collect();
    let o = run(&a, &b);
    show("inverted hp300", &o);
    assert!(o.ideal - o.after < 0.3);
}

#[test]
fn unrelated_layers_are_left_alone() {
    let a = bass();
    let b = noise(a.len(), 7);
    let o = run(&a, &b);
    show("unrelated", &o);
    assert!(o.accepted.is_none());
    assert_eq!(o.status, Status::Loose);
}

#[test]
fn a_late_main_layer_is_reported_and_not_made_worse() {
    let a = bass();
    let b = delay(&a, -14);
    let o = run(&a, &b);
    show("ahead", &o);
    assert_eq!(o.status, Status::Late);
    assert!((o.lag_ms - 14. / 48.).abs() < 0.03, "{}", o.lag_ms);
    assert!(o.after > o.before - 0.05);
}

#[test]
fn one_fit_is_fast() {
    let a = bass();
    let b = delay(&a, 48);
    let mut analyzer = Analyzer::new(RATE, Fit::NONE);
    analyzer.set_settings(Settings {
        mode: Mode::Learn,
        ..Settings::default()
    });
    for (x, y) in a.chunks(256).zip(b.chunks(256)) {
        analyzer.push(x, y);
    }
    let start = std::time::Instant::now();
    analyzer.analyse(Some(0.));
    let elapsed = start.elapsed();
    eprintln!("fit took {elapsed:?}");
    let limit = if cfg!(debug_assertions) { 1000 } else { 50 };
    assert!(elapsed.as_millis() < limit, "{elapsed:?}");
}