//! The draft's decoder network in floating point (draft-ietf-mlcodec-opus-dred-07 §2.2.2), the
//! reference the integer profile is measured against (extensions/opus-dred.md §6).
//!
//! It computes in `f64` from the float32 weights, with `tanh` and the sigmoid formed from an
//! exponential written here out of IEEE-754 basic operations in a fixed order, so that it gives
//! the same bits on every machine: the corpus's negative control that presents a floating-point
//! decode as exact is then byte-stable too. The draft's reference implementation computes in
//! `float32` with a rational approximation of `tanh`; §6 measures against both.

use std::collections::BTreeMap;

use super::payload::{self, Payload};
use super::weights::Layer;
use super::{Row, FEATURES, FRAMES_PER_LATENT};

/// `ln 2` in two parts: the high part has its low 20 bits zero, so `k·LN2_HI` is exact for
/// the `k` here, and the low part holds the rest.
const LN2_HI: f64 = 0.693_147_180_369_123_8;
const LN2_LO: f64 = 1.908_214_929_270_587_7e-10;

/// `e^(−y)` for `y ≥ 0`: `y = k·ln 2 + f` with `0 ≤ f < ln 2`, `e^(−f)` by its Taylor series
/// to `f^25`, then the exact scaling by `2^(−k)`. Basic operations only, in a fixed order.
pub fn exp_neg(y: f64) -> f64 {
    if y > 700.0 {
        return 0.0;
    }
    let k = (y / LN2_HI).floor();
    let f = (y - k * LN2_HI) - k * LN2_LO;
    let (mut term, mut sum) = (1.0, 1.0);
    for i in 1..=25 {
        term = term * (-f) / i as f64;
        sum += term;
    }
    sum * f64::from_bits(((1023 - k as i64) as u64) << 52)
}

pub fn tanh(x: f64) -> f64 {
    let e = exp_neg(2.0 * x.abs());
    let t = (1.0 - e) / (1.0 + e);
    if x < 0.0 { -t } else { t }
}

pub fn sigmoid(x: f64) -> f64 {
    let e = exp_neg(x.abs());
    if x >= 0.0 { 1.0 / (1.0 + e) } else { e / (1.0 + e) }
}

struct FLayer {
    inputs: usize,
    w: Vec<f64>,
    b: Vec<f64>,
    /// The reference implementation's 8-bit arithmetic, when emulated: each input rounded to
    /// `⌊0.5 + 127·x⌋`, the 8-bit weights, and each output scaled by its scale.
    q: Option<(Vec<f64>, Vec<f64>)>,
}

impl FLayer {
    fn of(l: &Layer) -> FLayer {
        let f = |v: &u32| f32::from_bits(*v) as f64;
        FLayer { inputs: l.inputs, w: l.w.iter().map(f).collect(), b: l.b.iter().map(f).collect(), q: None }
    }
    fn affine(&self, x: &[f64]) -> Vec<f64> {
        if let Some((w8, scale)) = &self.q {
            let xq: Vec<f64> = x.iter().map(|v| (0.5 + 127.0 * v).floor()).collect();
            return (0..self.b.len())
                .map(|o| {
                    let row = &w8[o * self.inputs..(o + 1) * self.inputs];
                    let mut s = 0.0;
                    for (w, v) in row.iter().zip(&xq) {
                        s += w * v;
                    }
                    s * scale[o] + self.b[o]
                })
                .collect();
        }
        (0..self.b.len())
            .map(|o| {
                let row = &self.w[o * self.inputs..(o + 1) * self.inputs];
                let mut s = 0.0;
                for (w, v) in row.iter().zip(x) {
                    s += w * v;
                }
                s + self.b[o]
            })
            .collect()
    }
}

/// The floating-point decoder over the layers of `weights::decoder_layers`.
pub struct FloatModel {
    l: BTreeMap<&'static str, FLayer>,
}

impl FloatModel {
    pub fn new(layers: &BTreeMap<&'static str, Layer>) -> FloatModel {
        FloatModel { l: layers.iter().map(|(k, v)| (*k, FLayer::of(v))).collect() }
    }

