//! The lowering `dred-int/1` (extensions/opus-dred.md §4.4–§4.6): the draft's float decoder
//! weights to two `WHG1` graphs of `wai-ops/1`, `init` and `step`, in one canonical byte
//! sequence each.
//!
//! - **Weights** become integers at `2^24` and biases at the scale of the accumulator they join,
//!   each rounded half to even from the float32 value's exact bits ([`quantize`]): no float
//!   arithmetic is done, so every implementation gets the same integers.
//! - **Activations** are integers at `2^15`. `tanh` and the sigmoid are `lut16` tables, read at
//!   `2^12` and `2^11`, whose entries are an integer computation fixed here ([`tanh_table`],
//!   [`sigmoid_table`]).
//! - **The layers** are `conv` 1×1 over `[C, 1, 1]` tensors with a `pre` bias and one shift
//!   rounding half up; a GRU is its gates' convolution over `[x, h]`, the candidate's two
//!   convolutions, `mul`, `add` and the state update `hc + z·(h − hc)`; a GLU is a sigmoid gate
//!   and `mul`; a causal convolution is `concat` of the previous chunk's input and this one's.

use std::collections::BTreeMap;

use wai_ops::graph::{opcode, BiasKind, Graph, InputRec, Op, Params, Round, Rq, TensorRec, CAP};

use super::weights::Layer;
use super::Reason;

/// The lowering's identifier, as the registry's `ops` member names it.
pub const LOWERING: &str = "dred-int/1";
/// Fractional bits: weights, activations, the graph inputs and the output features, the
/// candidate's pre-activation, and the two tables' inputs.
pub const W_FRAC: i32 = 24;
pub const A_FRAC: i32 = 15;
pub const IN_FRAC: i32 = 16;
pub const PRE_FRAC: i32 = 20;
pub const TANH_IN: i32 = 12;
pub const SIG_IN: i32 = 11;
/// The largest magnitude of an activation, and of a table's input (`lut16`, §5.6).
pub const ACT_MAX: i64 = 32767;
/// The range of a latent and a state input (Q16) and of an output feature.
pub const IN_MAX: i64 = (1 << 31) - 1;
/// `⌊e^(−1/2048)·2^62⌋`, the one constant of the tables (extensions/opus-dred.md §4.5).
pub const EXP_C0: u128 = 4_609_434_768_280_048_980;

const N: u32 = 64;

/// `R_half_even(x·2^f)` for the float32 whose bits are `bits`, from the bits alone: the value
/// is `(−1)^s·M·2^E` with `M` the significand (with its hidden bit when the exponent field is
/// not 0) and `E` the exponent; the result is `M·2^(E+f)` rounded half to even, signed. `None`
/// for a value that is not finite, or whose result would pass `2^62`.
pub fn quantize(bits: u32, f: i32) -> Option<i64> {
    let e = ((bits >> 23) & 0xff) as i32;
    let m = (bits & 0x7f_ffff) as i64;
    if e == 255 {
        return None;
    }
    let (mant, exp) = if e == 0 { (m, -149) } else { (m | 0x80_0000, e - 150) };
    let k = exp + f;
    let mag = if k >= 0 {
        if k > 38 {
            return None;
        }
        mant << k
    } else {
        let s = (-k) as u32;
        if s >= 63 {
            0
        } else {
            let q = mant >> s;
            let r = mant - (q << s);
            let h = 1i64 << (s - 1);
            if r > h || (r == h && q & 1 == 1) { q + 1 } else { q }
        }
    };
    Some(if bits >> 31 == 1 { -mag } else { mag })
}

/// `⌊e^(−u/2048)·2^62⌋` approximately, as the profile defines it exactly: starting from
/// `2^62`, for each bit `k` of `u` from bit 0 that is set, multiply by `c_k` and floor by
/// `2^62`, where `c_0` is [`EXP_C0`] and `c_(k+1) = ⌊c_k² / 2^62⌋`. `u < 2^15`.
pub fn exp_q62(u: u32) -> u128 {
    let mut c = EXP_C0;
    let mut e: u128 = 1 << 62;
    for k in 0..15 {
        if (u >> k) & 1 == 1 {
            e = (e * c) >> 62;
        }
        c = (c * c) >> 62;
    }
    e
}

