//! Integer-exact synthesis of a LEARNED 1-D audio codec — the decode side of a
//! `NeuralIntegerExact` audio codec, the audio analog of [`crate::int_hyper_synth`].
//! Pure i64, no float, no `ort`: the reconstructed int16 PCM is byte-identical on
//! every machine (*sample-equivalence*, a portable PCM hash).
//!
//! A decode through a FLOAT network lets ~1e-5 round-off make the waveform — and
//! its hash — diverge between machines. Here the synthesis `g_s` is a stack of integer 1-D
//! transposed convolutions (per-output-channel `(mult,shift)` requant) + integer
//! IGDN ([`crate::int_transform::igdn`]); the latent is recovered by the integer
//! rANS ([`crate::rans`]) under a per-channel factorized prior. Every operation is
//! i64 + an arithmetic shift, so two sinks honoring the spec produce the same PCM.
//!
//! It runs the model the export tool emits (`tools/wai_int_audio_export.py`): a
//! compact binary (`WAM1`) with int12 transposed-conv weights as `i16` and IGDN
//! `gamma` as `i32`.

use serde_json::Value;

use crate::int_transform::igdn;
use crate::rans::RansDecoder;

/// One layer of the integer synthesis `g_s`. `ConvT1d` weight is `[cin·cout·k]`
/// (layout `ic, oc, kx`); `Igdn` carries a full `c×c` gamma (`2^gf`) + beta
/// (`2^(2·afrac)`). Widened to i64 for the decode (arithmetic stays inside i64).
pub enum ALayer {
    ConvT1d { out_ch: usize, k: usize, stride: usize, pad: usize, opad: usize, weight: Vec<i64>, mult: Vec<i64>, bias: Vec<i64> },
    Igdn { c: usize, gamma: Vec<i64>, beta: Vec<i64> },
}

/// A parsed integer 1-D audio synthesis model (`WAM1`).
pub struct AudioModel {
    pub afrac: u32,
    pub gf: u32,
    pub shift: u32,
    pub layers: Vec<ALayer>,
}

/// Little-endian cursor over the binary model, bounds-checked: every read is of bytes
/// that are present, with its end formed in checked arithmetic, so a damaged model is
/// refused and never makes the parser allocate more than the bytes it holds.
struct Cur<'a> {
    b: &'a [u8],
    p: usize,
}
impl<'a> Cur<'a> {
    /// The next `n` elements of `size` bytes each.
    fn take(&mut self, n: usize, size: usize) -> Result<&'a [u8], &'static str> {
        let e = n.checked_mul(size).and_then(|len| self.p.checked_add(len)).ok_or("model truncated")?;
        let s = self.b.get(self.p..e).ok_or("model truncated")?;
        self.p = e;
        Ok(s)
    }
    fn u32(&mut self) -> Result<usize, &'static str> {
        Ok(u32::from_le_bytes(self.take(1, 4)?.try_into().unwrap()) as usize)
    }
    fn u8(&mut self) -> Result<u8, &'static str> {
        Ok(self.take(1, 1)?[0])
    }
    fn i16v(&mut self, n: usize) -> Result<Vec<i64>, &'static str> {
        Ok(self.take(n, 2)?.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]]) as i64).collect())
    }
    fn i32v(&mut self, n: usize) -> Result<Vec<i64>, &'static str> {
        Ok(self.take(n, 4)?.chunks_exact(4).map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64).collect())
    }
    fn i64v(&mut self, n: usize) -> Result<Vec<i64>, &'static str> {
        Ok(self.take(n, 8)?.chunks_exact(8).map(|c| i64::from_le_bytes(c.try_into().unwrap())).collect())
    }
}

/// The product of a layer's dimensions, refused where it overflows `usize`.
fn count(dims: &[usize]) -> Result<usize, &'static str> {
    dims.iter().try_fold(1usize, |a, &d| a.checked_mul(d)).ok_or("layer size overflows")
}

