//! The decoder weights (draft-ietf-mlcodec-opus-dred-07 §2.2.2.1, Appendix A): the record
//! file parsed, the decoder's float arrays found by name, and each layer's matrix expanded from
//! its storage layout to a dense matrix of the float32 values' bit patterns. Nothing here does
//! float arithmetic: a weight is carried as its 32 bits until [`super::lower`] quantizes it.

use std::collections::BTreeMap;

use sha2::{Digest, Sha256};

use super::Reason;

/// The SHA-256 of the weights file the draft names ([DRED-Weights], §7.1). The profile's
/// parameter set is this file and no other (extensions/opus-dred.md §4.1).
pub const WEIGHTS_SHA256: &str = "394f0ecf19926cebd678019dc94a54c88ae078340c16325069707f286cdf28a6";

/// One record of Appendix A.1.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Array {
    /// Type 0: float32 values, as their bit patterns.
    F32(Vec<u32>),
    /// Type 1: int32.
    I32(Vec<i32>),
    /// Type 3: int8 (read, never used by the float decoder).
    I8(Vec<i8>),
}

/// Parse the record file (Appendix A.1): each record a 64-byte header (magic `DNNw`, version 0,
/// type, size, block size, a NUL-terminated name of at most 43 bytes) and its data padded to
/// the block size, a multiple of 64. A record that breaks a rule, a name twice, or bytes after
/// the last record make the file malformed (`weights_malformed`).
pub fn parse(bytes: &[u8]) -> Result<BTreeMap<String, Array>, Reason> {
    let bad = Err(Reason::WeightsMalformed);
    let mut out = BTreeMap::new();
    let mut at = 0usize;
    let le = |b: &[u8], o: usize| i32::from_le_bytes(b[o..o + 4].try_into().unwrap());
    while at < bytes.len() {
        if bytes.len() - at < 64 || &bytes[at..at + 4] != b"DNNw" {
            return bad;
        }
        let (version, kind, size, block) = (le(bytes, at + 4), le(bytes, at + 8), le(bytes, at + 12), le(bytes, at + 16));
        let name_field = &bytes[at + 20..at + 64];
        let Some(nul) = name_field.iter().position(|&c| c == 0) else { return bad };
        let Ok(name) = std::str::from_utf8(&name_field[..nul]) else { return bad };
        if version != 0 || size < 0 || block < size || block % 64 != 0 || block - size >= 64 || name.is_empty() || !name.is_ascii() {
            return bad;
        }
        let (size, block) = (size as usize, block as usize);
        if bytes.len() - at - 64 < block {
            return bad;
        }
        let data = &bytes[at + 64..at + 64 + size];
        if bytes[at + 64 + size..at + 64 + block].iter().any(|&c| c != 0) {
            return bad;
        }
        let a = match kind {
            0 if size.is_multiple_of(4) => Array::F32(data.as_chunks::<4>().0.iter().map(|c| u32::from_le_bytes(*c)).collect()),
            1 if size.is_multiple_of(4) => Array::I32(data.as_chunks::<4>().0.iter().map(|c| i32::from_le_bytes(*c)).collect()),
            3 => Array::I8(data.iter().map(|&c| c as i8).collect()),
            _ => return bad,
        };
        if out.insert(name.to_string(), a).is_some() {
            return bad;
        }
        at += 64 + block;
    }
    Ok(out)
}

/// A layer's affine map as the float decoder computes it: `W·x + b`, `W` dense `[out][in]`,
/// every value as its float32 bit pattern.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Layer {
    pub inputs: usize,
    pub outputs: usize,
    /// `w[o * inputs + i]`.
    pub w: Vec<u32>,
    pub b: Vec<u32>,
}