fn round_ratio(num: u128, den: u128) -> i64 {
    ((2 * num + den) / (2 * den)).min(ACT_MAX as u128) as i64
}

/// The `tanh` table: entry `v + 32768` for `v` in `−32767..=32767` is `tanh(v / 2^12)` at
/// `2^15`, computed as `sgn(v)·min(32767, ⌊(2·num + den) / (2·den)⌋)` with `E = exp_q62(|v|)`,
/// `num = 2^15·(2^62 − E)` and `den = 2^62 + E`. Entry 0 (`v = −32768`) is 0, as `lut16`
/// requires; no certified input reaches it.
pub fn tanh_table() -> Vec<i16> {
    let mut t = vec![0i16; 65536];
    for v in -32767i32..=32767 {
        let e = exp_q62(v.unsigned_abs());
        let mag = round_ratio((1u128 << 15) * ((1u128 << 62) - e), (1u128 << 62) + e);
        t[(v + 32768) as usize] = (v.signum() as i64 * mag) as i16;
    }
    t
}

/// The sigmoid table: entry `v + 32768` is `σ(v / 2^11)` at `2^15`, `min(32767, ⌊(2·num + den)
/// / (2·den)⌋)` with `E = exp_q62(|v|)`, `den = 2^62 + E`, and `num = 2^15·2^62` for `v ≥ 0`,
/// `2^15·E` for `v < 0`. Entry 0 is 0.
pub fn sigmoid_table() -> Vec<i16> {
    let mut t = vec![0i16; 65536];
    for v in -32767i32..=32767 {
        let e = exp_q62(v.unsigned_abs());
        let num = if v >= 0 { 1u128 << 77 } else { e << 15 };
        t[(v + 32768) as usize] = round_ratio(num, (1u128 << 62) + e) as i16;
    }
    t
}

/// A graph under construction, numbered as integer-operators §13.0 numbers a lowered graph:
/// inputs first, then each operator's output in operator order.
struct Builder {
    inputs: Vec<InputRec>,
    tensors: Vec<TensorRec>,
    ops: Vec<Op>,
    tanh: Vec<i16>,
    sigmoid: Vec<i16>,
}

fn act(c: u32) -> TensorRec {
    TensorRec::clamp(c, -ACT_MAX, ACT_MAX)
}

fn rq_pre(shift: i32, bias: Vec<i64>) -> Rq {
    Rq { round: Round::HalfUp, bias_kind: BiasKind::Pre, m_kind: 0, s_kind: 1, m_bytes: 0, b_bytes: 8, mult: vec![], shift: vec![shift as u8], bias }
}

fn rq_shift(shift: u8) -> Rq {
    Rq::plain(Round::HalfUp, shift)
}

/// One layer's integers: `w[o][i]` at `2^24`, `b[o]` at `2^(in_frac + 24)`.
struct QLayer {
    inputs: usize,
    w: Vec<i64>,
    b: Vec<i64>,
}

fn qlayer(l: &Layer, in_frac: i32) -> Result<QLayer, Reason> {
    let q = |bits: &u32, f: i32| quantize(*bits, f).ok_or(Reason::WeightsMalformed);
    let w = l.w.iter().map(|b| q(b, W_FRAC)).collect::<Result<Vec<_>, _>>()?;
    let b = l.b.iter().map(|b| q(b, in_frac + W_FRAC)).collect::<Result<Vec<_>, _>>()?;
    if w.iter().any(|&v| i32::try_from(v).is_err()) {
        return Err(Reason::WeightsMalformed);
    }
    Ok(QLayer { inputs: l.inputs, w, b })
}

impl QLayer {
    /// Rows `rows` of the layer.
    fn rows(&self, rows: std::ops::Range<usize>) -> (Vec<i64>, Vec<i64>) {
        let n = self.inputs;
        (self.w[rows.start * n..rows.end * n].to_vec(), self.b[rows].to_vec())
    }
}

