//! The DRED extension payload (draft-ietf-mlcodec-opus-dred-07 §3, §3.1), parsed to its
//! integers: the header fields, the quantized initial state and the quantized latent vectors,
//! each latent with the quantizer it was coded at. Then the profile's dequantization
//! (extensions/opus-dred.md §4.3), which turns each quantized index into the graph input.
//!
//! Every value here is an exact integer function of the payload's bytes. The parse follows the
//! draft's reference decoder where the draft's prose leaves a detail open, and the extension
//! says where (§4.2): a magnitude is 1 plus the symbols decoded, and the quantizer of chunk `k`
//! is `min(qmax, q0 + ⌊(16·dQ_k + 8) / 16⌋)` with the slope in sixteenths.

use sha2::{Digest, Sha256};

use super::ec::{Dec, Enc};
use super::tables;
use super::Reason;

/// Coefficients in the initial state and in one latent vector (§2.2.2).
pub const STATE_DIM: usize = 50;
pub const LATENT_DIM: usize = 25;
/// The slope table of §3, in sixteenths of a quantizer step per chunk.
pub const DQ_SIXTEENTHS: [u32; 8] = [0, 2, 3, 4, 6, 8, 12, 16];
/// The largest magnitude a quantized index may reach (extensions/opus-dred.md §4.2): the
/// decode of a coefficient whose magnitude passes it is refused (`dred_latent_range`). A
/// realistic stream stays far below it; without a bound, a crafted stream could make the
/// escape loop of §3.1 read symbols without end past the payload's last byte.
pub const MAGNITUDE_CAP: i32 = 1 << 16;

/// The header of §3.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Header {
    pub q0: u32,
    pub dq: u32,
    /// The 8-bit extended offset when `X = 1`.
    pub extended: Option<u32>,
    /// The 5-bit offset indicator.
    pub indicator: u32,
    pub qmax: u32,
}

impl Header {
    /// `encoded_offset = extended·32 + indicator`, in units of 2.5 ms.
    pub fn encoded_offset(&self) -> i64 {
        self.extended.unwrap_or(0) as i64 * 32 + self.indicator as i64
    }
    /// The offset of §3, relative to the start of the frame that carries the extension, in
    /// units of 2.5 ms: `16 − encoded_offset + frame_offset`. The newest feature vector ends
    /// this far after that frame's start (a negative value is before it).
    pub fn dred_offset(&self, frame_offset: i64) -> i64 {
        16 - self.encoded_offset() + frame_offset
    }
    /// The quantizer of chunk `k` (§3).
    pub fn quantizer(&self, k: usize) -> u32 {
        let step = (DQ_SIXTEENTHS[self.dq as usize] as u64 * k as u64 + 8) / 16;
        (self.q0 as u64 + step).min(self.qmax as u64) as u32
    }
}

/// A parsed payload: its header, the quantized initial state, and each latent vector with its
/// quantizer, newest first.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Payload {
    pub header: Header,
    pub state: [i32; STATE_DIM],
    pub latents: Vec<([i32; LATENT_DIM], u32)>,
}

impl Payload {
    /// The bytes the symbol digest of the corpus covers: `q0`, `dQ`, the offset fields and
    /// `qmax` as bytes, the state's and each latent's indices as `i32` little-endian, each
    /// latent preceded by its quantizer as one byte (extensions/opus-dred.md §6.2).
    pub fn symbol_bytes(&self) -> Vec<u8> {
        let h = &self.header;
        let mut b = vec![h.q0 as u8, h.dq as u8, h.extended.is_some() as u8, h.extended.unwrap_or(0) as u8, h.indicator as u8, h.qmax as u8];
        b.extend_from_slice(&(self.latents.len() as u32).to_le_bytes());
        for v in self.state {
            b.extend_from_slice(&v.to_le_bytes());
        }
        for (l, q) in &self.latents {
            b.push(*q as u8);
            for v in l {
                b.extend_from_slice(&v.to_le_bytes());
            }
        }
        b
    }
}

fn icdf_of(decay: u32) -> [u16; 8] {
    let mut icdf = [0u16; 8];
    icdf[0] = decay.max(7) as u16;
    for i in 1..7 {
        icdf[i] = (7 - i as u32).max((icdf[i - 1] as u32 * decay) >> 15) as u16;
    }
    icdf
}

