//! FP8 E4M3FN encode/decode and the per-output-row FP8 quantiser used for the MTP draft head (quantised at load from its bf16 checkpoint).
//! The served checkpoint itself is stored pre-quantised (NVFP4 experts, FP8 linears and tied table) and is uploaded as is.

pub const FP8_MAX: f32 = 448.0;

pub fn fp8_e4m3_to_f32(b: u8) -> f32 {
    let s = if b & 0x80 != 0 { -1.0 } else { 1.0 };
    let e = ((b >> 3) & 0xF) as i32;
    let m = (b & 7) as f32;
    if e == 0xF && (b & 7) == 7 {
        return f32::NAN;
    }
    if e == 0 {
        s * m / 8.0 * 2f32.powi(-6)
    } else {
        s * (1.0 + m / 8.0) * 2f32.powi(e - 7)
    }
}

/// round-to-nearest-even encode (torch `.to(float8_e4m3fn)` semantics), saturating at +-448
pub fn f32_to_fp8_e4m3(x: f32) -> u8 {
    if x.is_nan() {
        return 0x7F;
    }
    let sign = if x.is_sign_negative() { 0x80u8 } else { 0 };
    let a = x.abs().min(FP8_MAX);
    let code = if a < 0.015625 {
        // subnormal range (step 2^-9); 8 steps reaches the smallest normal 0x08 correctly
        (a * 512.0).round_ties_even() as u8
    } else {
        let bits = a.to_bits();
        let mut e = ((bits >> 23) & 0xff) as i32 - 127;
        let m = bits & 0x7f_ffff;
        let mut q = m >> 20;
        let rem = m & 0xf_ffff;
        if rem > 0x8_0000 || (rem == 0x8_0000 && q & 1 == 1) {
            q += 1;
        }
        if q == 8 {
            q = 0;
            e += 1;
        }
        (((e + 7) as u32) << 3 | q) as u8
    };
    sign | code.min(0x7E)
}

pub struct Fp8Rows {
    pub w: Vec<u8>,
    pub scales: Vec<f32>,
}

/// FP8-E4M3 per-output-row quantiser: scale = row_amax/448, w = fp8(x/scale). No calibration.
pub fn quantize_fp8_rows(w: &[f32], rows: usize, cols: usize) -> Fp8Rows {
    assert_eq!(w.len(), rows * cols);
    let mut q = vec![0u8; rows * cols];
    let mut scales = vec![0f32; rows];
    for r in 0..rows {
        let row = &w[r * cols..(r + 1) * cols];
        let amax = row.iter().fold(0.0f32, |m, x| m.max(x.abs())).max(1e-12);
        let s = amax / FP8_MAX;
        scales[r] = s;
        for (o, &x) in q[r * cols..(r + 1) * cols].iter_mut().zip(row) {
            *o = f32_to_fp8_e4m3(x / s);
        }
    }
    Fp8Rows { w: q, scales }
}

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

    #[test]
    fn fp8_known_values() {
        assert_eq!(fp8_e4m3_to_f32(0x00), 0.0);
        assert_eq!(fp8_e4m3_to_f32(0x38), 1.0); // e=7,m=0
        assert_eq!(fp8_e4m3_to_f32(0x7E), 448.0);
        assert_eq!(fp8_e4m3_to_f32(0xB8), -1.0);
        assert_eq!(fp8_e4m3_to_f32(0x01), 2f32.powi(-9)); // smallest subnormal
        assert!(fp8_e4m3_to_f32(0x7F).is_nan());
        for b in 0u8..0x7F {
            assert_eq!(f32_to_fp8_e4m3(fp8_e4m3_to_f32(b)), b, "roundtrip {b:#x}");
        }
        assert_eq!(f32_to_fp8_e4m3(1e9), 0x7E);
    }

    #[test]
    fn fp8_encode_matches_bruteforce_rne() {
        // brute-force nearest with ties-to-even-code
        let brute = |x: f32| -> u8 {
            let a = x.abs().min(FP8_MAX);
            let (mut best, mut bd) = (0u8, f32::INFINITY);
            for c in 0u8..0x7F {
                let d = (fp8_e4m3_to_f32(c) - a).abs();
                if d < bd || (d == bd && c % 2 == 0 && best % 2 == 1) {
                    bd = d;
                    best = c;
                }
            }
            (if x.is_sign_negative() { 0x80 } else { 0 }) | best
        };
        let mut s = 1u64;
        for i in 0..20000 {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            let x = ((s >> 33) as f32 / (1u64 << 31) as f32 - 0.5) * if i % 3 == 0 { 600.0 } else if i % 3 == 1 { 4.0 } else { 0.05 };
            assert_eq!(f32_to_fp8_e4m3(x), brute(x), "x={x}");
        }
        for c in 0u8..0x7E {
            let m = (fp8_e4m3_to_f32(c) + fp8_e4m3_to_f32(c + 1)) / 2.0; // exact tie
            assert_eq!(f32_to_fp8_e4m3(m), brute(m), "tie {m}");
        }
    }
}