use ballast_core::{part_gain, Engine, Grid, Segment, Settings, Source, RAMP_MS, WINDOW_MS};
use common_dsp::loudness::{lufs, KWeighting};
use std::f64::consts::TAU;

const RATE: f32 = 48000.;
const BPM: f64 = 120.;
/// One quarter at 120 BPM.
const QUARTER: usize = 24000;

/// Two patches a switch-up might alternate: a loud saw and a quiet sine, an octave apart.
fn patch(i: usize, which: usize) -> f32 {
    let t = i as f64 / RATE as f64;
    if which.is_multiple_of(2) {
        let phase = (55. * t).fract();
        (0.5 * (2. * phase - 1.)) as f32
    } else {
        (0.1 * (TAU * 110. * t).sin()) as f32
    }
}
fn grid_engine(s: Settings) -> Engine {
    let mut e = Engine::new(RATE, 2).unwrap();
    e.set_settings(Settings {
        segment: Segment::Grid,
        grid: Grid::Quarter,
        ..s
    })
    .unwrap();
    e
}
/// Runs `input` through the engine in 512-sample blocks with the transport playing from 0.
fn run(e: &mut Engine, input: &[f32]) -> Vec<f32> {
    let mut out = Vec::with_capacity(input.len());
    for (i, &v) in input.iter().enumerate() {
        if i % 512 == 0 {
            e.set_transport(Some(BPM), Some(i as f64 / RATE as f64 * BPM / 60.));
        }
        out.push(e.process_sample([v, v])[0]);
    }
    out
}
/// K-weighted loudness of both channels (the same signal) over `window` samples from `start`.
fn loudness(x: &[f32], start: usize, window: usize) -> f32 {
    let mut k = KWeighting::new(RATE);
    // Settle the filter on what came before, as the engine's own measurement does.
    for &v in &x[start.saturating_sub(4800)..start] {
        k.process(v);
    }
    let sum: f64 = x[start..start + window]
        .iter()
        .map(|&v| (k.process(v) as f64).powi(2))
        .sum();
    lufs(2. * sum / window as f64)
}
fn window() -> usize {
    (RATE * WINDOW_MS * 0.001) as usize
}

#[test]
fn every_part_lands_at_the_target() {
    let target = -18.;
    let mut e = grid_engine(Settings {
        target,
        ..Settings::default()
    });
    let latency = e.latency_samples() as usize;
    let parts = 10;
    let input: Vec<f32> = (0..parts * QUARTER + latency)
        .map(|i| patch(i, i / QUARTER))
        .collect();
    let out = run(&mut e, &input);
    for p in 2..parts {
        let start = p * QUARTER;
        let before = loudness(&input, start, window());
        let after = loudness(&out, start + latency, window());
        assert!(
            (after - target).abs() < 0.3,
            "part {p}: {before:.1} → {after:.1}"
        );
    }
    let f = e.take_analyzer();
    assert!(f.spread_in > 5., "{}", f.spread_in);
    assert!(f.spread_out < 0.1, "{}", f.spread_out);
}

#[test]
fn the_gain_holds_still_inside_a_part_and_is_in_place_on_its_attack() {
    let mut e = grid_engine(Settings {
        target: -18.,
        ..Settings::default()
    });
    let latency = e.latency_samples() as usize;
    let ramp = (RATE * RAMP_MS * 0.001) as usize;
    let input: Vec<f32> = (0..6 * QUARTER + latency)
        .map(|i| patch(i, i / QUARTER) + 0.2)
        .collect();
    let out = run(&mut e, &input);
    for p in 2..5 {
        let start = p * QUARTER;
        assert!(input[start].abs() > 0.05);
        let gain = out[start + latency] / input[start];
        for n in start..start + QUARTER - ramp {
            if input[n].abs() < 0.05 {
                continue;
            }
            let g = out[n + latency] / input[n];
            assert!(
                (g / gain - 1.).abs() < 1e-4,
                "part {p}, sample {n}: {g} vs {gain}"
            );
        }
    }
}

#[test]
fn note_starts_from_the_audio_cut_the_parts() {
    let target = -18.;
    let mut e = Engine::new(RATE, 2).unwrap();
    e.set_settings(Settings {
        target,
        ..Settings::default()
    })
    .unwrap();
    let latency = e.latency_samples() as usize;
    // 300 ms notes with 100 ms gaps, alternating the two patches.
    let note = 19200;
    let input: Vec<f32> = (0..8 * note + latency)
        .map(|i| {
            if i % note < 14400 && i < 8 * note {
                patch(i, i / note)
            } else {
                0.
            }
        })
        .collect();
    let out = run(&mut e, &input);
    let f = e.take_analyzer();
    let parts = f.history[..f.count].iter().filter(|p| !p.gated).count();
    assert_eq!(parts, 8, "{:?}", &f.history[..f.count]);
    for n in 2..8 {
        // The tracker finds the start a few milliseconds in; measure past that.
        let start = n * note + 960;
        let after = loudness(&out, start + latency, window());
        assert!((after - target).abs() < 0.5, "note {n}: {after:.1}");
    }
}

#[test]
fn midi_note_starts_cut_legato_parts() {
    let target = -18.;
    let mut e = Engine::new(RATE, 2).unwrap();
    e.set_settings(Settings {
        target,
        source: Source::Midi,
        ..Settings::default()
    })
    .unwrap();
    let latency = e.latency_samples() as usize;
    let note = QUARTER;
    let mut out = Vec::new();
    let input: Vec<f32> = (0..6 * note + latency)
        .map(|i| patch(i, i / note))
        .collect();
    for (i, &v) in input.iter().enumerate() {
        if i % note == 0 {
            // Legato: the new note starts before the old one ends, on the same pitch.
            e.note_on(33);
            if i > 0 {
                e.note_off(33);
                e.note_on(33);
            }
        }
        out.push(e.process_sample([v, v])[0]);
    }
    for n in 2..6 {
        let after = loudness(&out, n * note + latency, window());
        assert!((after - target).abs() < 0.3, "note {n}: {after:.1}");
    }
}

