Blame for src/quantizador.rs

//! Quantizador GSQ por nível (Teos A.1/A.4/D.4) — o coração da ponte
//! NCM×KHTST: codebook Stiefel com seleção Gumbel-Softmax (gradiente exato,
//! sem STE) no treino e NN duro (argmax) em eval.
//!
//! Dois modos (r-01 §4.4): "gsq_rco" (padrão, Teo D.4 — relaxamento
//! diferenciável com τ-annealing) e "ste_ema" (Teo A.0 — STE + commitment +
//! codebook por EMA). A escala R_ℓ é mantida por EMA (Cor. A.1.1).

use candle_core::{CpuStorage, CustomOp3, Device, Layout, Result, Shape, Tensor, Var, D};
use crate::stiefel;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ModoQuant {
    GsqRco,
    SteEma,
}

impl ModoQuant {
    pub fn de_str(s: &str) -> ModoQuant {
        if s == "ste_ema" { ModoQuant::SteEma } else { ModoQuant::GsqRco }
    }
}

/// CustomOp3: (z, C, ruído_gumbel) → r_suave = R·softmax((zCᵀ/τ + g))·C.
/// O ruído entra como tensor EXPLÍCITO (constante no grafo) para que o
/// backward recompute exatamente o mesmo y do forward (gradiente exato do
/// relaxamento — Teo A.4(iii), sem re-amostragem no adjunto).
pub struct GsqSuave {
    pub r: f64,
    pub tau: f64,
}

fn softmax_ultima_dim(v: &[f32], n: usize, k: usize) -> Vec<f32> {
    let mut out = vec![0f32; v.len()];
    for i in 0..n {
        let linha = &v[i * k..(i + 1) * k];
        let mx = linha.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let mut soma = 0f32;
        for (j, &x) in linha.iter().enumerate() {
            let e = (x - mx).exp();
            out[i * k + j] = e;
            soma += e;
        }
        for j in 0..k {
            out[i * k + j] /= soma.max(1e-20);
        }
    }
    out
}

impl CustomOp3 for GsqSuave {
    fn name(&self) -> &'static str { "gsq_suave_stiefel" }

    fn cpu_fwd(
        &self,
        s1: &CpuStorage, l1: &Layout,
        _s2: &CpuStorage, _l2: &Layout,
        s3: &CpuStorage, l3: &Layout,
    ) -> Result<(CpuStorage, Shape)> {
        if !(l1.is_contiguous() && _l2.is_contiguous() && l3.is_contiguous()) {
            return Err(candle_core::Error::Msg(
                "gsq_suave: tensors de entrada devem ser contíguos (chame .contiguous())".into(),
            ));
        }
        let z = s1.as_slice::<f32>()?;
        let (n, d) = (l1.shape().dims()[0], l1.shape().dims()[1]);
        let c = _s2.as_slice::<f32>()?;
        let (k, _dc) = (_l2.shape().dims()[0], _l2.shape().dims()[1]);
        let ruido = s3.as_slice::<f32>()?;
        // v0.8.0 (bug v080-2, pré-existente): a fórmula anterior
        // `z.len()/d*k*d/d*k` avaliava para n·k²·(…) sem sentido (320 ≠ 40
        // no teste) e disparava em TODO build debug — os gates da v0.7.0 só
        // passavam em release (debug_assert compilado fora). Correta: n·k.
        debug_assert_eq!(ruido.len(), n * k); // (n, k)
        // logits = z·Cᵀ/τ + g
        let mut logits = vec![0f32; n * k];
        for i in 0..n {
            for j in 0..k {
                let mut acc = 0f32;
                for t in 0..d {
                    acc += z[i * d + t] * c[j * d + t];
                }
                logits[i * k + j] = acc as f32 / self.tau as f32 + ruido[i * k + j];
            }
        }
        let y = softmax_ultima_dim(&logits, n, k);
        // r = R · y·C
        let mut r = vec![0f32; n * d];
        for i in 0..n {
            for t in 0..d {
                let mut acc = 0f32;
                for j in 0..k {
                    acc += y[i * k + j] * c[j * d + t];
                }
                r[i * d + t] = (self.r as f32) * acc;
            }
        }
        Ok((CpuStorage::F32(r), (n, d).into()))
    }