impl AudioModel {
    /// Parse the export tool's compact binary model: magic `WAM1`, then
    /// `afrac/gf/shift/n_layers` (u32) and per layer a `u8` tag — 0 convT1d, 2 igdn
    /// — with int12 weights as `i16`, IGDN `gamma` as `i32`, the rest `i64`, and
    /// nothing after the last layer.
    pub fn from_bytes(b: &[u8]) -> Result<Self, &'static str> {
        if b.get(0..4) != Some(b"WAM1") {
            return Err("bad model magic (want WAM1)");
        }
        let mut c = Cur { b, p: 4 };
        let (afrac, gf, shift) = (c.u32()? as u32, c.u32()? as u32, c.u32()? as u32);
        let n = c.u32()?;
        // Each layer is at least its tag byte, so a count beyond the bytes left is a
        // truncated model, and reserving for it would let four damaged bytes ask for a
        // large allocation.
        let mut layers = Vec::with_capacity(n.min(b.len() - c.p));
        for _ in 0..n {
            layers.push(match c.u8()? {
                0 => {
                    let inc = c.u32()?;
                    let out = c.u32()?;
                    let (k, stride, pad, opad) = (c.u32()?, c.u32()?, c.u32()?, c.u32()?);
                    ALayer::ConvT1d { out_ch: out, k, stride, pad, opad,
                        weight: c.i16v(count(&[inc, out, k])?)?, mult: c.i64v(out)?, bias: c.i64v(out)? }
                }
                2 => {
                    let cc = c.u32()?;
                    ALayer::Igdn { c: cc, gamma: c.i32v(count(&[cc, cc])?)?, beta: c.i64v(cc)? }
                }
                _ => return Err("bad layer tag"),
            });
        }
        // Nothing follows the last layer (`extensions/integer-payloads.md` §9).
        if c.p != b.len() {
            return Err("bytes follow the last layer");
        }
        Ok(Self { afrac, gf, shift, layers })
    }
}

/// One integer 1-D transposed conv with per-output-channel requant — see
/// [`reference::conv_t_1d`] for the definition. Runs the fast kernel
/// ([`crate::int_kernels`], as a 1-row 2-D transposed conv) unless the current
/// [`crate::int_kernels::Policy`] selects the reference; both give the same integer
/// for every output element.
#[allow(clippy::too_many_arguments)]
fn conv_t_1d(
    x: &[i64], cin: usize, t: usize, weight: &[i64], cout: usize, k: usize,
    stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
) -> (Vec<i64>, usize) {
    use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
    if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
        && let Some(g) = ConvGeom::transposed(cin, 1, t, cout, 1, k, 1, stride, 0, pad, 0, opad, x.len(), weight.len())
        && let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc])
    {
        return (out, g.wo);
    }
    reference::conv_t_1d(x, cin, t, weight, cout, k, stride, pad, opad, mult, bias, shift)
}

/// The scalar reference kernel the fast path is checked against, kept verbatim.
#[doc(hidden)]
pub mod reference {
    /// One integer 1-D transposed conv with per-output-channel requant
    /// `((raw·mult)>>shift)+bias`. `x` is `[cin, t]`; `weight` is `[cin][cout][k]`.
    /// Returns `([cout, to], to)`. Pure i64 (i128 only for the requant multiply) —
    /// mirrors the export byte-for-byte.
    #[allow(clippy::too_many_arguments)]
    pub fn conv_t_1d(
        x: &[i64], cin: usize, t: usize, weight: &[i64], cout: usize, k: usize,
        stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
    ) -> (Vec<i64>, usize) {
        let to = (t - 1) * stride + k + opad - 2 * pad;
        let tp = to + 2 * pad;
        let mut raw = vec![0i64; cout * tp];
        for ic in 0..cin {
            let xc = &x[ic * t..ic * t + t];
            for oc in 0..cout {
                let wk = &weight[(ic * cout + oc) * k..(ic * cout + oc) * k + k];
                let plane = oc * tp;
                for (it, &v) in xc.iter().enumerate() {
                    if v == 0 {
                        continue;
                    }
                    let base = plane + it * stride;
                    for kx in 0..k {
                        raw[base + kx] += wk[kx] * v;
                    }
                }
            }
        }
        let mut out = vec![0i64; cout * to];
        for oc in 0..cout {
            let (m, b) = (mult[oc] as i128, bias[oc]);
            for ot in 0..to {
                let r = raw[oc * tp + ot + pad] as i128;
                out[oc * to + ot] = ((r * m) >> shift) as i64 + b;
            }
        }
        (out, to)
    }
}

fn i64s(v: &Value) -> Vec<i64> {
    v.as_array().map(|a| a.iter().filter_map(|n| n.as_i64()).collect()).unwrap_or_default()
}