#[test]
fn quiet_parts_keep_the_last_gain() {
    let mut e = grid_engine(Settings {
        target: -18.,
        gate: -40.,
        ..Settings::default()
    });
    // A loud part, then a part far under the gate, then the loud patch again. The parts change
    // on zero crossings of a 110 Hz sine, so no step rings into the quiet part.
    let input: Vec<f32> = (0..4 * QUARTER)
        .map(|i| {
            let v = 0.5 * (TAU * 110. * i as f64 / RATE as f64).sin() as f32;
            if i / QUARTER == 2 {
                v * 0.001
            } else {
                v
            }
        })
        .collect();
    run(&mut e, &input);
    let f = e.take_analyzer();
    let parts = &f.history[..f.count];
    // The first part is the silence before the input reaches the measurement.
    let quiet = 1 + parts[1..]
        .iter()
        .position(|p| p.gated)
        .expect("a gated part");
    assert_eq!(parts[quiet].gain, parts[quiet - 1].gain);
}

#[test]
fn amount_and_range_scale_the_correction() {
    assert_eq!(part_gain(-10., -14., 1., 12.), -4.);
    assert_eq!(part_gain(-10., -14., 0.5, 12.), -2.);
    assert_eq!(part_gain(-40., -14., 1., 12.), 12.);
    assert_eq!(part_gain(0., -14., 1., 6.), -6.);
    let mut e = grid_engine(Settings {
        target: -18.,
        amount: 50.,
        ..Settings::default()
    });
    let input: Vec<f32> = (0..4 * QUARTER).map(|i| patch(i, i / QUARTER)).collect();
    run(&mut e, &input);
    let f = e.take_analyzer();
    for p in &f.history[1..f.count] {
        let full = -18. - p.loudness;
        assert!((p.gain - 0.5 * full).abs() < 1e-3, "{p:?}");
    }
}

#[test]
fn learn_takes_the_next_part_as_the_target() {
    let mut e = grid_engine(Settings::default());
    let input: Vec<f32> = (0..3 * QUARTER).map(|i| patch(i, i / QUARTER)).collect();
    run(&mut e, &input[..QUARTER + 100]);
    e.learn();
    assert!(e.take_analyzer().learning);
    let before = e.take_analyzer().learn_id;
    run(&mut e, &input[QUARTER + 100..]);
    let f = e.take_analyzer();
    assert!(!f.learning);
    assert_ne!(f.learn_id, before);
    let taken = f.history[..f.count]
        .iter()
        .find(|p| p.loudness == f.learned)
        .expect("learned a part's loudness");
    assert!(!taken.gated);
}

#[test]
fn the_measurement_can_leave_out_the_sub() {
    // The same top, once with a loud sub under it.
    let top = |i: usize| (0.05 * (TAU * 1000. * i as f64 / RATE as f64).sin()) as f32;
    let sub = |i: usize| (0.5 * (TAU * 50. * i as f64 / RATE as f64).sin()) as f32;
    let gains = |from: f32| {
        let mut e = grid_engine(Settings {
            target: -24.,
            range: 24.,
            from,
            ..Settings::default()
        });
        let input: Vec<f32> = (0..4 * QUARTER)
            .map(|i| {
                top(i)
                    + if (i / QUARTER).is_multiple_of(2) {
                        sub(i)
                    } else {
                        0.
                    }
            })
            .collect();
        run(&mut e, &input);
        let f = e.take_analyzer();
        let p = &f.history[..f.count];
        (p[p.len() - 2].gain - p[p.len() - 1].gain).abs()
    };
    assert!(gains(20.) > 10.);
    assert!(gains(200.) < 1., "{}", gains(200.));
}

#[test]
fn bypass_is_the_input_delayed_by_the_latency() {
    let mut e = grid_engine(Settings {
        bypass: true,
        target: -30.,
        ..Settings::default()
    });
    let latency = e.latency_samples() as usize;
    let lookahead = latency - window() - (RATE * RAMP_MS * 0.001) as usize;
    // The tracker's lookahead: two periods of 30 Hz and a margin, about 67 ms.
    assert!((lookahead as f32 / RATE - 0.067).abs() < 0.005);
    let input: Vec<f32> = (0..2 * QUARTER).map(|i| patch(i, i / QUARTER)).collect();
    let out = run(&mut e, &input);
    for i in latency..input.len() {
        assert_eq!(out[i], input[i - latency]);
    }
}

#[test]
fn settings_are_validated_and_serialise() {
    assert!(Settings::default().validate().is_ok());
    for bad in [
        Settings {
            target: 3.,
            ..Settings::default()
        },
        Settings {
            range: 30.,
            ..Settings::default()
        },
        Settings {
            from: 10.,
            ..Settings::default()
        },
        Settings {
            amount: f32::NAN,
            ..Settings::default()
        },
    ] {
        assert!(bad.validate().is_err());
        assert!(Engine::new(RATE, 2).unwrap().set_settings(bad).is_err());
    }
    let s = Settings {
        segment: Segment::Grid,
        grid: Grid::Eighth,
        source: Source::Midi,
        target: -9.,
        ..Settings::default()
    };
    let json = serde_json::to_string(&s).unwrap();
    assert_eq!(serde_json::from_str::<Settings>(&json).unwrap(), s);
    assert!(serde_json::from_str::<Settings>(r#"{"target":0,"oops":1}"#).is_err());
}