/// The decoder's layers (§2.2.2, Appendix A.2), each `(name, inputs, outputs, sparse)`: the
/// arrays `dec_<name>_weights_float` and `dec_<name>_bias`, and `dec_<name>_weights_idx` for a
/// block-sparse one.
pub const LAYERS: [(&str, usize, usize, bool); 28] = [
    ("dense1", 26, 96, false),
    ("hidden_init", 50, 128, false),
    ("gru_init", 128, 320, true),
    ("gru1_input", 96, 192, true),
    ("gru1_recurrent", 64, 192, false),
    ("gru2_input", 192, 192, true),
    ("gru2_recurrent", 64, 192, false),
    ("gru3_input", 288, 192, true),
    ("gru3_recurrent", 64, 192, false),
    ("gru4_input", 384, 192, true),
    ("gru4_recurrent", 64, 192, false),
    ("gru5_input", 480, 192, true),
    ("gru5_recurrent", 64, 192, false),
    ("glu1", 64, 64, false),
    ("glu2", 64, 64, false),
    ("glu3", 64, 64, false),
    ("glu4", 64, 64, false),
    ("glu5", 64, 64, false),
    ("conv_dense1", 160, 32, true),
    ("conv_dense2", 256, 32, true),
    ("conv_dense3", 352, 32, true),
    ("conv_dense4", 448, 32, true),
    ("conv_dense5", 544, 32, true),
    ("conv1", 64, 32, false),
    ("conv2", 64, 32, false),
    ("conv3", 64, 32, false),
    ("conv4", 64, 32, false),
    ("conv5", 64, 32, false),
];
/// The output layer, block-sparse.
pub const OUTPUT: (&str, usize, usize, bool) = ("output", 576, 80, true);

/// The decoder's layers by name, expanded to dense matrices (Appendix A.3). A dense matrix is
/// stored input-major (`weights[in·outputs + out]`). A block-sparse one is stored as blocks of
/// 8 outputs by 4 inputs, input-major inside the block (`w[in·8 + out]`), listed group by
/// group of 8 outputs, each group's blocks named by the index array (a count, then each block's
/// first input). A block that starts at an input that is not a multiple of 4 or passes the
/// layer's inputs, a group that names one input block twice, a missing array, an array of the
/// wrong length, or a value that is not finite makes the set malformed (`weights_malformed`).
pub fn decoder_layers(arrays: &BTreeMap<String, Array>) -> Result<BTreeMap<&'static str, Layer>, Reason> {
    let bad = || Reason::WeightsMalformed;
    let f32s = |name: String| match arrays.get(&name) {
        Some(Array::F32(v)) => Ok(v.clone()),
        _ => Err(bad()),
    };
    let mut out = BTreeMap::new();
    for &(name, inputs, outputs, sparse) in LAYERS.iter().chain(std::iter::once(&OUTPUT)) {
        let raw = f32s(format!("dec_{name}_weights_float"))?;
        let b = f32s(format!("dec_{name}_bias"))?;
        if b.len() != outputs || (sparse && (outputs % 8 != 0 || inputs % 4 != 0)) {
            return Err(bad());
        }
        let mut w = vec![0u32; inputs * outputs];
        if sparse {
            let Some(Array::I32(idx)) = arrays.get(&format!("dec_{name}_weights_idx")) else { return Err(bad()) };
            let (mut p, mut wp) = (0usize, 0usize);
            for og in (0..outputs).step_by(8) {
                let &n = idx.get(p).ok_or_else(bad)?;
                p += 1;
                if n < 0 {
                    return Err(bad());
                }
                let mut seen = vec![false; inputs / 4];
                for _ in 0..n {
                    let &pos = idx.get(p).ok_or_else(bad)?;
                    p += 1;
                    if pos < 0 || pos % 4 != 0 || pos as usize >= inputs || std::mem::replace(&mut seen[pos as usize / 4], true) {
                        return Err(bad());
                    }
                    let blk = raw.get(wp..wp + 32).ok_or_else(bad)?;
                    for i in 0..4 {
                        for o in 0..8 {
                            w[(og + o) * inputs + pos as usize + i] = blk[i * 8 + o];
                        }
                    }
                    wp += 32;
                }
            }
            if p != idx.len() || wp != raw.len() {
                return Err(bad());
            }
        } else {
            if raw.len() != inputs * outputs {
                return Err(bad());
            }
            for i in 0..inputs {
                for o in 0..outputs {
                    w[o * inputs + i] = raw[i * outputs + o];
                }
            }
        }
        if w.iter().chain(&b).any(|&v| (v >> 23) & 0xff == 0xff) {
            return Err(bad());
        }
        out.insert(name, Layer { inputs, outputs, w, b });
    }
    Ok(out)
}