/// One coefficient (§3.1): nothing when `r = 0` or `p0 = 255`; otherwise the zero / positive /
/// negative symbol, then the magnitude in runs of the 8-symbol table.
fn laplace(d: &mut Dec, p0: u8, r: u8) -> Result<i32, Reason> {
    if r == 0 || p0 == 255 {
        return Ok(0);
    }
    let p0 = (p0 as u32) << 7;
    let s0 = 32768 - p0;
    let sign = d.icdf16(&[s0 as u16, (s0 / 2) as u16, 0], 15);
    if sign == 0 {
        return Ok(0);
    }
    let icdf = icdf_of((r as u32) << 7);
    let mut value: i32 = 1;
    loop {
        let v = d.icdf16(&icdf, 15) as i32;
        value += v;
        if value > MAGNITUDE_CAP {
            return Err(Reason::DredLatentRange);
        }
        if v != 7 {
            break;
        }
    }
    Ok(if sign == 2 { -value } else { value })
}

fn vector<const N: usize>(d: &mut Dec, q: u32, r: &[[u8; 16]; N], p0: &[[u8; 16]; N]) -> Result<[i32; N], Reason> {
    let mut out = [0i32; N];
    for (k, o) in out.iter_mut().enumerate() {
        *o = laplace(d, p0[k][q as usize], r[k][q as usize])?;
    }
    Ok(out)
}

/// Parse a DRED payload (the extension's data after any experimental prefix), decoding at most
/// `max_latents` latent vectors when given. Latent vectors are decoded until fewer than 8 bits
/// of the payload remain (§3), then no more.
pub fn parse(bytes: &[u8], max_latents: Option<usize>) -> Result<Payload, Reason> {
    let mut d = Dec::new(bytes);
    let q0 = d.uint(16);
    let dq = d.uint(8);
    let extended = if d.uint(2) == 1 { Some(d.uint(256)) } else { None };
    let indicator = d.uint(32);
    let mut qmax = 15;
    if q0 < 14 && dq != 0 {
        let nvals = 15 - (q0 + 1);
        let ft = 2 * nvals;
        let s = d.decode(ft);
        if s >= nvals {
            qmax = q0 + (s - nvals) + 1;
            d.update(s, s + 1, ft);
        } else {
            d.update(0, nvals, ft);
        }
    }
    let header = Header { q0, dq, extended, indicator, qmax };
    let state = vector(&mut d, q0, &tables::STATE_DECAY, &tables::STATE_P0)?;
    let mut latents = Vec::new();
    let total = 8 * bytes.len() as i64;
    while max_latents.is_none_or(|m| latents.len() < m) {
        if total - d.tell() <= 7 {
            break;
        }
        let q = header.quantizer(latents.len());
        let l = vector(&mut d, q, &tables::LATENT_DECAY, &tables::LATENT_P0)?;
        latents.push((l, q));
    }
    Ok(Payload { header, state, latents })
}

/// The body of a payload with no header (the draft's test bit files): the state at `q0`, then
/// `n` latents, each at `q0`.
pub fn parse_body(bytes: &[u8], q0: u32, n: usize) -> Result<Payload, Reason> {
    let mut d = Dec::new(bytes);
    let header = Header { q0, dq: 0, extended: None, indicator: 0, qmax: 15 };
    let state = vector(&mut d, q0, &tables::STATE_DECAY, &tables::STATE_P0)?;
    let mut latents = Vec::with_capacity(n);
    for _ in 0..n {
        latents.push((vector(&mut d, q0, &tables::LATENT_DECAY, &tables::LATENT_P0)?, q0));
    }
    Ok(Payload { header, state, latents })
}

/// Indices drawn from the draft's own distributions by decoding `bytes` (for the corpus
/// generator, which feeds it a seeded stream): a state vector at `q`.
pub fn sample_state(bytes: &[u8], q: u32) -> [i32; STATE_DIM] {
    vector(&mut Dec::new(bytes), q, &tables::STATE_DECAY, &tables::STATE_P0).unwrap_or([0; STATE_DIM])
}