use crate::int_hyper_synth::{refused, Refused, RefusalKind, MAX_AFRAC, MAX_TENSOR};

/// A clip's coded form, typed: what a `WIA1` payload carries
/// (`extensions/integer-payloads.md` §5), and what the tool's JSON form carries, read
/// by [`AudioJson`]. `cdfs` holds one table per latent channel, all of one length.
#[derive(Debug, Clone, Copy)]
pub struct AudioStream<'a> {
    /// The fixed-point scale the stream was coded at; it must equal the model's.
    pub afrac: u32,
    pub n_ch: usize,
    pub lat_len: usize,
    /// Samples presented: the first `t_out` of the synthesis.
    pub t_out: usize,
    /// The symbol offset every channel's table is coded at (the JSON form's `lo`).
    pub offset: i32,
    pub median_q: &'a [i64],
    pub cdfs: &'a [Vec<i32>],
    pub rans: &'a [u8],
}

/// The tool's JSON bitstream (`tools/wai_int_audio_export.py`'s
/// `audio_bitstream.json`), read into the typed form. A converter input and a
/// debugging aid; the payload is `WIA1`. What the old JSON decoder left implicit is
/// explicit here: its `afrac` is read (the old decoder ignored it) and must be the
/// model's; every channel's table must have the first table's length (the old decoder
/// read every table at that length).
#[derive(Debug, Clone)]
pub struct AudioJson {
    pub afrac: u32,
    pub n_ch: usize,
    pub lat_len: usize,
    pub t_out: usize,
    pub offset: i32,
    pub median_q: Vec<i64>,
    pub cdfs: Vec<Vec<i32>>,
    pub rans: Vec<u8>,
}

impl AudioJson {
    /// Read the tool's JSON bitstream.
    pub fn from_value(bs: &Value) -> Result<Self, &'static str> {
        let uz = |k: &'static str| bs[k].as_u64().and_then(|v| usize::try_from(v).ok()).ok_or(k);
        let (n_ch, lat_len, t_out) = (uz("n_ch")?, uz("lat_len")?, uz("t_out")?);
        let afrac = bs["afrac"].as_u64().and_then(|v| u32::try_from(v).ok()).ok_or("afrac")?;
        let offset = bs["lo"].as_i64().and_then(|v| i32::try_from(v).ok()).ok_or("lo")?;
        let median_q = i64s(&bs["median_q"]);
        let cdfs = bs["cdfs"].as_array().ok_or("cdfs")?.iter()
            .map(|row| row.as_array().and_then(|r| r.iter().map(|n| n.as_i64().and_then(|v| i32::try_from(v).ok())).collect::<Option<Vec<i32>>>()))
            .collect::<Option<Vec<_>>>().ok_or("cdfs holds a value outside i32")?;
        let first = cdfs.first().ok_or("empty cdfs")?.len();
        if cdfs.iter().any(|c| c.len() != first) {
            return Err("the channels' tables are not of one length");
        }
        let rans = b64_decode(bs["rans_b64"].as_str().ok_or("rans_b64")?)?;
        Ok(AudioJson { afrac, n_ch, lat_len, t_out, offset, median_q, cdfs, rans })
    }

    /// The typed form, borrowing this one's tables.
    pub fn stream(&self) -> AudioStream<'_> {
        AudioStream {
            afrac: self.afrac, n_ch: self.n_ch, lat_len: self.lat_len, t_out: self.t_out, offset: self.offset,
            median_q: &self.median_q, cdfs: &self.cdfs, rans: &self.rans,
        }
    }
}