impl Builder {
    fn new() -> Builder {
        Builder { inputs: vec![], tensors: vec![], ops: vec![], tanh: tanh_table(), sigmoid: sigmoid_table() }
    }
    fn input(&mut self, rec: TensorRec) -> u32 {
        let t = self.tensors.len() as u32;
        self.tensors.push(rec);
        self.inputs.push(InputRec { tensor: t, max_h: 1, max_w: 1, kind: 0, src: 0, init: 0 });
        t
    }
    fn op(&mut self, code: u8, ins: Vec<u32>, out: TensorRec, p: Params) -> u32 {
        let t = self.tensors.len() as u32;
        self.tensors.push(out);
        self.ops.push(Op { code, ins, outs: vec![t], p });
        t
    }
    fn channels(&self, t: u32) -> u32 {
        self.tensors[t as usize].channels
    }
    /// `conv` 1×1: `w` is `[cout][cin]`, the bias `pre`, one shift, half up.
    fn conv(&mut self, x: u32, w: Vec<i64>, b: Vec<i64>, shift: i32, out: TensorRec) -> u32 {
        let cout = b.len() as u32;
        debug_assert_eq!(w.len(), cout as usize * self.channels(x) as usize);
        let p = Params::Conv { cout, kh: 1, kw: 1, sh: 1, sw: 1, ph: 0, pw: 0, groups: 1, w_bytes: 4, rq: rq_pre(shift, b), w };
        self.op(opcode::CONV, vec![x], out, p)
    }
    fn lut(&mut self, x: u32, tanh: bool) -> u32 {
        let t = if tanh { self.tanh.clone() } else { self.sigmoid.clone() };
        let c = self.channels(x);
        self.op(opcode::LUT16, vec![x], TensorRec::derived(c), Params::Lut16 { e_in: -128, e_out: -128, t })
    }
    /// A dense layer with `tanh`, from activations at `2^in_frac`.
    fn dense_tanh(&mut self, x: u32, l: &QLayer, in_frac: i32) -> u32 {
        let (w, b) = l.rows(0..l.b.len());
        let c = b.len() as u32;
        let pre = self.conv(x, w, b, in_frac + W_FRAC - TANH_IN, act(c));
        self.lut(pre, true)
    }
    fn concat(&mut self, ins: Vec<u32>) -> u32 {
        let c = ins.iter().map(|&t| self.channels(t)).sum();
        self.op(opcode::CONCAT, ins, TensorRec::derived(c), Params::None)
    }
    fn split(&mut self, x: u32, start: u32, count: u32, out: TensorRec) -> u32 {
        self.op(opcode::SPLIT, vec![x], out, Params::Split { start, count })
    }
    fn mul(&mut self, a: u32, b: u32, shift: u8, out: TensorRec) -> u32 {
        self.op(opcode::MUL, vec![a, b], out, Params::Mul(rq_shift(shift)))
    }
    fn add(&mut self, a: u32, b: u32, out: TensorRec) -> u32 {
        self.op(opcode::ADD, vec![a, b], out, Params::None)
    }
    /// A GRU (draft §2.2.2): `h' = z·h + (1 − z)·hc`, written `hc + z·(h − hc)`.
    fn gru(&mut self, x: u32, h: u32, input: &QLayer, recurrent: &QLayer) -> u32 {
        let n = N as usize;
        let nx = input.inputs;
        // The gates: rows 0..2N of both layers over [x, h], one accumulation.
        let mut w = Vec::with_capacity(2 * n * (nx + n));
        let mut b = Vec::with_capacity(2 * n);
        for o in 0..2 * n {
            w.extend_from_slice(&input.w[o * nx..(o + 1) * nx]);
            w.extend_from_slice(&recurrent.w[o * n..(o + 1) * n]);
            b.push(input.b[o] + recurrent.b[o]);
        }
        let xh = self.concat(vec![x, h]);
        let gates_pre = self.conv(xh, w, b, A_FRAC + W_FRAC - SIG_IN, act(2 * N));
        let gates = self.lut(gates_pre, false);
        let z = self.split(gates, 0, N, TensorRec::derived(N));
        let r = self.split(gates, N, N, TensorRec::derived(N));
        // The candidate: tanh(h_in + r·h_rec), at 2^20 until the table's scale.
        let (wi, bi) = input.rows(2 * n..3 * n);
        let h_in = self.conv(x, wi, bi, A_FRAC + W_FRAC - PRE_FRAC, TensorRec::derived(N));
        let (wr, br) = recurrent.rows(2 * n..3 * n);
        let h_rec = self.conv(h, wr, br, A_FRAC + W_FRAC - PRE_FRAC, TensorRec::derived(N));
        let rh = self.mul(r, h_rec, A_FRAC as u8, TensorRec::derived(N));
        let pre = self.add(h_in, rh, TensorRec::derived(N));
        let pre12 = self.op(opcode::RESCALE, vec![pre], act(N), Params::Rescale(rq_shift((PRE_FRAC - TANH_IN) as u8)));
        let hc = self.lut(pre12, true);
        // h' = hc + z·(h − hc).
        let neg = Rq { round: Round::Floor, bias_kind: BiasKind::None, m_kind: 1, s_kind: 1, m_bytes: 4, b_bytes: 0, mult: vec![-1], shift: vec![0], bias: vec![] };
        let nhc = self.op(opcode::RESCALE, vec![hc], TensorRec::derived(N), Params::Rescale(neg));
        let d = self.add(h, nhc, TensorRec::derived(N));
        let zd = self.mul(z, d, A_FRAC as u8, TensorRec::derived(N));
        self.add(hc, zd, act(N))
    }
    /// A GLU (draft §2.2.2): `y·σ(W·y + b)`.
    fn glu(&mut self, y: u32, l: &QLayer) -> u32 {
        let (w, b) = l.rows(0..l.b.len());
        let pre = self.conv(y, w, b, A_FRAC + W_FRAC - SIG_IN, act(N));
        let g = self.lut(pre, false);
        self.mul(y, g, A_FRAC as u8, TensorRec::derived(N))
    }
    fn finish(self, outputs: Vec<u32>) -> Graph {
        Graph { max_elems: CAP as u32, scalars: vec![], inputs: self.inputs, outputs, tensors: self.tensors, ops: self.ops }
    }
}

