//! Deep redundancy (DRED) for Opus, decoded exactly — `wai.audio.opus.dred.int`
//! ([extensions/opus-dred.md]).
//!
//! draft-ietf-mlcodec-opus-dred-07 carries, in an Opus packet's padding, a compact neural latent
//! of the speech before the packet, so that a receiver can rebuild what a loss took. Its decoder
//! network is specified in floating point, so two receivers rebuild different features from the
//! same packets. This module is the integer profile `dred-int/1`: the same network, from the
//! draft's own weights file pinned by digest, as `wai-ops/1` graphs whose output is the same
//! integers on every machine and on every fabric whose operator implementation conforms.
//!
//! What this module implements, by section of the extension:
//!
//! - §3: finding DRED in a packet ([`packet`]), the payload's integers ([`payload`], over the
//!   range decoder of RFC 6716 §4.1, [`ec`], and the draft's tables, [`tables`]);
//! - §4: the parameter set ([`weights`]), the lowering `dred-int/1` and its tables ([`lower`]),
//!   and the driver ([`Model`]);
//! - §5: a stream under loss, its reconstruction and the record of it ([`scenario`]);
//! - §6: the measurement against the floating-point network ([`float_ref`], [`measure`]);
//! - §7: the presentation statement over a reconstruction ([`scenario::player_entry`]).
//!
//! The corpus `live-conformance/dred/` is written by [`corpus::generate`] and checked by
//! [`corpus::verify`], and by `live-conformance/dred_verify.py`, which shares no code with this
//! module.
//!
//! [extensions/opus-dred.md]: ../../extensions/opus-dred.md

pub mod corpus;
pub mod ec;
pub mod float_ref;
pub mod lower;
pub mod measure;
pub mod packet;
pub mod payload;
pub mod scenario;
pub mod tables;
pub mod weights;

use std::collections::BTreeMap;

use sha2::{Digest, Sha256};
use wai_ops::interp::Tensor;
use wai_ops::Loaded;

pub use payload::Payload;

/// The capability string of the draft's decoding, within the draft's own conformance
/// (§4.1 of the draft), and of the integer profile.
pub const CAPABILITY: &str = "wai.audio.opus.dred";
pub const CAPABILITY_INT: &str = "wai.audio.opus.dred.int";
/// Features per 10 ms frame, and frames per latent vector (draft §2.1, §2.2.2).
pub const FEATURES: usize = 20;
pub const FRAMES_PER_LATENT: usize = 4;

/// The refusals of extensions/opus-dred.md §8, and two of the integer operator set's.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum Reason {
    /// The packet breaks RFC 6716 §3.4: nothing in it is read.
    PacketMalformed,
    /// A quantized index's magnitude passes [`payload::MAGNITUDE_CAP`].
    DredLatentRange,
    /// The weights are not the pinned file, or a record names another digest.
    WeightsDigestMismatch,
    /// A record's lost packets are not those its loss pattern gives.
    LossPatternMismatch,
    /// A record's reconstruction is not the one its packets, loss and weights give.
    ReconstructionMismatch,
    /// The set does not lower or load (integer-operators §10.5): `model_mismatch`. Unreachable
    /// through the pinned digest, which names a set that lowers.
    WeightsMalformed,
}

impl Reason {
    pub fn code(self) -> &'static str {
        match self {
            Reason::PacketMalformed => "packet_malformed",
            Reason::DredLatentRange => "dred_latent_range",
            Reason::WeightsDigestMismatch => "weights_digest_mismatch",
            Reason::LossPatternMismatch => "loss_pattern_mismatch",
            Reason::ReconstructionMismatch => "reconstruction_mismatch",
            Reason::WeightsMalformed => "model_mismatch",
        }
    }
    /// The codes §8 lists, in its order.
    pub const LISTED: [Reason; 5] =
        [Reason::PacketMalformed, Reason::DredLatentRange, Reason::WeightsDigestMismatch, Reason::LossPatternMismatch, Reason::ReconstructionMismatch];
}

/// One 10 ms feature vector at `2^16`.
pub type Row = [i32; FEATURES];

/// The loaded graphs of `dred-int/1`, and their identities (the SHA-256 of each graph's
/// `WHG1` bytes, integer-operators §10.5).
pub struct Model {
    pub init: Loaded,
    pub step: Loaded,
    pub init_sha256: [u8; 32],
    pub step_sha256: [u8; 32],
}

impl Model {
    /// The pinned weights file: its digest checked, then lowered and loaded under `acc64`.
    pub fn from_weights(bytes: &[u8]) -> Result<Model, Reason> {
        Model::from_layers(&weights::pinned(bytes)?)
    }