/// [`sample_state`] for a latent vector.
pub fn sample_latent(bytes: &[u8], q: u32) -> [i32; LATENT_DIM] {
    vector(&mut Dec::new(bytes), q, &tables::LATENT_DECAY, &tables::LATENT_P0).unwrap_or([0; LATENT_DIM])
}

/// A latent's indices at `q0` requantized to `q`: `x·s_q / s_q0` rounded half up where `q`
/// codes the coefficient, 0 where it does not (an encoder's choice; the encoder is not
/// normative, draft §2.2.1).
pub fn requantize(x: &[i32; LATENT_DIM], q0: u32, q: u32) -> [i32; LATENT_DIM] {
    let mut out = [0i32; LATENT_DIM];
    for k in 0..LATENT_DIM {
        if coded(tables::LATENT_DECAY[k][q as usize], tables::LATENT_P0[k][q as usize]) {
            let (s0, s) = (tables::LATENT_SCALE[k][q0 as usize] as i64, tables::LATENT_SCALE[k][q as usize] as i64);
            out[k] = div_half_up(x[k] as i64 * s, s0.max(1)) as i32;
        }
    }
    out
}

/// `⌊(2a + b) / (2b)⌋` for `b > 0`: `a / b` rounded to nearest, ties toward +∞.
pub fn div_half_up(a: i64, b: i64) -> i64 {
    (2 * a + b).div_euclid(2 * b)
}

/// The profile's dequantization (extensions/opus-dred.md §4.3): the index `x` with scale `s`
/// stands for `x·256 / max(1, s)`, which the graph input holds in Q16, rounded half up.
pub fn dequantize(x: i32, s: u8) -> i64 {
    div_half_up((x as i64) << 24, (s as i64).max(1))
}

/// The state's graph input (Q16).
pub fn state_input(p: &Payload) -> Vec<i64> {
    let q = p.header.q0 as usize;
    (0..STATE_DIM).map(|k| dequantize(p.state[k], tables::STATE_SCALE[k][q])).collect()
}

/// Latent `l`'s graph input (Q16): its 25 dequantized coefficients, then `q/8 − 1` (§2.2.2).
pub fn latent_input(p: &Payload, l: usize) -> Vec<i64> {
    let (v, q) = &p.latents[l];
    let mut z: Vec<i64> = (0..LATENT_DIM).map(|k| dequantize(v[k], tables::LATENT_SCALE[k][*q as usize])).collect();
    z.push((*q as i64 - 8) << 13);
    z
}

/// SHA-256 of the six decoder tables, each serialized row by row (coefficient `k` from 0, the
/// sixteen quantizers in order, one byte each), in the order latent scale, latent decay, latent
/// P(0), state scale, state decay, state P(0) (extensions/opus-dred.md §4.2).
pub fn tables_digest() -> [u8; 32] {
    let mut h = Sha256::new();
    for t in [&tables::LATENT_SCALE[..], &tables::LATENT_DECAY[..], &tables::LATENT_P0[..], &tables::STATE_SCALE[..], &tables::STATE_DECAY[..], &tables::STATE_P0[..]] {
        for row in t {
            h.update(row);
        }
    }
    h.finalize().into()
}

// ---- the encoder side, for the corpus generator ----

fn laplace_encode(e: &mut Enc, value: i32, p0: u8, r: u8) {
    if r == 0 || p0 == 255 {
        return;
    }
    let p0 = (p0 as u32) << 7;
    let s0 = 32768 - p0;
    let sign = match value.signum() {
        0 => 0,
        1 => 1,
        _ => 2,
    };
    e.icdf16(sign, &[s0 as u16, (s0 / 2) as u16, 0], 15);
    if value != 0 {
        let icdf = icdf_of((r as u32) << 7);
        let mut m = value.abs() - 1;
        loop {
            e.icdf16(m.min(7) as usize, &icdf, 15);
            m -= 7;
            if m < 0 {
                break;
            }
        }
    }
}