/// The model's output `[c, t]` on a `[c, t]` latent, planned from its layers' geometry
/// and parameter lengths alone, before anything is allocated
/// (`extensions/integer-payloads.md` §7). Layer by layer in order: the layer fits the
/// running channel count (`ModelMismatch`: output channels, kernel and stride
/// non-zero, `c·out·k` weights, `out` multipliers and biases, a requantisation shift
/// under 128; an inverse GDN of the running channel count, with `c·c` and `c`
/// parameters and shifts under 128); the uncropped length `(t − 1)·stride + k + opad`
/// and the buffer of `out` channels of it are within [`MAX_TENSOR`] (`TooLarge`); and
/// the crop `2·pad` leaves an output (`ModelMismatch`). Every size is formed in `u64`
/// with checked arithmetic, a size past `u64` being over the cap, so every target
/// refuses alike. Returns the output `[c, t]`, the most elements one layer holds, and
/// the multiply-accumulates the layers comprise (`extensions/integer-payloads.md` §7):
/// a transposed convolution's every input times `out·k`, an inverse GDN's every
/// element times `c`, summed in `u128`.
fn plan(model: &AudioModel, c: usize, t: usize) -> Result<(usize, usize, u64, u128), Refused> {
    use RefusalKind::*;
    const CAP: u64 = MAX_TENSOR as u64;
    let (mut c, mut t) = (c as u64, t as u64);
    // The most elements the reference kernels hold while one layer runs, its input
    // included: a transposed convolution's input, uncropped buffer and output; an
    // inverse GDN's input, output and one position's squares (`c` values of 16 bytes).
    let mut live = c * t;
    let mut macs: u128 = 0;
    for layer in &model.layers {
        match layer {
            ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => {
                let (cout, k, s, p, op) = (*out_ch as u64, *k as u64, *stride as u64, *pad as u64, *opad as u64);
                if cout == 0 || k == 0 || s == 0 {
                    return Err(refused(ModelMismatch, "a layer has zero output channels, kernel or stride"));
                }
                if c.checked_mul(cout).and_then(|v| v.checked_mul(k)) != Some(weight.len() as u64)
                    || mult.len() as u64 != cout || bias.len() as u64 != cout || model.shift > 127
                {
                    return Err(refused(ModelMismatch, "a layer's parameters do not match the running channel count"));
                }
                let full = (t - 1).checked_mul(s).and_then(|v| v.checked_add(k)).and_then(|v| v.checked_add(op));
                let full = full.filter(|&n| n <= CAP).ok_or(refused(TooLarge, "a layer's uncropped length exceeds 2^26"))?;
                if cout.checked_mul(full).is_none_or(|n| n > CAP) {
                    return Err(refused(TooLarge, "a layer would form a buffer over 2^26 elements"));
                }
                let crop = p.checked_mul(2).unwrap_or(u64::MAX);
                if full <= crop {
                    return Err(refused(ModelMismatch, "a layer's padding leaves no output"));
                }
                live = live.max(c * t + cout * full + cout * (full - crop));
                macs = macs.saturating_add((c * t) as u128 * (cout * k) as u128);
                (c, t) = (cout, full - crop);
            }
            ALayer::Igdn { c: cc, gamma, beta } => {
                if *cc as u64 != c || c.checked_mul(c) != Some(gamma.len() as u64) || beta.len() as u64 != c
                    || model.gf > 127 || model.afrac > 127
                {
                    return Err(refused(ModelMismatch, "an IGDN layer does not match the running channel count"));
                }
                live = live.max(2 * c * t + 2 * c);
                macs = macs.saturating_add((c * t) as u128 * c as u128);
            }
        }
    }
    // Every running size is at most 2^26 here, so it fits usize on every target.
    Ok((c as usize, t as usize, live, macs))
}

