//! Quantização dinâmica int8 por canal de saída (equivalente Rust do
//! `quanta/torchao_quant.py` — Teo 28.1: camadas linearizáveis com cota
//! h/√12). EXCLUSÕES estruturais (T10(iii)): codebooks EMA, cabeças de
//! conceitos, embeddings de nível — mantidas em f32.

use candle_core::{Device, Result, Tensor};

pub struct PesoInt8 {
    pub q: Tensor,     // i32 com valores em [-127,127] (cota int8)
    pub escala: Tensor, // f32 (d_saida, 1)
    pub forma: Vec<usize>,
}

/// Quantiza W (d_saida, d_entrada) por LINHA (canal de saída).
pub fn quantizar_linha(w: &Tensor) -> Result<PesoInt8> {
    let dev = w.device();
    let dados = w.flatten_all()?.to_vec1::<f32>()?;
    let (s, e) = (w.dim(0)?, w.dim(1)?);
    let mut q = vec![0i32; dados.len()];
    let mut escala = vec![0f32; s];
    for i in 0..s {
        let linha = &dados[i * e..(i + 1) * e];
        let amax = linha.iter().fold(0f32, |m, x| m.max(x.abs()));
        let esc = (amax / 127.0).max(1e-12);
        escala[i] = esc;
        for j in 0..e {
            let v = (linha[j] / esc).round().clamp(-128.0, 127.0);
            q[i * e + j] = v as i32;
        }
    }
    Ok(PesoInt8 {
        q: Tensor::from_vec(q, (s, e), dev)?,
        escala: Tensor::from_vec(escala, (s, 1), dev)?,
        forma: vec![s, e],
    })
}

/// y = x·Wᵀ desquantizado on-the-fly (por canal).
pub fn linear_int8(x: &Tensor, p: &PesoInt8, dev: &Device) -> Result<Tensor> {
    let qf = p.q.to_dtype(candle_core::DType::F32)?;
    let w = qf.broadcast_mul(&p.escala)?; // (s, e)
    x.matmul(&w.transpose(0, 1)?)
}

/// Paridade int8 vs f32 (V14-análogo do Teo 28.1): erro relativo médio.
pub fn paridade(w: &Tensor, x: &Tensor, dev: &Device) -> Result<f64> {
    let p = quantizar_linha(w)?;
    let y_q = linear_int8(x, &p, dev)?;
    let y_f = x.matmul(&w.transpose(0, 1)?)?;
    let dif = (&y_q - &y_f)?.abs()?.mean_all()?.to_scalar::<f32>()? as f64;
    let ref_ = y_f.abs()?.mean_all()?.to_scalar::<f32>()? as f64;
    Ok(dif / ref_.max(1e-9))
}

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

    #[test]
    fn paridade_int8_dentro_da_cota() -> Result<()> {
        let dev = Device::Cpu;
        let w = Tensor::randn(0f32, 0.05, (64, 32), &dev)?;
        let x = Tensor::randn(0f32, 1f32, (4, 32), &dev)?;
        let erro = paridade(&w, &x, &dev)?;
        // cota empírica do Teo 28.1 (h/√12 por canal, |w| ≤ 0.05·3σ)
        assert!(erro < 0.05, "erro relativo int8 = {erro}");
        Ok(())
    }
}