//! The measurement of extensions/opus-dred.md §6: how far the integer profile's features lie
//! from a floating-point decode of the same latents.
//!
//! - [`compare`] is the feature comparison of draft-ietf-mlcodec-opus-dred-07 §4.1, as the
//!   draft's comparison tool computes it: per feature, the mean squared difference over the
//!   frames; `total` the square root of their sum and `max` of their largest, the pitch feature
//!   (18) left out of both; `pitch` the mean absolute pitch difference over the frames whose
//!   reference voicing (feature 19) passes 0.2. The draft's DRED decoding test passes when
//!   `total ≤ 0.5`, `max ≤ 0.15` and `pitch ≤ 0.02`.
//! - [`Stats`] adds the largest absolute difference, and a band log-spectral distance: the 18
//!   cepstral features are the orthonormal DCT-II of the log10 band energies (draft §2.1.3), so
//!   the inverse DCT of their difference is the difference of the log band energies, and ten
//!   times its root mean square over the bands is a distance in decibels per frame.
//!
//! It also reads the draft's test vectors ([DRED-Test-Vectors], draft §7.1): the DRED bit files
//! (records of a 4-byte big-endian quantizer, chunk count and byte count, then the payload),
//! the packet files (records of a 4-byte big-endian length and final range, then the packet,
//! length 0 for a lost one) and the float32 feature files. Nothing of them is committed.

use super::payload::{self, Payload};
use super::FEATURES;

/// The draft's three figures (§4.1), for a reference `x` and a decode `y`.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Compare {
    pub total: f64,
    pub max: f64,
    pub pitch: f64,
}

impl Compare {
    /// The draft's DRED decoding thresholds.
    pub fn passes(&self) -> bool {
        self.total <= 0.5 && self.max <= 0.15 && self.pitch <= 0.02
    }
}

pub fn compare(x: &[[f64; FEATURES]], y: &[[f64; FEATURES]]) -> Compare {
    let mut mse = [0f64; FEATURES];
    let (mut pitch, mut pitch_n) = (0f64, 0usize);
    for (a, b) in x.iter().zip(y) {
        for i in 0..FEATURES {
            let e = a[i] - b[i];
            mse[i] += e * e;
        }
        if a[FEATURES - 1] > 0.2 {
            pitch += (a[FEATURES - 2] - b[FEATURES - 2]).abs();
            pitch_n += 1;
        }
    }
    let n = x.len().max(1) as f64;
    let (mut tot, mut max) = (0f64, 0f64);
    for (i, m) in mse.iter().enumerate() {
        if i != FEATURES - 2 {
            tot += m / n;
            max = max.max(m / n);
        }
    }
    Compare { total: tot.sqrt(), max: max.sqrt(), pitch: if pitch_n == 0 { 0.0 } else { pitch / pitch_n as f64 } }
}

/// Every figure of §6 for one comparison.
#[derive(Clone, Copy, Debug, Default, PartialEq)]
pub struct Stats {
    pub frames: usize,
    pub draft: Compare,
    /// The largest absolute difference of any feature but pitch, and of pitch.
    pub max_abs: f64,
    pub max_abs_pitch: f64,
    /// The band log-spectral distance in dB: its mean over the frames, and its largest.
    pub lsd_mean_db: f64,
    pub lsd_max_db: f64,
}

pub fn stats(x: &[[f64; FEATURES]], y: &[[f64; FEATURES]]) -> Stats {
    let nb = 18usize;
    let mut s = Stats { frames: x.len(), draft: compare(x, y), ..Stats::default() };
    let mut lsd_sum = 0.0;
    for (a, b) in x.iter().zip(y) {
        for i in 0..FEATURES {
            let d = (a[i] - b[i]).abs();
            if i == FEATURES - 2 {
                s.max_abs_pitch = s.max_abs_pitch.max(d);
            } else {
                s.max_abs = s.max_abs.max(d);
            }
        }
        // Inverse orthonormal DCT-II of the cepstral difference.
        let mut ss = 0.0;
        for n in 0..nb {
            let mut v = 0.0;
            for k in 0..nb {
                let w = if k == 0 { (1.0 / nb as f64).sqrt() } else { (2.0 / nb as f64).sqrt() };
                v += w * (a[k] - b[k]) * cos_pi((n as f64 + 0.5) * k as f64 / nb as f64);
            }
            ss += (10.0 * v) * (10.0 * v);
        }
        let lsd = (ss / nb as f64).sqrt();
        lsd_sum += lsd;
        s.lsd_max_db = s.lsd_max_db.max(lsd);
    }
    s.lsd_mean_db = lsd_sum / x.len().max(1) as f64;
    s
}