/// A layer's 8-bit weights and per-output scales, as the draft's reference implementation's
/// quantized arithmetic reads them (Appendix A.3.2, A.3.3): `w8[o·inputs + i]`, every 8×4 block
/// stored output-major inside (`w[out·4 + in]`); a dense one's blocks input-block by input-block
/// within each group of 8 outputs. Only for the measurement of §6, which emulates that arithmetic
/// to show where the draft's test vectors come from; no decode reads them.
pub fn int8_layer(arrays: &BTreeMap<String, Array>, name: &str, inputs: usize, outputs: usize, sparse: bool) -> Option<(Vec<i8>, Vec<u32>)> {
    let Array::I8(raw) = arrays.get(&format!("dec_{name}_weights_int8"))? else { return None };
    let Array::F32(scale) = arrays.get(&format!("dec_{name}_scale"))? else { return None };
    let mut w = vec![0i8; inputs * outputs];
    let mut wp = 0usize;
    if sparse {
        let Array::I32(idx) = arrays.get(&format!("dec_{name}_weights_idx"))? else { return None };
        let mut p = 0usize;
        for og in (0..outputs).step_by(8) {
            let n = *idx.get(p)? as usize;
            p += 1;
            for _ in 0..n {
                let pos = *idx.get(p)? as usize;
                p += 1;
                let blk = raw.get(wp..wp + 32)?;
                for o in 0..8 {
                    for i in 0..4 {
                        w[(og + o) * inputs + pos + i] = blk[o * 4 + i];
                    }
                }
                wp += 32;
            }
        }
    } else {
        for og in (0..outputs).step_by(8) {
            for ig in (0..inputs).step_by(4) {
                let blk = raw.get(wp..wp + 32)?;
                for o in 0..8 {
                    for i in 0..4 {
                        w[(og + o) * inputs + ig + i] = blk[o * 4 + i];
                    }
                }
                wp += 32;
            }
        }
    }
    (wp == raw.len() && scale.len() == outputs).then(|| (w, scale.clone()))
}

/// SHA-256, lowercase hexadecimal.
pub fn sha256_hex(bytes: &[u8]) -> String {
    Sha256::digest(bytes).iter().map(|b| format!("{b:02x}")).collect()
}

