//! Levels each part of a bassline, a note or a step of the beat grid, to one loudness with a
//! single gain per part. No compression: the gain does not move inside a part, so a patch
//! keeps its own envelope and punch.
//!
//! A part's loudness is the K-weighted (ITU-R BS.1770) mean square of its first 100 ms, or
//! of all of it when it is shorter, counted in 10 ms blocks; blocks 20 dB under the loudest
//! one are left out, so a gap at the start does not pull the reading down. Each part is
//! weighted from rest, as if it played alone. The audio path is delayed by the note tracker's lookahead
//! plus that window plus a 5 ms ramp, so every part's gain is known before its attack leaves
//! the plugin, and the ramp to it ends exactly on the attack.
mod settings;
use common_dsp::loudness::{lufs, KWeighting};
use common_dsp::svf::{Svf, SvfCoef};
use common_dsp::tracker::{NoteEvent, Tracker, TrackerConfig};
use common_dsp::{coefficient, db_to_gain, gain_to_db};
pub use settings::{
    Grid, Segment, Settings, Source, FROM_MAX, FROM_MIN, RANGE_MAX, TARGET_MAX, TARGET_MIN,
};

/// Each part is measured over this much of its start.
pub const WINDOW_MS: f32 = 100.;
/// The gain glides to a new part's value over this long, ending on its first sample.
pub const RAMP_MS: f32 = 5.;
/// A boundary closer than this to the last one is ignored: the parts are merged.
const MIN_PART_MS: f32 = 10.;
/// The window is measured in blocks of this length.
const BLOCK_MS: f32 = 10.;
const BLOCKS: usize = (WINDOW_MS / BLOCK_MS) as usize;
/// Blocks this far under the part's loudest block do not count: a note start found by a
/// pitch change can come early, inside the gap before the note.
const RELATIVE_GATE_DB: f64 = 20.;
/// Parts shown in the editor.
pub const HISTORY: usize = 32;
/// Parts used for the spread readout.
const SPREAD: usize = 16;
const PENDING: usize = 32;
const Q: f32 = std::f32::consts::FRAC_1_SQRT_2;
/// Tempo the grid runs at before the host reports one.
const FALLBACK_BPM: f64 = 120.;

#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Part {
    /// Loudness of the part's measured start, in LUFS.
    pub loudness: f32,
    /// Gain given to the part, in dB.
    pub gain: f32,
    /// Under the gate: the part kept the previous gain.
    pub gated: bool,
}

#[derive(Clone, Copy, Debug)]
pub struct AnalyzerFrame {
    /// Latest parts, oldest first; the first `count` entries are valid.
    pub history: [Part; HISTORY],
    pub count: usize,
    /// Gain applied at the output right now, in dB.
    pub gain: f32,
    /// Standard deviation of the loudness of the last 16 parts over the gate, before and after
    /// Ballast. NaN with fewer than two.
    pub spread_in: f32,
    pub spread_out: f32,
    /// Waiting for the next part to take as the target.
    pub learning: bool,
    /// The loudness last taken by Learn, and a number that changes each time one is taken.
    pub learned: f32,
    pub learn_id: u32,
    pub tempo: f64,
    pub synced: bool,
    pub latency: u32,
    pub sample_rate: f32,
}
impl Default for AnalyzerFrame {
    fn default() -> Self {
        Self {
            history: [Part::default(); HISTORY],
            count: 0,
            gain: 0.,
            spread_in: f32::NAN,
            spread_out: f32::NAN,
            learning: false,
            learned: f32::NAN,
            learn_id: 0,
            tempo: FALLBACK_BPM,
            synced: false,
            latency: 0,
            sample_rate: 48000.,
        }
    }
}

/// Gain in dB that brings `loudness` to the target, at `amount` (0–1), within ±`range`.
pub fn part_gain(loudness: f32, target: f32, amount: f32, range: f32) -> f32 {
    ((target - loudness) * amount).clamp(-range, range)
}

#[derive(Clone, Copy, Default)]
struct Pending {
    start: u64,
    gain: f32,
}