    fn bwd(
        &self,
        arg1: &Tensor, arg2: &Tensor, arg3: &Tensor,
        _res: &Tensor, grad_res: &Tensor,
    ) -> Result<(Option<Tensor>, Option<Tensor>, Option<Tensor>)> {
        // recompute y (mesmo ruído — arg3): y = softmax(zCᵀ/τ + g)
        let logits = (&arg1.matmul(&arg2.transpose(0, 1)?)?.affine(1.0 / self.tau, 0.0)? + arg3)?;
        let y = candle_nn::ops::softmax_last_dim(&logits)?;
        // gy = R·(g_·Cᵀ)
        let gy = (&grad_res.matmul(&arg2.transpose(0, 1)?)? * self.r)?;
        // gl = y ⊙ (gy − Σ(y⊙gy)) — media como coluna p/ broadcast (n,1)
        let media = (&y * &gy)?.sum(D::Minus1)?;
        let media_c = media.reshape((media.dims()[0], 1))?;
        let gl = (&y * &gy.broadcast_sub(&media_c.broadcast_div(&ones_like(&media_c)?)?)?)?;
        // ∂L/∂z = (1/τ)·gl·C ; ∂L/∂C = R·yᵀ·g_ ; ruído constante
        let gz = gl.matmul(arg2)?.affine(1.0 / self.tau, 0.0)?;
        let gc = (&y.transpose(0, 1)?.matmul(grad_res)? * self.r)?;
        Ok((Some(gz), Some(gc), None))
    }
}

fn ones_like(t: &Tensor) -> Result<Tensor> {
    Tensor::ones(t.shape(), t.dtype(), t.device())
}

/// Quantizador de um nível ℓ do VQ residual.
pub struct QuantizadorNivel {
    pub w: Var,             // (K, d) bruto — codebook = R(Wᵀ)ᵀ
    pub r_escala: f64,      // R_ℓ (EMA)
    pub tau: f64,           // τ corrente (agenda analítica)
    pub modo: ModoQuant,
    pub ema: f64,           // fator EMA (0,99)
    pub escala_modo: String, // "ema_projecao" | "otima_teo"
    pub iters_polar: usize,
    pub delta_ort_ultimo: f64,
    pub uso: Vec<f64>,      // contagem de uso por código (telemetria/colapso)
    pub semente: u64,       // RNG dedicado do nível (pareamento — Teo 31.7)
    chamadas: std::cell::Cell<u64>, // contador p/ ruído fresco por chamada
}

impl QuantizadorNivel {
    pub fn novo(k: usize, d: usize, modo: ModoQuant, ema: f64, escala_modo: &str, dev: &Device) -> Result<Self> {
        let w0 = stiefel::init_stiefel(k, d, dev)?;
        Ok(QuantizadorNivel {
            w: Var::from_tensor(&w0)?,
            r_escala: 1.0 / (k as f64).sqrt(),
            tau: 1.0,
            modo,
            ema,
            escala_modo: escala_modo.to_string(),
            iters_polar: 30,
            delta_ort_ultimo: 0.0,
            uso: vec![0.0; k],
            semente: 0xC0FFEE + k as u64 * 7919 + d as u64,
            chamadas: std::cell::Cell::new(0),
        })
    }

    /// Codebook C = R(Wᵀ)ᵀ com linhas ortonormais (diferenciável — Teo 26.3).
    /// SEMPRE contíguo (index_select/arquivo exigem layout padrão).
    pub fn codebook(&self) -> Result<Tensor> {
        let c = stiefel::codebook_stiefel(self.w.as_tensor(), self.iters_polar)?;
        c.contiguous()
    }