    /// Any decoder layers (the tests' synthetic sets): lowered and loaded under `acc64`.
    pub fn from_layers(layers: &BTreeMap<&'static str, weights::Layer>) -> Result<Model, Reason> {
        let g = lower::lower(layers)?;
        let (ib, sb) = (g.init.to_bytes(), g.step.to_bytes());
        let load = |b: &[u8]| wai_ops::graph::load(b, wai_ops::Profile::Acc64).map_err(|_| Reason::WeightsMalformed);
        Ok(Model { init: load(&ib)?, step: load(&sb)?, init_sha256: Sha256::digest(&ib).into(), step_sha256: Sha256::digest(&sb).into() })
    }

    /// The features of the first `n` latent vectors of `p` (all of them when `n` passes their
    /// count), newest first: row `4l + j` is vector `j` of latent `l`'s output, so row 0 is the
    /// newest 10 ms and each row is 10 ms older than the one before (§4.7).
    pub fn decode(&self, p: &Payload, n: usize) -> Result<Vec<Row>, Reason> {
        let run = |g: &Loaded, ins: Vec<Tensor>| crate::int_ops::run_graph(g, ins, &[]).map_err(|_| Reason::WeightsMalformed);
        let col = |c: u32, v: Vec<i64>| Tensor::new(c, 1, 1, v);
        let mut state = run(&self.init, vec![col(50, payload::state_input(p))])?;
        let mut mem: Vec<Tensor> = (0..5).map(|_| Tensor::filled(32, 1, 1, 0)).collect();
        let mut rows = Vec::new();
        for l in 0..n.min(p.latents.len()) {
            let mut ins = vec![col(26, payload::latent_input(p, l))];
            ins.append(&mut state);
            ins.append(&mut mem);
            let mut out = run(&self.step, ins)?;
            let x = out.remove(0);
            for j in 0..FRAMES_PER_LATENT {
                let mut r = [0i32; FEATURES];
                for (i, v) in r.iter_mut().enumerate() {
                    *v = x.data[j * FEATURES + i] as i32;
                }
                rows.push(r);
            }
            mem = out.split_off(5);
            state = out;
        }
        Ok(rows)
    }
}

/// The bytes of feature rows: each value as `i32` little-endian, row by row.
pub fn row_bytes(rows: &[Row]) -> Vec<u8> {
    rows.iter().flat_map(|r| r.iter().flat_map(|v| v.to_le_bytes())).collect()
}

/// SHA-256, lowercase hexadecimal.
pub fn hex32(d: &[u8; 32]) -> String {
    d.iter().map(|b| format!("{b:02x}")).collect()
}

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

    /// The whole driver on a synthetic set: the graphs lower and load under `acc64`, a payload
    /// decodes to 4 rows a latent, the first `n` latents' rows do not depend on the latents after
    /// them, and two decodes are the same integers.
    #[test]
    fn the_driver_decodes_a_synthetic_set() {
        let file = weights::tests::synthetic_file(11, 3.0);
        let layers = weights::decoder_layers(&weights::parse(&file).unwrap()).unwrap();
        let m = Model::from_layers(&layers).unwrap();
        let h = payload::Header { q0: 3, dq: 2, extended: None, indicator: 8, qmax: 15 };
        let mut lat = Vec::new();
        for c in 0..5 {
            let q = h.quantizer(c);
            let mut l = [0i32; payload::LATENT_DIM];
            for (k, v) in l.iter_mut().enumerate() {
                if payload::coded(tables::LATENT_DECAY[k][q as usize], tables::LATENT_P0[k][q as usize]) {
                    *v = ((k * 7 + c * 3) % 9) as i32 - 4;
                }
            }
            lat.push((l, q));
        }
        let mut st = [0i32; payload::STATE_DIM];
        for (k, v) in st.iter_mut().enumerate() {
            if payload::coded(tables::STATE_DECAY[k][3], tables::STATE_P0[k][3]) {
                *v = (k % 5) as i32 - 2;
            }
        }
        let p = Payload { header: h, state: st, latents: lat };
        let all = m.decode(&p, usize::MAX).unwrap();
        assert_eq!(all.len(), 20);
        assert!(all.iter().flatten().any(|&v| v != 0));
        let two = m.decode(&p, 2).unwrap();
        assert_eq!(&all[..8], &two[..]);
        assert_eq!(m.decode(&p, usize::MAX).unwrap(), all);
        // The graph identities are fixed by the set.
        let again = Model::from_layers(&layers).unwrap();
        assert_eq!((again.init_sha256, again.step_sha256), (m.init_sha256, m.step_sha256));
    }
}