    /// The network as the draft's reference implementation computes it when it reads the 8-bit
    /// weights the weights file also holds (every layer but `dense1` and `hidden_init`): the
    /// arithmetic the draft's test vectors were made with, as §6 finds. Measurement only.
    pub fn int8_reference(layers: &BTreeMap<&'static str, Layer>, arrays: &BTreeMap<String, super::weights::Array>) -> FloatModel {
        let mut m = FloatModel::new(layers);
        for &(name, inputs, outputs, sparse) in super::weights::LAYERS.iter().chain(std::iter::once(&super::weights::OUTPUT)) {
            if let Some((w8, scale)) = super::weights::int8_layer(arrays, name, inputs, outputs, sparse) {
                let l = m.l.get_mut(name).unwrap();
                l.q = Some((w8.iter().map(|&v| v as f64).collect(), scale.iter().map(|&b| f32::from_bits(b) as f64).collect()));
            }
        }
        m
    }

    fn dense_tanh(&self, name: &str, x: &[f64]) -> Vec<f64> {
        self.l[name].affine(x).into_iter().map(tanh).collect()
    }

    fn gru(&self, k: usize, x: &[f64], h: &[f64]) -> Vec<f64> {
        let gi = self.l[format!("gru{k}_input").as_str()].affine(x);
        let gr = self.l[format!("gru{k}_recurrent").as_str()].affine(h);
        (0..64)
            .map(|i| {
                let z = sigmoid(gi[i] + gr[i]);
                let r = sigmoid(gi[64 + i] + gr[64 + i]);
                let hc = tanh(gi[128 + i] + r * gr[128 + i]);
                z * h[i] + (1.0 - z) * hc
            })
            .collect()
    }

    /// The features of the first `n` latents of `p`, newest first, as `f64`.
    pub fn decode(&self, p: &Payload, n: usize) -> Vec<[f64; FEATURES]> {
        let q16 = |v: i64| v as f64 / 65536.0;
        let is: Vec<f64> = payload::state_input(p).into_iter().map(q16).collect();
        let s0 = self.dense_tanh("gru_init", &self.dense_tanh("hidden_init", &is));
        let mut h: Vec<Vec<f64>> = (0..5).map(|k| s0[64 * k..64 * (k + 1)].to_vec()).collect();
        let mut mem: Vec<Vec<f64>> = vec![vec![0.0; 32]; 5];
        let mut rows = Vec::new();
        for l in 0..n.min(p.latents.len()) {
            let z: Vec<f64> = payload::latent_input(p, l).into_iter().map(q16).collect();
            let mut t = self.dense_tanh("dense1", &z);
            for k in 0..5 {
                h[k] = self.gru(k + 1, &t, &h[k]);
                let g = self.l[format!("glu{}", k + 1).as_str()].affine(&h[k]);
                t.extend(h[k].iter().zip(&g).map(|(y, a)| y * sigmoid(*a)));
                let c = self.dense_tanh(&format!("conv_dense{}", k + 1), &t);
                let mut cc = mem[k].clone();
                cc.extend_from_slice(&c);
                t.extend(self.dense_tanh(&format!("conv{}", k + 1), &cc));
                mem[k] = c;
            }
            let x = self.l["output"].affine(&t);
            for j in 0..FRAMES_PER_LATENT {
                let mut r = [0f64; FEATURES];
                r.copy_from_slice(&x[j * FEATURES..(j + 1) * FEATURES]);
                rows.push(r);
            }
        }
        rows
    }
}

/// Float features at `2^16`, rounded half to even: the rows a floating-point decode would claim.
pub fn to_q16(rows: &[[f64; FEATURES]]) -> Vec<Row> {
    rows.iter()
        .map(|r| {
            let mut o = [0i32; FEATURES];
            for (d, v) in o.iter_mut().zip(r) {
                *d = (v * 65536.0).round_ties_even().clamp(i32::MIN as f64, i32::MAX as f64) as i32;
            }
            o
        })
        .collect()
}

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

    #[test]
    fn the_exponential_is_accurate() {
        for y in [0.0, 1e-9, 0.01, 0.5, 0.69, 0.7, 1.0, 2.0, 7.3, 16.0, 40.0, 699.0] {
            let (got, want) = (exp_neg(y), (-y).exp());
            assert!((got - want).abs() <= want * 1e-15, "{y}: {got} {want}");
        }
        for x in [-9.0, -1.0, -0.001, 0.0, 0.3, 4.0] {
            assert!((tanh(x) - x.tanh()).abs() < 1e-14);
            assert!((sigmoid(x) - 1.0 / (1.0 + (-x).exp())).abs() < 1e-14);
        }
    }
}