pub struct Engine {
    rate: f32,
    channels: usize,
    settings: Settings,
    tracker: Tracker,
    lookahead: usize,
    window: usize,
    ramp: usize,
    min_part: usize,
    delay: Vec<[f32; 2]>,
    delay_at: usize,
    weighting: [KWeighting; 2],
    from: [[Svf; 2]; 2],
    from_coef: SvfCoef,
    /// Samples of the aligned stream: the input delayed by the tracker's lookahead.
    clock: u64,
    open: bool,
    part_start: u64,
    block: usize,
    blocks: [(f64, usize); BLOCKS],
    measured: usize,
    last_gain: f32,
    queue: [Pending; PENDING],
    queue_head: usize,
    queue_len: usize,
    gain: f32,
    ramp_from: f32,
    ramp_to: f32,
    ramp_left: usize,
    history: [Part; HISTORY],
    history_at: usize,
    history_len: usize,
    learn: bool,
    learned: f32,
    learn_id: u32,
    tempo: f64,
    /// Beat position of the aligned stream.
    beats: f64,
    synced: bool,
    cell: i64,
    output: f32,
    bypass: f32,
    smooth: f32,
    fresh: bool,
}

impl Engine {
    pub fn new(rate: f32, channels: usize) -> Result<Self, &'static str> {
        if !(8000.0..=384000.).contains(&rate) {
            return Err("unsupported sample rate");
        }
        if !(1..=2).contains(&channels) {
            return Err("unsupported channel count");
        }
        let tracker = Tracker::new(rate, TrackerConfig::BASS);
        let lookahead = tracker.latency();
        let samples = |ms: f32| ((rate * ms * 0.001).round() as usize).max(1);
        let (window, ramp) = (samples(WINDOW_MS), samples(RAMP_MS));
        let mut e = Self {
            rate,
            channels,
            settings: Settings::default(),
            tracker,
            lookahead,
            window,
            ramp,
            min_part: samples(MIN_PART_MS),
            delay: vec![[0.; 2]; lookahead + window + ramp + 1],
            delay_at: 0,
            weighting: [KWeighting::new(rate); 2],
            from: [[Svf::default(); 2]; 2],
            from_coef: SvfCoef::IDENTITY,
            clock: 0,
            open: false,
            part_start: 0,
            block: samples(BLOCK_MS),
            blocks: [(0., 0); BLOCKS],
            measured: 0,
            last_gain: 0.,
            queue: [Pending::default(); PENDING],
            queue_head: 0,
            queue_len: 0,
            gain: 1.,
            ramp_from: 1.,
            ramp_to: 1.,
            ramp_left: 0,
            history: [Part::default(); HISTORY],
            history_at: 0,
            history_len: 0,
            learn: false,
            learned: f32::NAN,
            learn_id: 0,
            tempo: FALLBACK_BPM,
            beats: 0.,
            synced: false,
            cell: 0,
            output: 1.,
            bypass: 0.,
            smooth: coefficient(rate, 10.),
            fresh: true,
        };
        e.set_settings(Settings::default())?;
        e.fresh = true;
        Ok(e)
    }
    /// Lookahead of the tracker, the measuring window and the ramp.
    pub fn latency_samples(&self) -> u32 {
        (self.lookahead + self.window + self.ramp) as u32
    }
    pub fn channels(&self) -> usize {
        self.channels
    }
    pub fn settings(&self) -> Settings {
        self.settings
    }
    pub fn set_settings(&mut self, s: Settings) -> Result<(), &'static str> {
        s.validate()?;
        self.tracker.set(s.source == Source::Midi, s.gate, false);
        if s.from != self.settings.from || self.fresh {
            self.from_coef = if s.from > FROM_MIN {
                SvfCoef::highpass(self.rate, s.from, Q)
            } else {
                SvfCoef::IDENTITY
            };
        }
        self.settings = s;
        if std::mem::take(&mut self.fresh) {
            self.output = db_to_gain(s.output);
            self.bypass = f32::from(s.bypass);
        }
        Ok(())
    }
    /// Host tempo in BPM and, while playing, the beat position at the block's first sample.
    pub fn set_transport(&mut self, tempo: Option<f64>, position: Option<f64>) {
        if let Some(t) = tempo.filter(|t| t.is_finite() && *t > 0.) {
            self.tempo = t;
        }
        match position.filter(|p| p.is_finite()) {
            Some(p) => {
                // The parts are cut on the aligned stream, `lookahead` samples behind the input.
                self.beats = p - self.lookahead as f64 / self.rate as f64 * self.tempo / 60.;
                self.synced = true;
            }
            None => self.synced = false,
        }
    }
    pub fn note_on(&mut self, note: u8) {
        self.tracker.note_on(note);
    }
    pub fn note_off(&mut self, note: u8) {
        self.tracker.note_off(note);
    }
    /// Takes the loudness of the next part over the gate as the new target.
    pub fn learn(&mut self) {
        self.learn = true;
    }
    pub fn reset(&mut self) {
        self.tracker.reset();
        self.delay.fill([0.; 2]);
        self.delay_at = 0;
        for k in &mut self.weighting {
            k.reset();
        }
        self.from = [[Svf::default(); 2]; 2];
        self.clock = 0;
        self.open = false;
        self.blocks = [(0., 0); BLOCKS];
        self.measured = 0;
        self.last_gain = 0.;
        self.queue_len = 0;
        self.gain = 1.;
        self.ramp_left = 0;
        self.history_len = 0;
        self.learn = false;
        self.fresh = true;
        let s = self.settings;
        let _ = self.set_settings(s);
    }

    /// Mean square of the measured blocks, without those far under the loudest.
    fn gated_mean(&self) -> f64 {
        let used = self.measured.div_ceil(self.block).min(BLOCKS);
        let blocks = &self.blocks[..used];
        let loudest = blocks
            .iter()
            .map(|(s, n)| s / (*n).max(1) as f64)
            .fold(0., f64::max);
        let floor = loudest * 10f64.powf(-RELATIVE_GATE_DB / 10.);
        let (sum, n) = blocks
            .iter()
            .filter(|(s, n)| s / (*n).max(1) as f64 >= floor)
            .fold((0., 0), |(a, b), (s, n)| (a + s, b + n));
        sum / n.max(1) as f64
    }
    fn close(&mut self) {
        self.open = false;
        let s = self.settings;
        let loudness = lufs(self.gated_mean());
        let gated = loudness < s.gate;
        let gain = if gated {
            self.last_gain
        } else {
            part_gain(loudness, s.target, s.amount * 0.01, s.range)
        };
        if !gated {
            self.last_gain = gain;
            if std::mem::take(&mut self.learn) {
                self.learned = loudness;
                self.learn_id = self.learn_id.wrapping_add(1);
            }
        }
        if self.queue_len < PENDING {
            let at = (self.queue_head + self.queue_len) % PENDING;
            self.queue[at] = Pending {
                start: self.part_start,
                gain,
            };
            self.queue_len += 1;
        }
        self.history[self.history_at] = Part {
            loudness,
            gain,
            gated,
        };
        self.history_at = (self.history_at + 1) % HISTORY;
        self.history_len = (self.history_len + 1).min(HISTORY);
    }

    // Channels index the delayed frame, the filters and the output at once.
    #[allow(clippy::needless_range_loop)]
    pub fn process_sample(&mut self, x: [f32; 2]) -> [f32; 2] {
        let mono = if self.channels == 2 {
            0.5 * (x[0] + x[1])
        } else {
            x[0]
        };
        let mut started = false;
        self.tracker.process(mono, |e| {
            if matches!(e, NoteEvent::Start { .. }) {
                started = true;
            }
        });
        let len = self.delay.len();
        self.delay[self.delay_at] = x;
        let aligned = self.delay[(self.delay_at + len - self.lookahead) % len];
        self.delay_at = (self.delay_at + 1) % len;
        let delayed = self.delay[self.delay_at];

        // Grid lines are counted on the aligned stream even in Notes mode, so switching modes
        // does not cut a part.
        let cell = (self.beats / self.settings.grid.quarters()).floor() as i64;
        self.beats += self.tempo / 60. / self.rate as f64;
        let crossed = cell != self.cell;
        self.cell = cell;
        let boundary = match self.settings.segment {
            Segment::Notes => started,
            Segment::Grid => crossed,
        };

        if boundary && !(self.open && self.measured < self.min_part) {
            if self.open {
                self.close();
            }
            self.open = true;
            self.part_start = self.clock;
            self.blocks = [(0., 0); BLOCKS];
            self.measured = 0;
            // Each part is weighted from rest, as if it played alone: the filters' memory of
            // the last part would otherwise ring into this one.
            for k in &mut self.weighting {
                k.reset();
            }
            self.from = [[Svf::default(); 2]; 2];
        }
        let mut power = 0f64;
        for c in 0..self.channels {
            let f = &mut self.from[c];
            let v = f[0].process(&self.from_coef, aligned[c]);
            let v = f[1].process(&self.from_coef, v);
            let w = self.weighting[c].process(v) as f64;
            power += w * w;
        }
        if self.open {
            let b = &mut self.blocks[(self.measured / self.block).min(BLOCKS - 1)];
            b.0 += power;
            b.1 += 1;
            self.measured += 1;
            if self.measured >= self.window {
                self.close();
            }
        }
        // A part's ramp starts `ramp` samples before its first sample reaches the output.
        if self.queue_len > 0 {
            let next = self.queue[self.queue_head];
            if self.clock >= next.start + self.window as u64 {
                self.queue_head = (self.queue_head + 1) % PENDING;
                self.queue_len -= 1;
                self.ramp_from = self.gain;
                self.ramp_to = db_to_gain(next.gain);
                self.ramp_left = self.ramp;
            }
        }
        if self.ramp_left > 0 {
            self.ramp_left -= 1;
            let t = 1. - self.ramp_left as f32 / self.ramp as f32;
            let t = t * t * (3. - 2. * t);
            self.gain = self.ramp_from + (self.ramp_to - self.ramp_from) * t;
        }
        self.clock += 1;

        let s = self.settings;
        self.output += (db_to_gain(s.output) - self.output) * self.smooth;
        self.bypass += (f32::from(s.bypass) - self.bypass) * self.smooth;
        if (f32::from(s.bypass) - self.bypass).abs() < 1e-6 {
            self.bypass = f32::from(s.bypass);
        }
        let mut out = [0.; 2];
        for c in 0..self.channels {
            let wet = delayed[c] * self.gain * self.output;
            out[c] = if self.bypass >= 1. {
                delayed[c]
            } else {
                wet + (delayed[c] - wet) * self.bypass
            };
        }
        out
    }

    fn spread(&self, after: bool) -> f32 {
        let (mut n, mut sum, mut sq) = (0usize, 0f64, 0f64);
        for i in 0..self.history_len {
            if n == SPREAD {
                break;
            }
            let p = self.history[(self.history_at + HISTORY - 1 - i) % HISTORY];
            if p.gated {
                continue;
            }
            let v = (p.loudness + if after { p.gain } else { 0. }) as f64;
            n += 1;
            sum += v;
            sq += v * v;
        }
        if n < 2 {
            return f32::NAN;
        }
        let mean = sum / n as f64;
        ((sq / n as f64 - mean * mean).max(0.)).sqrt() as f32
    }

    /// Readings for the editor. Never allocates.
    pub fn take_analyzer(&self) -> AnalyzerFrame {
        let mut f = AnalyzerFrame {
            count: self.history_len,
            gain: gain_to_db(self.gain),
            spread_in: self.spread(false),
            spread_out: self.spread(true),
            learning: self.learn,
            learned: self.learned,
            learn_id: self.learn_id,
            tempo: self.tempo,
            synced: self.synced,
            latency: self.latency_samples(),
            sample_rate: self.rate,
            ..AnalyzerFrame::default()
        };
        let first = (self.history_at + HISTORY - self.history_len) % HISTORY;
        for i in 0..self.history_len {
            f.history[i] = self.history[(first + i) % HISTORY];
        }
        f
    }
}