/// Write a payload: the inverse of [`parse`]. A coefficient whose decay is 0 or whose P(0) is
/// 255 is not coded, so it must be 0 (the decoder gives 0 there).
pub fn write(p: &Payload) -> Vec<u8> {
    let h = &p.header;
    let mut e = Enc::new();
    e.uint(h.q0, 16);
    e.uint(h.dq, 8);
    match h.extended {
        Some(x) => {
            e.uint(1, 2);
            e.uint(x, 256);
        }
        None => e.uint(0, 2),
    }
    e.uint(h.indicator, 32);
    if h.q0 < 14 && h.dq != 0 {
        let nvals = 15 - (h.q0 + 1);
        if h.qmax >= 15 {
            e.encode(0, nvals, 2 * nvals);
        } else {
            let fl = nvals + h.qmax - (h.q0 + 1);
            e.encode(fl, fl + 1, 2 * nvals);
        }
    }
    let q0 = h.q0 as usize;
    for (k, &v) in p.state.iter().enumerate() {
        laplace_encode(&mut e, v, tables::STATE_P0[k][q0], tables::STATE_DECAY[k][q0]);
    }
    for (l, q) in &p.latents {
        for (k, &v) in l.iter().enumerate() {
            laplace_encode(&mut e, v, tables::LATENT_P0[k][*q as usize], tables::LATENT_DECAY[k][*q as usize]);
        }
    }
    e.finish()
}

/// Whether coefficient `k` of a vector at quantizer `q` is coded (§3.1).
pub fn coded(r: u8, p0: u8) -> bool {
    r != 0 && p0 != 255
}

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

    #[test]
    fn quantizer_follows_the_slope_table() {
        // §3's example: Q0 = 5, dQ = 2 (3/16), chunk 20 uses round(5 + 3·20/16) = 9.
        let h = Header { q0: 5, dq: 2, extended: None, indicator: 0, qmax: 15 };
        assert_eq!(h.quantizer(20), 9);
        let h = Header { q0: 5, dq: 7, extended: None, indicator: 0, qmax: 11 };
        assert_eq!((h.quantizer(0), h.quantizer(3), h.quantizer(10)), (5, 8, 11));
    }

    #[test]
    fn dequantization_rounds_half_up() {
        assert_eq!(dequantize(1, 255), div_half_up(1 << 24, 255));
        assert_eq!(dequantize(-3, 0), -3 << 24);
        assert_eq!(div_half_up(-5, 2), -2);
        assert_eq!(div_half_up(5, 2), 3);
        assert_eq!(div_half_up(-7, 2), -3);
    }

    #[test]
    fn payloads_round_trip() {
        let mut seed = 99u64;
        let mut rnd = |m: i32| {
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((seed >> 33) as i32).rem_euclid(2 * m + 1) - m
        };
        for (q0, dq, qmax, ext) in [(2u32, 0u32, 15u32, None), (6, 3, 12, Some(3u32)), (13, 7, 15, None), (14, 5, 15, Some(255))] {
            let header = Header { q0, dq, extended: ext, indicator: 17, qmax };
            let mut state = [0i32; STATE_DIM];
            for (k, s) in state.iter_mut().enumerate() {
                if coded(tables::STATE_DECAY[k][q0 as usize], tables::STATE_P0[k][q0 as usize]) {
                    *s = rnd(20);
                }
            }
            let mut latents = Vec::new();
            for c in 0..9 {
                let q = header.quantizer(c);
                let mut l = [0i32; LATENT_DIM];
                for (k, v) in l.iter_mut().enumerate() {
                    if coded(tables::LATENT_DECAY[k][q as usize], tables::LATENT_P0[k][q as usize]) {
                        *v = rnd(30);
                    }
                }
                latents.push((l, q));
            }
            let p = Payload { header, state, latents };
            let bytes = write(&p);
            assert_eq!(parse(&bytes, None).unwrap(), p, "q0 {q0} dq {dq}");
            assert_eq!(parse(&bytes, Some(4)).unwrap().latents[..], p.latents[..4]);
        }
    }

    /// An empty payload reads zeros: a header and a state, and no latent vector, since fewer
    /// than 8 bits remain from the start.
    #[test]
    fn an_empty_payload_has_no_latents() {
        assert!(parse(&[], None).unwrap().latents.is_empty());
    }
}