    /// δ_ort do codebook corrente (V9).
    pub fn delta_ort(&self) -> Result<f64> {
        let c = self.codebook()?.detach();
        stiefel::erro_ortogonalidade(&c.transpose(0, 1)?)
    }

    /// Amostra ruído de Gumbel (n, K) — fora do grafo (constante).
    pub fn ruido_gumbel(&self, n: usize, k: usize, dev: &Device) -> Result<Tensor> {
        use rand::prelude::*;
        let mut rng = StdRng::seed_from_u64(self.semente_ruido());
        let mut v = Vec::with_capacity(n * k);
        for _ in 0..n * k {
            let u: f32 = rng.gen::<f32>().clamp(1e-9, 1.0 - 1e-9);
            v.push((-(-u.ln()).ln()) as f32);
        }
        Tensor::from_vec(v, (n, k), dev)
    }

    fn semente_ruido(&self) -> u64 {
        let c = self.chamadas.get();
        self.chamadas.set(c + 1);
        self.semente.wrapping_add(c.wrapping_mul(0x9E3779B97F4A7C15))
    }

    /// Quantização de TREINO (…, d) → (…, d):
    /// - gsq_rco: relaxamento suave diferenciável (gradiente exato);
    /// - ste_ema: STE (r = z + sg[r_duro − z]).
    pub fn quantizar_treino(&self, z: &Tensor) -> Result<Tensor> {
        let c = self.codebook()?.contiguous()?;
        match self.modo {
            ModoQuant::GsqRco => {
                let z2 = contiguo_2d(z)?;
                let (n, _d) = (z2.dims()[0], z2.dims()[1]);
                let k = c.dims()[0];
                let ruido = self.ruido_gumbel(n, k, z.device())?;
                let r = Tensor::apply_op3(&z2, &c, &ruido, GsqSuave { r: self.r_escala, tau: self.tau.max(1e-3) })?;
                redimensionar_de_volta(&r, z)
            }
            ModoQuant::SteEma => {
                let r_duro = self.valores_duros_int(z, &self.codebook()?.contiguous()?)?.detach();
                // STE: z + (r_duro − z) com r_duro constante no grafo
                Ok(z.add(&r_duro.sub(z)?)?)
            }
        }
    }

    fn valores_duros_int(&self, z: &Tensor, c: &Tensor) -> Result<Tensor> {
        // c deve vir contíguo do caller (custom op exige layout padrão)
        // ids = argmax ⟨z, u_k⟩ (Teo A.4(i): uma GEMM + top-1)
        let z2 = contiguo_2d(z)?;
        let logits = z2.matmul(&c.transpose(0, 1)?)?;
        let ids = logits.argmax(D::Minus1)?;
        let linhas = c.index_select(&ids, 0)?;
        let r = linhas.affine(self.r_escala, 0.0)?;
        redimensionar_de_volta(&r, z)
    }

    /// Quantização DURA (eval/telemetria): (ids, valores R·c_id).
    pub fn quantizar_duro(&self, z: &Tensor) -> Result<(Tensor, Tensor)> {
        let c = self.codebook()?.detach();
        let z2 = contiguo_2d(z)?;
        let logits = z2.matmul(&c.transpose(0, 1)?)?;
        let ids = logits.argmax(D::Minus1)?;
        let valores = self.valores_duros_int(z, &c)?;
        Ok((ids, valores))
    }

    /// Materializa códigos ids → vetores R·c_id (para reconstrução na amostragem).
    pub fn materializar(&self, ids: &Tensor) -> Result<Tensor> {
        let c = self.codebook()?.detach();
        let linhas = c.index_select(&ids, 0)?;
        Ok(linhas.affine(self.r_escala, 0.0)?)
    }