/// What [`decode_typed`] checks before it decodes a symbol, and the buffers it then
/// forms. Returns the synthesis's length, the most bytes the decode holds at once
/// on the reference kernels beyond the payload and the parsed model (the latent and
/// each layer's buffers, `i64` each, then the synthesis's output, `i64`, with the
/// PCM, `i16`), and the synthesis's multiply-accumulates on the latent.
///
/// The latent's length is the one `t_out` needs and no longer: the synthesis of
/// `lat_len` latent samples gives at least `t_out` samples (`GeometryMismatch`
/// otherwise), and that of `lat_len − 1` gives fewer (`GeometryMismatch`), so a
/// payload cannot declare a large latent to present a short clip. A length for which
/// the model leaves no output gives none.
pub(crate) fn audio_plan(model: &AudioModel, s: &AudioStream<'_>) -> Result<(usize, u64, u128), Refused> {
    use RefusalKind::*;
    let afrac = model.afrac;
    if !(1..=MAX_AFRAC).contains(&afrac) {
        return Err(refused(ModelMismatch, "the model's fixed-point scale is outside 1..=30"));
    }
    if s.afrac != afrac {
        return Err(refused(ModelMismatch, "the stream's afrac is not the model's"));
    }
    let (c, l) = (s.n_ch, s.lat_len);
    if c == 0 || l == 0 || s.t_out == 0 || s.median_q.len() != c || s.cdfs.len() != c {
        return Err(refused(FieldRange, "a dimension is zero, or the medians or tables are not one per channel"));
    }
    if s.median_q.iter().any(|m| m.unsigned_abs() > 1 << 47) {
        return Err(refused(FieldRange, "a median's magnitude exceeds 2^47"));
    }
    let len = s.cdfs[0].len();
    let ok = |t: &Vec<i32>| t.len() == len && (3..=4097).contains(&len) && t[0] == 0 && t[len - 1] == 1 << 16
        && t.windows(2).all(|p| p[0] < p[1]);
    if !s.cdfs.iter().all(ok) {
        return Err(refused(CdfMalformed, "a table is not a CDF from 0 to 65536 of the channels' one length"));
    }
    if c.checked_mul(l).is_none_or(|n| n > MAX_TENSOR) || s.t_out > MAX_TENSOR {
        return Err(refused(TooLarge, "the latent or the clip exceeds 2^26 elements"));
    }
    let (oc, t, live, macs) = plan(model, c, l)?;
    if oc != 1 {
        return Err(refused(ModelMismatch, "the model does not give one channel"));
    }
    if t < s.t_out {
        return Err(refused(GeometryMismatch, "the synthesis gives fewer samples than t_out"));
    }
    let shorter = if l > 1 { plan(model, c, l - 1).map_or(0, |(_, t, ..)| t) } else { 0 };
    if shorter >= s.t_out {
        return Err(refused(GeometryMismatch, "the latent is longer than t_out needs: one sample fewer gives t_out samples"));
    }
    let working = (8 * live).max(8 * t as u64 + 2 * s.t_out as u64);
    Ok((t, working, macs))
}

/// One layer's arithmetic on the values `x` holds, checked before it runs, as
/// [`crate::int_hyper_synth::run_net_checked`] checks a 2-D layer: every product and
/// partial sum and the requantised value of a convolution fit `i64`; IGDN's terms,
/// norm and root fit their types.
fn layer_in_range(model: &AudioModel, layer: &ALayer, x: &[i64], c: usize) -> Result<(), Refused> {
    use crate::int_kernels::max_abs;
    const I64: u128 = i64::MAX as u128;
    match layer {
        ALayer::ConvT1d { k, weight, mult, bias, .. } => {
            if model.shift > 127 {
                return Err(refused(RefusalKind::ModelMismatch, "requantisation shift of 128 or more"));
            }
            let terms = (c as u128).saturating_mul(*k as u128);
            let acc = terms.saturating_mul(max_abs(weight) as u128).saturating_mul(max_abs(x) as u128);
            if acc > I64 {
                return Err(refused(RefusalKind::FieldRange, "a latent could take an accumulation outside i64"));
            }
            let q = ((acc * max_abs(mult) as u128) >> model.shift) + 1;
            if q.saturating_add(max_abs(bias) as u128) > I64 {
                return Err(refused(RefusalKind::FieldRange, "a latent could take a requantised value outside i64"));
            }
            Ok(())
        }
        ALayer::Igdn { gamma, beta, .. } => crate::int_hyper_synth::igdn_check(x, c, gamma, beta, model.gf, model.afrac)
            .map_err(|(why, values)| refused(if values { RefusalKind::FieldRange } else { RefusalKind::ModelMismatch }, why)),
    }
}