/// The two graphs of `dred-int/1`, from the decoder's layers.
pub struct Lowered {
    /// Input 0: the initial state `[50, 1, 1]` at `2^16`. Outputs: the five GRU states.
    pub init: Graph,
    /// Inputs: the latent `[26, 1, 1]` at `2^16`, the five GRU states, the five convolution
    /// memories. Outputs: the 80 features at `2^16`, the five states, the five memories.
    pub step: Graph,
}

/// Lower the decoder's layers (`weights::decoder_layers`) to the graphs of §4.6.
pub fn lower(layers: &BTreeMap<&'static str, Layer>) -> Result<Lowered, Reason> {
    let q = |name: &str, in_frac: i32| qlayer(&layers[name], in_frac);
    // init: IS → tanh(hidden_init) → tanh(gru_init) → five states.
    let mut b = Builder::new();
    let is = b.input(TensorRec::clamp(50, -IN_MAX, IN_MAX));
    let h0 = b.dense_tanh(is, &q("hidden_init", IN_FRAC)?, IN_FRAC);
    let s0 = b.dense_tanh(h0, &q("gru_init", A_FRAC)?, A_FRAC);
    let outs: Vec<u32> = (0..5).map(|k| b.split(s0, k * N, N, TensorRec::derived(N))).collect();
    let init = b.finish(outs);

    // step.
    let mut b = Builder::new();
    let z = b.input(TensorRec::clamp(26, -IN_MAX, IN_MAX));
    let hs: Vec<u32> = (0..5).map(|_| b.input(act(N))).collect();
    let cs: Vec<u32> = (0..5).map(|_| b.input(act(32))).collect();
    let mut t = b.dense_tanh(z, &q("dense1", IN_FRAC)?, IN_FRAC);
    let (mut new_h, mut new_c) = (vec![], vec![]);
    for k in 0..5 {
        let gi = q(&format!("gru{}_input", k + 1), A_FRAC)?;
        let gr = q(&format!("gru{}_recurrent", k + 1), A_FRAC)?;
        let h = b.gru(t, hs[k], &gi, &gr);
        let g = b.glu(h, &q(&format!("glu{}", k + 1), A_FRAC)?);
        t = b.concat(vec![t, g]);
        let c = b.dense_tanh(t, &q(&format!("conv_dense{}", k + 1), A_FRAC)?, A_FRAC);
        let cc = b.concat(vec![cs[k], c]);
        let cv = b.dense_tanh(cc, &q(&format!("conv{}", k + 1), A_FRAC)?, A_FRAC);
        t = b.concat(vec![t, cv]);
        new_h.push(h);
        new_c.push(c);
    }
    let out = q("output", A_FRAC)?;
    let (w, bias) = out.rows(0..80);
    let x = b.conv(t, w, bias, A_FRAC + W_FRAC - IN_FRAC, TensorRec::clamp(80, -IN_MAX - 1, IN_MAX));
    let mut outs = vec![x];
    outs.extend(new_h);
    outs.extend(new_c);
    let step = b.finish(outs);
    Ok(Lowered { init, step })
}

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

    #[test]
    fn quantization_is_exact_from_the_bits() {
        let cases: [(f32, i32, i64); 9] = [
            (1.0, 24, 1 << 24),
            (-0.5, 1, -1),
            (0.75, 1, 2),     // 1.5 → 2 (half to even)
            (0.25, 1, 0),     // 0.5 → 0 (half to even)
            (1.25, 1, 2),     // 2.5 → 2
            (-1.25, 1, -2),   // −2.5 → −2
            (3.0e-8, 24, 1),  // 0.503 → 1
            (1.0e-30, 24, 0),
            (-36.71875, 24, -616_038_400),
        ];
        for (v, f, want) in cases {
            assert_eq!(quantize(v.to_bits(), f), Some(want), "{v} at 2^{f}");
        }
        assert_eq!(quantize(f32::NAN.to_bits(), 24), None);
        assert_eq!(quantize(f32::INFINITY.to_bits(), 24), None);
        // A subnormal.
        assert_eq!(quantize(1, 149), Some(1));
    }

    #[test]
    fn the_tables_are_tanh_and_the_sigmoid() {
        let (t, s) = (tanh_table(), sigmoid_table());
        assert_eq!((t[0], s[0]), (0, 0));
        assert_eq!((t[32768], s[32768]), (0, 16384));
        assert_eq!((t[32768 + 32767], t[1]), (32767, -32767));
        assert_eq!((s[32768 + 32767], s[1]), (32767, 0));
        // Odd and complementary.
        for v in 1..32768usize {
            assert_eq!(t[32768 + v], -t[32768 - v]);
            assert!((s[32768 + v] as i32 + s[32768 - v] as i32 - 32768).abs() <= 1, "{v}");
        }
        // Within one step of the functions they stand for.
        for v in (-32767i32..=32767).step_by(97) {
            let x = v as f64 / 4096.0;
            assert!((t[(v + 32768) as usize] as f64 - (x.tanh() * 32768.0).clamp(-32767.0, 32767.0)).abs() <= 0.5 + 1e-6, "tanh {v}");
            let x = v as f64 / 2048.0;
            let want = (32768.0 / (1.0 + (-x).exp())).min(32767.0);
            assert!((s[(v + 32768) as usize] as f64 - want).abs() <= 0.5 + 1e-6, "sigmoid {v}");
        }
        // Monotone.
        assert!(t[1..].windows(2).all(|w| w[0] <= w[1]) && s[1..].windows(2).all(|w| w[0] <= w[1]));
    }

    #[test]
    fn exp_is_close_to_the_function() {
        for u in [0u32, 1, 2, 1000, 2048, 20000, 32767] {
            let want = (-(u as f64) / 2048.0).exp() * 2f64.powi(62);
            let got = exp_q62(u) as f64;
            assert!((got - want).abs() <= want * 1e-14 + 64.0, "{u}");
        }
    }
}