    /// Atualiza R_ℓ por EMA da escala de projeção (Cor. A.1.1):
    /// "ema_projecao": R ← EMA(E[‖P_C z‖]); "otima_teo": R ← EMA(E[‖P_C z‖/√K]).
    pub fn atualizar_escala(&mut self, z: &Tensor) -> Result<()> {
        let c = self.codebook()?.detach();
        let z2 = contiguo_2d(z)?;
        let cz = z2.matmul(&c.transpose(0, 1)?)?;
        let pz = cz.matmul(&c)?;
        let alvo = pz.sqr()?.sum(D::Minus1)?.sqrt()?.mean_all()?.to_scalar::<f32>()? as f64;
        let alvo = match self.escala_modo.as_str() {
            "otima_teo" => alvo / (self.uso.len() as f64).sqrt(),
            _ => alvo,
        };
        self.r_escala = self.ema * self.r_escala + (1.0 - self.ema) * alvo.max(1e-4);
        Ok(())
    }

    /// Registra uso de códigos (telemetria de colapso — R2 do n-04 §9).
    pub fn registrar_uso(&mut self, ids: &Tensor) -> Result<()> {
        let v = ids.flatten_all()?.to_vec1::<u32>()?;
        for id in v {
            if (id as usize) < self.uso.len() {
                self.uso[id as usize] += 1.0;
            }
        }
        Ok(())
    }

    /// Perplexidade de uso do codebook (utilização efetiva de K — telemetria V2/V6).
    pub fn perplexidade_uso(&self) -> f64 {
        let total: f64 = self.uso.iter().sum();
        if total <= 0.0 {
            return 0.0;
        }
        let h: f64 = self.uso.iter().map(|&u| {
            let p = u / total;
            if p > 0.0 { -p * p.ln() } else { 0.0 }
        }).sum();
        h.exp()
    }

    /// Agenda analítica de τ (Teo 30.6): decaimento de τ0 até τ_min em t_alvo passos.
    pub fn agendar_tau(&mut self, passo: usize, tau0: f64, tau_min: f64, t_alvo: usize) {
        let t = passo.min(t_alvo) as f64;
        let tau = tau_min + (tau0 - tau_min) * (1.0 - t / t_alvo.max(1) as f64);
        self.tau = tau.max(tau_min);
    }
}

// ---------- utilidades de forma ----------

/// Achata (…, d) → (n, d) e devolve também a forma original.
fn contiguo_2d(t: &Tensor) -> Result<Tensor> {
    let d = t.dim(D::Minus1)?;
    let n = t.elem_count() / d;
    t.reshape((n, d))?.contiguous()
}

fn redimensionar_de_volta(r: &Tensor, modelo: &Tensor) -> Result<Tensor> {
    r.reshape(modelo.shape())
}

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

    #[test]
    fn gsq_gradiente_vs_finitas_diferencas() -> Result<()> {
        // V-análogo do V8: gradiente exato do relaxamento vs. FD
        let dev = Device::Cpu;
        let mut q = QuantizadorNivel::novo(8, 6, ModoQuant::GsqRco, 0.99, "ema_projecao", &dev)?;
        q.tau = 1.0;
        q.r_escala = 0.7;
        let z = Tensor::randn(0f32, 1f32, (5, 6), &dev)?;
        let z2 = contiguo_2d(&z)?;
        let c = q.codebook()?.contiguous()?;
        let ruido = q.ruido_gumbel(5, 8, &dev)?;
        // valor da perda com pesos fixos
        let pesos = Tensor::randn(0f32, 1f32, (5, 6), &dev)?;
        let perda = |zz: &Tensor| -> Result<f64> {
            let r = Tensor::apply_op3(zz, &c, &ruido, GsqSuave { r: q.r_escala, tau: q.tau })?;
            let l = (&r * &pesos)?.sum_all()?;
            Ok(l.to_scalar::<f32>()? as f64)
        };
        // gradiente via op (aplicado a Var)
        let zv = Var::from_tensor(&z2)?;
        let r = Tensor::apply_op3(zv.as_tensor(), &c, &ruido, GsqSuave { r: q.r_escala, tau: q.tau })?;
        let l = (&r * &pesos)?.sum_all()?;
        let gs = l.backward()?;
        let g = gs.get(&zv).ok_or_else(|| candle_core::Error::Msg("sem gradiente".into()))?;
        let gv = g.flatten_all()?.to_vec1::<f32>()?;
        // FD central
        let eps = 1e-3f32;
        let mut erro_max = 0f64;
        for i in 0..z2.elem_count() {
            let mut zp = z2.flatten_all()?.to_vec1::<f32>()?;
            let mut zm = zp.clone();
            zp[i] += eps;
            zm[i] -= eps;
            let tp = Tensor::from_vec(zp, (5, 6), &dev)?;
            let tm = Tensor::from_vec(zm, (5, 6), &dev)?;
            let fd = (perda(&tp)? - perda(&tm)?) / (2.0 * eps as f64);
            erro_max = erro_max.max((fd - gv[i] as f64).abs());
        }
        assert!(erro_max < 5e-3, "erro máx FD vs exato = {erro_max}");
        Ok(())
    }

    #[test]
    fn perplexidade_uso_reflete_histograma() {
        let mut q = QuantizadorNivel {
            w: Var::from_tensor(&Tensor::zeros((4, 4), candle_core::DType::F32, &Device::Cpu).unwrap()).unwrap(),
            r_escala: 1.0, tau: 1.0, modo: ModoQuant::GsqRco, ema: 0.99,
            escala_modo: "ema_projecao".into(), iters_polar: 10, delta_ort_ultimo: 0.0,
            uso: vec![10.0, 10.0, 10.0, 10.0],
            semente: 7,
            chamadas: std::cell::Cell::new(0),
        };
        let p = q.perplexidade_uso();
        assert!((p - 4.0).abs() < 1e-9);
        q.uso = vec![40.0, 0.0, 0.0, 0.0];
        assert!(q.perplexidade_uso() < 1.01);
    }
}