/// Decode a typed clip → int16 PCM, mono, byte-exact on every machine, refusing
/// rather than panicking or overflowing on any input:
///
/// 1. the stream's `afrac` is the model's, in `1..=30`; every table is well formed;
///    the model's layers fit the latent's shape, give one channel, and give at least
///    `t_out` samples from a latent no longer than that needs, with no buffer over
///    [`MAX_TENSOR`] elements — all before anything is allocated ([`audio_plan`]);
/// 2. per channel, rANS-decode the symbol `s`; the latent is `s·2^afrac + median_q`
///    ([`crate::int_entropy::dequantize`]), and a stream that runs out before its last
///    symbol is refused;
/// 3. integer synthesis `g_s`, each layer's arithmetic checked before it runs;
/// 4. the first `t_out` samples, `(x·32767 + 2^(afrac−1)) >> afrac` clamped to `i16`.
pub fn decode_typed(model: &AudioModel, s: &AudioStream<'_>) -> Result<Vec<i16>, Refused> {
    use crate::int_entropy::dequantize;
    use RefusalKind::*;
    audio_plan(model, s)?;
    let (afrac, c, l, len) = (model.afrac, s.n_ch, s.lat_len, s.cdfs[0].len());

    // 1. factorized rANS-decode the per-channel latent (channel-major), rebuild y_q.
    // The latent grows a symbol at a time, and the first symbol that runs the stream
    // out refuses the payload before its value is used.
    let mut y = Vec::new();
    let mut dec = RansDecoder::new(s.rans).map_err(|_| refused(FieldRange, "the latent stream is not a rANS stream"))?;
    for ch in 0..c {
        for _ in 0..l {
            let sym = dec.decode(&s.cdfs[ch], len, s.offset);
            if dec.is_exhausted() {
                return Err(refused(StreamExhausted, "the latent stream ended before its last symbol"));
            }
            crate::int_hyper_synth::push_capped(&mut y, dequantize(sym, afrac, s.median_q[ch]).ok_or(refused(FieldRange, "a latent value leaves i64"))?, c * l);
        }
    }
    if !dec.is_canonical_end() {
        return Err(refused(StreamNoncanonical, "the latent stream does not end canonically after its last symbol"));
    }
    // 2. integer synthesis → waveform, each layer checked on the values it is given.
    let (mut x, mut cc, mut tt) = (y, c, l);
    for layer in &model.layers {
        layer_in_range(model, layer, &x, cc)?;
        (x, cc, tt) = run_layer(model, layer, x, cc, tt);
    }
    // 3. waveform → int16 PCM (round-half-up, clamp): the first t_out samples.
    let half = 1i64 << (afrac - 1);
    let mut pcm = Vec::with_capacity(s.t_out);
    for &v in &x[..s.t_out] {
        let num = v.checked_mul(32767).and_then(|n| n.checked_add(half))
            .ok_or(refused(FieldRange, "a synthesised sample leaves i64 when scaled"))?;
        pcm.push((num >> afrac).clamp(-32768, 32767) as i16);
    }
    Ok(pcm)
}

/// One layer of the synthesis on `x` (`[c, t]`): the new activations and shape.
fn run_layer(model: &AudioModel, layer: &ALayer, x: Vec<i64>, c: usize, t: usize) -> (Vec<i64>, usize, usize) {
    match layer {
        ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => {
            let (nx, to) = conv_t_1d(&x, c, t, weight, *out_ch, *k, *stride, *pad, *opad, mult, bias, model.shift);
            (nx, *out_ch, to)
        }
        ALayer::Igdn { gamma, beta, .. } => (igdn(&x, c, t, model.afrac, gamma, beta, model.gf), c, t),
    }
}

/// Decode the tool's JSON bitstream (`audio_bitstream.json`) → int16 PCM:
/// [`AudioJson`] then [`decode_typed`]. A debugging and converter path; the payload
/// is `WIA1`. Byte-identical on every machine.
pub fn decode(model: &AudioModel, bs: &Value) -> Result<Vec<i16>, &'static str> {
    decode_typed(model, &AudioJson::from_value(bs)?.stream()).map_err(|r| r.detail)
}

/// FNV-1a 64-bit over the little-endian int16 PCM bytes — the dependency-free
/// checksum the export tool and this decoder share to prove byte-for-byte agreement
/// across languages (sample-equivalence).
pub fn pcm_fnv1a(pcm: &[i16]) -> u64 {
    let mut hash: u64 = 14695981039346656037;
    for &s in pcm {
        for &b in &s.to_le_bytes() {
            hash = (hash ^ b as u64).wrapping_mul(1099511628211);
        }
    }
    hash
}

/// Dependency-free base64 decode (standard alphabet, no padding required) — the
/// rANS bytes ride in JSON as base64, same as `int_video`.
fn b64_decode(s: &str) -> Result<Vec<u8>, &'static str> {
    fn val(c: u8) -> Result<u8, &'static str> {
        Ok(match c {
            b'A'..=b'Z' => c - b'A',
            b'a'..=b'z' => c - b'a' + 26,
            b'0'..=b'9' => c - b'0' + 52,
            b'+' => 62,
            b'/' => 63,
            _ => return Err("bad base64 char"),
        })
    }
    let s: Vec<u8> = s.bytes().filter(|&c| c != b'=' && !c.is_ascii_whitespace()).collect();
    let mut out = Vec::with_capacity(s.len() * 3 / 4);
    for chunk in s.chunks(4) {
        let mut acc = 0u32;
        let mut bits = 0;
        for &c in chunk {
            acc = (acc << 6) | val(c)? as u32;
            bits += 6;
        }
        // each full byte sits in the high `bits-8k` positions
        for k in 0..bits / 8 {
            out.push((acc >> (bits - 8 * (k + 1))) as u8);
        }
    }
    Ok(out)
}