/// The weights file, checked against the pinned digest before anything reads it
/// (`weights_digest_mismatch`).
pub fn pinned(bytes: &[u8]) -> Result<BTreeMap<&'static str, Layer>, Reason> {
    if sha256_hex(bytes) != WEIGHTS_SHA256 {
        return Err(Reason::WeightsDigestMismatch);
    }
    decoder_layers(&parse(bytes)?)
}

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

    /// A record file of the decoder's arrays with values from a fixed seed: the same layouts as
    /// the draft's file, so every path of the expansion runs.
    pub(crate) fn synthetic_file(seed: u64, scale: f32) -> Vec<u8> {
        let mut s = seed;
        let mut rnd = move || {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            ((s >> 40) as f32 / (1u64 << 24) as f32 - 0.5) * 2.0
        };
        let mut file = Vec::new();
        let mut rec = |name: &str, kind: i32, data: Vec<u8>| {
            let size = data.len();
            let block = size.div_ceil(64) * 64;
            file.extend_from_slice(b"DNNw");
            file.extend_from_slice(&0i32.to_le_bytes());
            file.extend_from_slice(&kind.to_le_bytes());
            file.extend_from_slice(&(size as i32).to_le_bytes());
            file.extend_from_slice(&(block as i32).to_le_bytes());
            let mut nm = [0u8; 44];
            nm[..name.len()].copy_from_slice(name.as_bytes());
            file.extend_from_slice(&nm);
            file.extend_from_slice(&data);
            file.resize(file.len() + block - size, 0);
        };
        let f32b = |v: &[f32]| v.iter().flat_map(|x| x.to_le_bytes()).collect::<Vec<u8>>();
        for &(name, inputs, outputs, sparse) in LAYERS.iter().chain(std::iter::once(&OUTPUT)) {
            let b: Vec<f32> = (0..outputs).map(|_| rnd() * 0.2).collect();
            if sparse {
                // Every other input block of each group, so blocks are skipped and present.
                let mut idx = Vec::new();
                let mut w = Vec::new();
                for g in 0..outputs / 8 {
                    let blocks: Vec<i32> = (0..inputs / 4).filter(|k| (k + g) % 2 == 0).map(|k| 4 * k as i32).collect();
                    idx.push(blocks.len() as i32);
                    idx.extend(&blocks);
                    for _ in 0..blocks.len() * 32 {
                        w.push(rnd() * scale / (inputs as f32).sqrt());
                    }
                }
                rec(&format!("dec_{name}_weights_float"), 0, f32b(&w));
                rec(&format!("dec_{name}_weights_idx"), 1, idx.iter().flat_map(|x| x.to_le_bytes()).collect());
            } else {
                let w: Vec<f32> = (0..inputs * outputs).map(|_| rnd() * scale / (inputs as f32).sqrt()).collect();
                rec(&format!("dec_{name}_weights_float"), 0, f32b(&w));
            }
            rec(&format!("dec_{name}_bias"), 0, f32b(&b));
        }
        file
    }

    #[test]
    fn a_record_file_parses_and_expands() {
        let file = synthetic_file(7, 2.0);
        let arrays = parse(&file).unwrap();
        let layers = decoder_layers(&arrays).unwrap();
        assert_eq!(layers.len(), 29);
        let g = &layers["gru2_input"];
        assert_eq!((g.inputs, g.outputs, g.w.len()), (192, 192, 192 * 192));
        // Group 0 holds the even input blocks only: input 4 (block 1) is zero in output 0.
        assert_eq!(g.w[4], 0);
        assert_ne!(g.w[0], 0);
        assert_eq!(pinned(&file), Err(Reason::WeightsDigestMismatch));
    }

    #[test]
    fn malformed_files_are_refused() {
        let file = synthetic_file(7, 2.0);
        // A trailing byte, a bad magic, a non-zero pad byte.
        let mut f = file.clone();
        f.push(0);
        assert_eq!(parse(&f), Err(Reason::WeightsMalformed));
        let mut f = file.clone();
        f[0] = b'X';
        assert_eq!(parse(&f), Err(Reason::WeightsMalformed));
        // A missing array.
        let mut arrays = parse(&file).unwrap();
        arrays.remove("dec_glu3_bias");
        assert_eq!(decoder_layers(&arrays), Err(Reason::WeightsMalformed));
        // A non-finite weight.
        let mut arrays = parse(&file).unwrap();
        if let Some(Array::F32(v)) = arrays.get_mut("dec_conv2_weights_float") {
            v[5] = 0x7fc0_0000;
        }
        assert_eq!(decoder_layers(&arrays), Err(Reason::WeightsMalformed));
        // A group naming one input block twice.
        let mut arrays = parse(&file).unwrap();
        if let Some(Array::I32(v)) = arrays.get_mut("dec_output_weights_idx") {
            v[2] = v[1];
        }
        assert_eq!(decoder_layers(&arrays), Err(Reason::WeightsMalformed));
    }
}