#[cfg(test)]
mod testes_debug {
    use super::*;
    use candle_core::{Device, Var};

    /// Compara o adjunto do GsqSuave contra implementação nativa (ops candle).
    #[test]
    fn gsq_bwd_vs_nativo() -> Result<()> {
        let dev = Device::Cpu;
        let (n, d, k) = (5usize, 6, 8);
        let c0 = stiefel::init_stiefel(k, d, &dev)?.contiguous()?;
        let zv = Var::from_tensor(&Tensor::randn(0f32, 1f32, (n, d), &dev)?)?;
        let ruido = Tensor::randn(0f32, 1f32, (n, k), &dev)?;
        let pesos = Tensor::randn(0f32, 1f32, (n, d), &dev)?;
        let (r_op, tau) = (0.7f64, 1.3f64);
        // caminho NATIVO
        let logits_n = (zv.as_tensor().matmul(&c0.transpose(0, 1)?)?.affine(1.0 / tau, 0.0)? + &ruido)?;
        let y_n = candle_nn::ops::softmax(&logits_n, candle_core::D::Minus1)?;
        let r_n = y_n.matmul(&c0)?.affine(r_op, 0.0)?;
        let l_n = (&r_n * &pesos)?.sum_all()?;
        let gs_n = l_n.backward()?;
        let g_n = gs_n.get(&zv).unwrap().flatten_all()?.to_vec1::<f32>()?;
        // caminho OP
        let r_o = Tensor::apply_op3(zv.as_tensor(), &c0, &ruido, GsqSuave { r: r_op, tau })?;
        let l_o = (&r_o * &pesos)?.sum_all()?;
        let gs_o = l_o.backward()?;
        let g_o = gs_o.get(&zv).unwrap().flatten_all()?.to_vec1::<f32>()?;
        // forward deve ser idêntico
        let dif_fwd = (&r_n - &r_o)?.abs()?.max_all()?.to_scalar::<f32>()?;
        assert!(dif_fwd < 1e-5, "forward difere: {dif_fwd}");
        let mut erro = 0f64;
        for i in 0..g_n.len() {
            erro = erro.max((g_n[i] - g_o[i]).abs() as f64);
        }
        assert!(erro < 1e-4, "grad nativo vs op = {erro}\n n={g_n:?}\n o={g_o:?}");
        Ok(())
    }
}