/// The fast 1-D transposed conv against the scalar reference, element by element.
#[cfg(test)]
mod plan_tests {
    use super::*;

    /// The committed synthesis's planned multiply-accumulates on the committed latent are
    /// its layers' terms at the lengths the reference kernel gives (run here on a zero
    /// latent): a transposed convolution's every input times `out·k`, an inverse GDN's
    /// every element times `c`.
    #[test]
    fn the_planned_macs_are_every_layers_terms() {
        let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_audio");
        let model = AudioModel::from_bytes(&std::fs::read(dir.join("audio_model.bin")).unwrap()).unwrap();
        let wia1 = std::fs::read(dir.join("audio_bitstream.wia1")).unwrap();
        let a = crate::int_wire::parse_wia1(&wia1).unwrap();
        let (c0, l) = (a.n_ch as usize, a.lat_len as usize);
        let (.., macs) = plan(&model, c0, l).unwrap();
        let (mut c, mut t, mut terms) = (c0, l, 0u128);
        for layer in &model.layers {
            match layer {
                ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => {
                    let (_, to) = reference::conv_t_1d(&vec![0; c * t], c, t, weight, *out_ch, *k, *stride, *pad, *opad, mult, bias, model.shift);
                    terms += (c * t) as u128 * (out_ch * k) as u128;
                    (c, t) = (*out_ch, to);
                }
                ALayer::Igdn { .. } => terms += (c * t) as u128 * c as u128,
            }
        }
        assert!(terms > 0);
        assert_eq!(macs, terms);
    }
}

#[cfg(test)]
mod kernel_equivalence {
    use super::*;
    use crate::int_kernels::testkit::{Rng, FAST};
    use crate::int_kernels::with_policy;