/// `cos(π·x)` for `x ≥ 0` from basic operations only (reduced to `[0, 1/4]` by the
/// function's symmetries, then a Taylor series), so that every machine prints the same figures.
fn cos_pi(x: f64) -> f64 {
    let x = x % 2.0;
    let (x, s) = if x > 1.0 { (2.0 - x, 1.0) } else { (x, 1.0) };
    let (x, s) = if x > 0.5 { (1.0 - x, -s) } else { (x, s) };
    let t = std::f64::consts::PI * x;
    let (mut term, mut sum) = (1.0, 1.0);
    for k in 1..=12 {
        term = -term * t * t / ((2 * k - 1) as f64 * (2 * k) as f64);
        sum += term;
    }
    s * sum
}

/// Integer rows at `2^16` as `f64`.
pub fn rows_f64(rows: &[super::Row]) -> Vec<[f64; FEATURES]> {
    rows.iter()
        .map(|r| {
            let mut o = [0f64; FEATURES];
            for (d, v) in o.iter_mut().zip(r) {
                *d = *v as f64 / 65536.0;
            }
            o
        })
        .collect()
}

/// One record of a DRED bit file: the quantizer, the chunk count and the payload. The draft's
/// test decoder decodes its initial state at the quantizer, then `chunks / 2` latents at the
/// same quantizer, and writes `2·chunks` feature vectors oldest first.
pub struct BitRecord {
    pub q0: u32,
    pub chunks: usize,
    pub payload: Vec<u8>,
}

pub fn read_dred_bits(b: &[u8]) -> Option<Vec<BitRecord>> {
    let mut out = Vec::new();
    let mut at = 0usize;
    let be = |o: usize| u32::from_be_bytes(b[o..o + 4].try_into().unwrap());
    while at + 12 <= b.len() {
        let (q0, chunks, len) = (be(at), be(at + 4) as usize, be(at + 8) as usize);
        at += 12;
        if b.len() - at < len || q0 > 15 {
            return None;
        }
        out.push(BitRecord { q0, chunks, payload: b[at..at + len].to_vec() });
        at += len;
    }
    (at == b.len()).then_some(out)
}

/// A bit-file record parsed as its decoder parses it: no header (the record gives `q0`, with
/// `dQ = 0`), the state, then `chunks / 2` latents, each at `q0`, read through the payload's own
/// coefficient rules.
pub fn record_payload(r: &BitRecord) -> Result<Payload, super::Reason> {
    payload::parse_body(&r.payload, r.q0, r.chunks / 2)
}

/// The packets of a packet file: `None` for a lost packet.
pub fn read_packets(b: &[u8]) -> Option<Vec<Option<Vec<u8>>>> {
    let mut out = Vec::new();
    let mut at = 0usize;
    while at + 8 <= b.len() {
        let len = u32::from_be_bytes(b[at..at + 4].try_into().unwrap()) as usize;
        at += 8;
        if b.len() - at < len {
            return None;
        }
        out.push((len > 0).then(|| b[at..at + len].to_vec()));
        at += len;
    }
    (at == b.len()).then_some(out)
}

/// A float32 feature file, 20 features a frame.
pub fn read_features(b: &[u8]) -> Option<Vec<[f64; FEATURES]>> {
    if !b.len().is_multiple_of(4 * FEATURES) {
        return None;
    }
    Some(
        b.as_chunks::<{ 4 * FEATURES }>()
            .0
            .iter()
            .map(|c| {
                let mut r = [0f64; FEATURES];
                for (i, v) in r.iter_mut().enumerate() {
                    *v = f32::from_le_bytes(c[4 * i..4 * i + 4].try_into().unwrap()) as f64;
                }
                r
            })
            .collect(),
    )
}

/// Reverse a record's rows (newest first) into the oldest-first order of the bit file's
/// feature output: latent `l`'s four rows are written at `2·chunks − 4 − 4l` onwards, oldest
/// first.
pub fn oldest_first<T: Copy>(rows: &[T]) -> Vec<T> {
    rows.iter().rev().copied().collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn the_comparison_is_the_drafts() {
        let x = vec![[0.0; FEATURES]; 4];
        let mut y = x.clone();
        y[0][0] = 0.4; // mse 0.04 on feature 0
        y[1][18] = 1.0; // pitch, left out of total and max
        let c = compare(&x, &y);
        assert!((c.total - 0.2).abs() < 1e-12 && (c.max - 0.2).abs() < 1e-12);
        assert_eq!(c.pitch, 0.0); // no frame is voiced
        let mut xv = x.clone();
        for r in &mut xv {
            r[19] = 0.5;
        }
        assert!((compare(&xv, &y).pitch - 0.25).abs() < 1e-12);
        // A constant cepstral difference in c0 only is a flat band difference: 10·sqrt(1/18)·δ.
        let mut z = x.clone();
        z[0][0] = 1.0;
        let s = stats(&x[..1], &z[..1]);
        assert!((s.lsd_mean_db - 10.0 / 18f64.sqrt()).abs() < 1e-12);
        for x in [0.0, 0.1, 0.25, 0.5, 0.9, 1.0, 1.3, 7.7, 16.95] {
            assert!((cos_pi(x) - (std::f64::consts::PI * x).cos()).abs() < 1e-14, "{x}");
        }
    }
}
