//! Note tracker with lookahead: pitch and note starts of a monophonic bass, from audio or MIDI.
//!
//! Detection runs on the undelayed input. Events are reported at *output* time, `latency()`
//! samples later, so a host that delays its audio path by the same amount sees each note start
//! exactly on its attack, with the pitch already known.
use crate::{coefficient, gain_to_db, hz_to_midi, Biquad, Yin};

const EVENTS: usize = 64;
const HOP_MS: f32 = 5.;
const GATE_HOLD_MS: f32 = 20.;
const GATE_HYSTERESIS_DB: f32 = 6.;
/// A pitch move this large (semitones) counts as a new note.
const JUMP: f32 = 0.7;
const CLARITY: f32 = 0.5;

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TrackerConfig {
    /// Lowest detectable fundamental. Sets the lookahead (two periods plus a margin).
    pub min_hz: f32,
    pub max_hz: f32,
    /// Lowpass in front of detection, removing upper harmonics that confuse YIN.
    pub detect_lowpass: f32,
    /// Approximate decimated analysis rate.
    pub detect_rate: f32,
}
impl TrackerConfig {
    /// Sub range: 25–250 Hz.
    pub const SUB: Self = Self {
        min_hz: 25.,
        max_hz: 250.,
        detect_lowpass: 400.,
        detect_rate: 3000.,
    };
    /// Bass range for harmonic tools: 30–500 Hz.
    pub const BASS: Self = Self {
        min_hz: 30.,
        max_hz: 500.,
        detect_lowpass: 1000.,
        detect_rate: 6000.,
    };
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum NoteEvent {
    /// A new note. `glide` asks to slide from the previous pitch instead of restarting.
    Start {
        pitch: f32,
        glide: bool,
    },
    /// The sounding note's smoothed pitch moved within a semitone (bends, vibrato).
    Move {
        pitch: f32,
    },
    Stop,
}

#[derive(Clone, Copy, Debug, PartialEq)]
enum Kind {
    Onset,
    Release,
    NoteOn(u8),
    NoteOff(u8),
}
#[derive(Clone, Copy, Debug)]
struct Event {
    time: u64,
    kind: Kind,
}

/// Allocation-free after `new`. Pitches are fractional MIDI notes.
pub struct Tracker {
    rate: f32,
    cfg: TrackerConfig,
    latency: usize,
    midi: bool,
    gate_db: f32,
    glide: bool,
    detect: [Biquad; 2],
    decimate: usize,
    decimate_count: usize,
    detect_rate: f32,
    ring: Vec<f32>,
    ring_at: usize,
    scratch: Vec<f32>,
    yin: Yin,
    hop: usize,
    hop_count: usize,
    gate_env: f32,
    gate_release: f32,
    gate_open: bool,
    gate_hold: usize,
    clock: u64,
    events: [Event; EVENTS],
    event_head: usize,
    event_len: usize,
    /// Output-time note state.
    sounding: bool,
    pending: bool,
    tracked: f32,
    candidate: Option<(f32, u8)>,
    held: [u8; 16],
    held_len: usize,
}

impl Tracker {
    pub fn new(rate: f32, cfg: TrackerConfig) -> Self {
        let decimate = (rate / cfg.detect_rate).round().max(1.) as usize;
        let detect_rate = rate / decimate as f32;
        let max_lag = (detect_rate / cfg.min_hz).ceil() as usize;
        // YIN needs one comparison window of at least the longest period after the lags.
        let analysis = 2 * max_lag + 3;
        let q = std::f32::consts::FRAC_1_SQRT_2;
        Self {
            rate,
            cfg,
            latency: analysis * decimate,
            midi: false,
            gate_db: -42.,
            glide: false,
            detect: [Biquad::lowpass(rate, cfg.detect_lowpass, q); 2],
            decimate,
            decimate_count: 0,
            detect_rate,
            ring: vec![0.; analysis],
            ring_at: 0,
            scratch: vec![0.; analysis],
            yin: Yin::new(max_lag + 2),
            hop: ((detect_rate * HOP_MS * 0.001) as usize).max(1),
            hop_count: 0,
            gate_env: 0.,
            gate_release: 1. - coefficient(rate, 30.),
            gate_open: false,
            gate_hold: 0,
            clock: 0,
            events: [Event {
                time: 0,
                kind: Kind::Release,
            }; EVENTS],
            event_head: 0,
            event_len: 0,
            sounding: false,
            pending: false,
            tracked: 0.,
            candidate: None,
            held: [0; 16],
            held_len: 0,
        }
    }
    /// Lookahead in samples. Delay the audio path by exactly this much.
    pub fn latency(&self) -> usize {
        self.latency
    }
    pub fn config(&self) -> TrackerConfig {
        self.cfg
    }
    /// `midi` selects note input over detection; `gate_db` is the detection threshold;
    /// `glide` lets legato notes slide. Returns true when the source changed, which drops
    /// every pending and held note: the caller should release its sound.
    pub fn set(&mut self, midi: bool, gate_db: f32, glide: bool) -> bool {
        let changed = midi != self.midi;
        self.midi = midi;
        self.gate_db = gate_db;
        self.glide = glide;
        if changed {
            self.release_all();
        }
        changed
    }
    fn release_all(&mut self) {
        self.sounding = false;
        self.pending = false;
        self.held_len = 0;
        self.event_len = 0;
        self.candidate = None;
    }
    pub fn reset(&mut self) {
        for f in &mut self.detect {
            f.reset();
        }
        self.ring.fill(0.);
        self.decimate_count = 0;
        self.hop_count = 0;
        self.gate_env = 0.;
        self.gate_open = false;
        self.gate_hold = 0;
        self.release_all();
    }
    /// MIDI note start at the current input sample. Used when the source is MIDI.
    pub fn note_on(&mut self, note: u8) {
        self.push(Kind::NoteOn(note.min(127)));
    }
    pub fn note_off(&mut self, note: u8) {
        self.push(Kind::NoteOff(note.min(127)));
    }
    /// True while a note sounds at output time.
    pub fn sounding(&self) -> bool {
        self.sounding
    }
    fn push(&mut self, kind: Kind) {
        if self.event_len == EVENTS {
            // Dropping the oldest keeps the newest note state; overflow only happens on floods.
            self.event_head = (self.event_head + 1) % EVENTS;
            self.event_len -= 1;
        }
        let at = (self.event_head + self.event_len) % EVENTS;
        self.events[at] = Event {
            time: self.clock,
            kind,
        };
        self.event_len += 1;
    }
    /// Detected pitch of the most recent analysis window, in fractional MIDI notes.
    fn estimate(&mut self) -> Option<f32> {
        let n = self.ring.len();
        for i in 0..n {
            self.scratch[i] = self.ring[(self.ring_at + i) % n];
        }
        let (hz, clarity) = self.yin.estimate(
            &self.scratch,
            self.detect_rate,
            self.cfg.min_hz,
            self.cfg.max_hz,
        )?;
        (clarity >= CLARITY).then(|| hz_to_midi(hz))
    }
    fn start(&mut self, pitch: f32, glide: bool, emit: &mut impl FnMut(NoteEvent)) {
        self.tracked = pitch;
        self.candidate = None;
        self.pending = false;
        self.sounding = true;
        emit(NoteEvent::Start { pitch, glide });
    }
    fn stop(&mut self, emit: &mut impl FnMut(NoteEvent)) {
        self.sounding = false;
        self.pending = false;
        self.candidate = None;
        emit(NoteEvent::Stop);
    }
    fn handle(&mut self, kind: Kind, emit: &mut impl FnMut(NoteEvent)) {
        let midi = self.midi;
        match kind {
            Kind::Onset if !midi => match self.estimate() {
                Some(p) => self.start(p, false, emit),
                None => {
                    self.pending = true;
                    self.sounding = true;
                }
            },
            Kind::Release if !midi => self.stop(emit),
            Kind::NoteOn(n) if midi => {
                let len = self.held_len;
                if let Some(i) = self.held[..len].iter().position(|h| *h == n) {
                    self.held.copy_within(i + 1..len, i);
                    self.held_len -= 1;
                }
                if self.held_len == self.held.len() {
                    self.held.copy_within(1.., 0);
                    self.held_len -= 1;
                }
                let legato = self.held_len > 0;
                self.held[self.held_len] = n;
                self.held_len += 1;
                self.start(n as f32, legato && self.glide, emit);
            }
            Kind::NoteOff(n) if midi => {
                let len = self.held_len;
                let Some(i) = self.held[..len].iter().position(|h| *h == n) else {
                    return;
                };
                let top = i + 1 == len;
                self.held.copy_within(i + 1..len, i);
                self.held_len -= 1;
                if self.held_len == 0 {
                    self.stop(emit);
                } else if top {
                    let next = self.held[self.held_len - 1];
                    self.start(next as f32, self.glide, emit);
                }
            }
            _ => {}
        }
    }
    /// Regular analysis; its window starts at the current output time.
    fn hop_analysis(&mut self, emit: &mut impl FnMut(NoteEvent)) {
        if self.midi || !self.sounding {
            return;
        }
        let Some(p) = self.estimate() else {
            return;
        };
        if self.pending {
            self.start(p, false, emit);
            return;
        }
        if (p - self.tracked).abs() >= JUMP {
            match self.candidate {
                Some((c, count)) if (c - p).abs() < 0.5 => {
                    if count + 1 >= 2 {
                        self.start(p, self.glide, emit);
                    } else {
                        self.candidate = Some((p, count + 1));
                    }
                }
                _ => self.candidate = Some((p, 1)),
            }
        } else {
            self.candidate = None;
            self.tracked += 0.3 * (p - self.tracked);
            emit(NoteEvent::Move {
                pitch: self.tracked,
            });
        }
    }
    /// Feeds one undelayed mono input sample. Events due at output time go to `emit`, in order.
    pub fn process(&mut self, mono: f32, mut emit: impl FnMut(NoteEvent)) {
        let d = self.detect[0].process(mono);
        let d = self.detect[1].process(d);
        if !self.midi {
            let a = d.abs();
            self.gate_env = if a > self.gate_env {
                a
            } else {
                self.gate_env * self.gate_release
            };
            let level = gain_to_db(self.gate_env);
            if level > self.gate_db {
                self.gate_hold = (self.rate * GATE_HOLD_MS * 0.001) as usize;
                if !self.gate_open {
                    self.gate_open = true;
                    self.push(Kind::Onset);
                }
            } else if self.gate_open && level < self.gate_db - GATE_HYSTERESIS_DB {
                if self.gate_hold == 0 {
                    self.gate_open = false;
                    self.push(Kind::Release);
                } else {
                    self.gate_hold -= 1;
                }
            }
        }
        self.decimate_count += 1;
        if self.decimate_count >= self.decimate {
            self.decimate_count = 0;
            self.ring[self.ring_at] = d;
            self.ring_at = (self.ring_at + 1) % self.ring.len();
            self.hop_count += 1;
            if self.hop_count >= self.hop {
                self.hop_count = 0;
                self.hop_analysis(&mut emit);
            }
        }
        while self.event_len > 0 {
            let e = self.events[self.event_head];
            if e.time + self.latency as u64 > self.clock {
                break;
            }
            self.event_head = (self.event_head + 1) % EVENTS;
            self.event_len -= 1;
            self.handle(e.kind, &mut emit);
        }
        self.clock += 1;
    }
}

/// Fundamental that follows tracker events: jumps on new notes, slides on bends and legato.
#[derive(Clone, Copy, Debug, Default)]
pub struct PitchFollower {
    target: f32,
    hz: f32,
    jump: bool,
}
impl PitchFollower {
    pub fn event(&mut self, snap: bool, e: NoteEvent) {
        match e {
            NoteEvent::Start { pitch, glide } => {
                let p = if snap { pitch.round() } else { pitch };
                self.target = crate::midi_to_hz(p);
                self.jump |= !glide;
            }
            NoteEvent::Move { pitch } if !snap => self.target = crate::midi_to_hz(pitch),
            // The last pitch is held so tails keep their tuning.
            _ => {}
        }
    }
    /// Advances one sample with one-pole coefficient `coef`. Returns the frequency, 0 before the
    /// first note, and whether it jumped this sample (a new note, or the first one).
    pub fn step(&mut self, coef: f32) -> (f32, bool) {
        let jump = std::mem::take(&mut self.jump) || (self.hz == 0. && self.target > 0.);
        if jump {
            self.hz = self.target;
        } else {
            self.hz += (self.target - self.hz) * coef;
        }
        (self.hz, jump)
    }
    pub fn hz(&self) -> f32 {
        self.hz
    }
}

/// A processor driven by note events that are aligned with its input audio. Lathe, Temper and
/// Knurl are built from one each, so the same code runs standalone (after its own tracker) and
/// inside a strip that shares one tracker between several of them.
pub trait Tracked {
    /// Delay of the output behind the input, in samples.
    fn latency(&self) -> usize;
    /// The undelayed input, one tracker lookahead ahead of the audio passed to `process`.
    /// Called before the sample's events.
    fn ahead(&mut self, _mono: f32) {}
    /// A note event due at the next `process` sample.
    fn event(&mut self, e: NoteEvent);
    /// One sample of audio; `sounding` is the tracker's note state at this sample.
    fn process(&mut self, x: [f32; 2], sounding: bool) -> [f32; 2];
}

/// A tracker and the lookahead delay in front of one [`Tracked`] module: the standalone form.
pub struct Lookahead {
    pub tracker: Tracker,
    delay: Vec<[f32; 2]>,
    at: usize,
}
impl Lookahead {
    pub fn new(rate: f32, cfg: TrackerConfig) -> Self {
        let tracker = Tracker::new(rate, cfg);
        let len = tracker.latency() + 1;
        Self {
            tracker,
            delay: vec![[0.; 2]; len],
            at: 0,
        }
    }
    /// The tracker's lookahead, in samples.
    pub fn latency(&self) -> usize {
        self.tracker.latency()
    }
    pub fn reset(&mut self) {
        self.tracker.reset();
        self.delay.fill([0.; 2]);
    }
    /// Feeds one input sample (`mono` is its channel average) and returns `module`'s output.
    #[inline]
    pub fn process(&mut self, x: [f32; 2], mono: f32, module: &mut impl Tracked) -> [f32; 2] {
        module.ahead(mono);
        self.tracker.process(mono, |e| module.event(e));
        let len = self.delay.len();
        self.delay[self.at] = x;
        self.at = (self.at + 1) % len;
        let y = self.delay[self.at];
        module.process(y, self.tracker.sounding())
    }
}