    #[test]
    fn conv_t_1d_equals_reference() {
        let mut g = Rng(0xA0D_0001);
        let mut ran = 0;
        for case in 0..500 {
            let k = 1 + g.below(9);
            let stride = 1 + g.below(5);
            let pad = g.below(k);
            let opad = if stride > 1 { g.below(stride) } else { 0 };
            let t = 1 + g.below(40);
            let (cin, cout) = (1 + g.below(6), 1 + g.below(6));
            if (t - 1) * stride + k + opad <= 2 * pad {
                continue;
            }
            let (xr, wr) = [(1 << 15, 2047), (1 << 40, 1 << 10)][case % 2];
            let x = g.vec(cin * t, xr, 30);
            let wt = g.vec(cin * cout * k, wr, 5);
            let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 16).abs()).collect();
            let bias = g.vec(cout, 1 << 16, 0);
            let want = reference::conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24);
            for p in FAST {
                let got = with_policy(p, || conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24));
                assert_eq!(got, want, "conv_t_1d {p:?} k={k} s={stride} p={pad} op={opad} t={t}");
            }
            ran += 1;
        }
        assert!(ran > 250, "too few valid cases ({ran})");
    }

    /// Overflow only at a position the reference accumulates into and then crops:
    /// a length-1 input, kernel 3, pad 1 keeps only the middle of the reference's
    /// 3-sample raw buffer (`2·2^40·1 = 2^41`); both outer samples are
    /// `2·2^40·2^22 = 2^63` and overflow `i64`. The fast path must panic where the
    /// reference does and otherwise give the same 2^41.
    #[test]
    fn overflow_only_at_a_cropped_border_behaves_like_reference() {
        use crate::int_kernels::testkit::overflow_checks;
        use crate::int_kernels::Policy;
        let (cin, cout, k) = (2usize, 1usize, 3usize);
        let x = vec![1i64 << 40; cin];
        let wt: Vec<i64> = (0..cin * cout).flat_map(|_| [1i64 << 22, 1, 1 << 22]).collect();
        let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
        let run = |p: Policy| {
            std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, cin, 1, &wt, cout, k, 1, 1, 0, &mult, &bias, 0))).ok()
        };
        let want = run(Policy::reference());
        if overflow_checks() {
            assert!(want.is_none(), "the reference overflows at a cropped sample in a checked build");
        } else {
            assert_eq!(want, Some((vec![1i64 << 41], 1)));
        }
        for p in FAST {
            assert_eq!(run(p), want, "fast must match the reference, panic included ({p:?})");
        }
    }

    /// A one-sample input upsampled by any stride is one output sample (k = 3,
    /// pad 1): a handful of operations for the reference. The fast path must return
    /// the same sample without overflowing, or allocating or looping per stride step.
    #[test]
    fn a_one_sample_input_upsampled_by_any_stride() {
        use crate::int_kernels::Policy;
        let (x, wt) = ([5i64], [1i64, 2, 3]);
        // usize::MAX first: a kernel allocating a table per phase fails there at once
        // ("capacity overflow") instead of committing memory for 2^29 tables later.
        for s in [usize::MAX, 1, 2, 1 << 13, 1 << 16, 1 << 29] {
            let run = |p: Policy| with_policy(p, || conv_t_1d(&x, 1, 1, &wt, 1, 3, s, 1, 0, &[1], &[0], 0));
            let want = run(Policy::reference());
            assert_eq!(want, (vec![10], 1), "stride {s}: the middle of 5·[1, 2, 3]");
            for p in FAST {
                assert_eq!(run(p), want, "stride {s} {p:?}");
            }
        }
    }

    /// Several input samples upsampled by a stride and cropped by padding so large
    /// that the reference's uncropped buffer `cout · (to + 2·pad)` cannot be sized:
    /// three samples, kernel 1, stride = pad = 2^62, one output sample. With two
    /// output channels the size overflows `usize` (the reference panics with
    /// overflow checks on, and with them off indexes out of the wrapped buffer);
    /// with one it is 2^63 + 1 samples, more than any allocation can hold ("capacity
    /// overflow"). The reference panics in every build, so every policy must, and
    /// the fast kernel must not be entered.
    #[cfg(target_pointer_width = "64")]
    #[test]
    fn several_samples_whose_uncropped_buffer_cannot_be_sized() {
        use crate::int_kernels::testkit::fast_conv_runs;
        use crate::int_kernels::Policy;
        let s = 1usize << 62;
        for cout in [2usize, 1] {
            let (x, wt) = ([1i64, 2, 3], vec![7i64; cout]);
            let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
            let run = |p: Policy| {
                fast_conv_runs(|| std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, 1, 3, &wt, cout, 1, s, s, 0, &mult, &bias, 0))).ok())
            };
            assert_eq!(run(Policy::reference()), (None, 0), "{cout} channels: the reference panics");
            for p in FAST {
                assert_eq!(run(p), (None, 0), "{cout} channels {p:?}");
            }
        }
    }

    /// Geometry whose padded input would be more than `PAD_FACTOR` (16) times what
    /// the reference reads and writes is refused by the fast path, and the reference
    /// runs: the fast kernel is not entered and the samples are the reference's. At
    /// exactly the factor the fast kernel runs. A one-sample input, kernel 1, stride
    /// 1, cropped by `pad` with output padding `2·pad` restoring one output sample
    /// that no tap reaches: the padded input is `cin · (1 + pad)` samples against
    /// `cin + cin + 1` the reference touches (input, weights, output).
    #[test]
    fn padding_beyond_the_pad_factor_runs_the_reference() {
        use crate::int_kernels::testkit::fast_conv_runs;
        // (channels, pad, expected fast-kernel runs): 1 · 48 = 16 · 3 and
        // 2 · 40 = 16 · 5 at the factor; 1 · 49 and 2 · 41 beyond it.
        for (cin, pad, runs) in [(1usize, 47usize, 1u64), (1, 48, 0), (2, 39, 1), (2, 40, 0)] {
            let (x, wt) = (vec![5i64; cin], vec![3i64; cin]);
            let want = reference::conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1);
            assert_eq!(want, (vec![7], 1), "reference, {cin} channels, pad {pad}");
            for p in FAST {
                let got = fast_conv_runs(|| with_policy(p, || conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1)));
                assert_eq!(got, (want.clone(), runs), "{cin} channels, pad {pad} {p:?}");
            }
        }
    }
}