Blame for src/tokenizador.rs

//! Tokenizador de conceitos NCM (Estágio 1 / fase A) — port de
//! `conceitos/tokenizador_conceitos.py` (n-04 §2.3) com a síntese GSQ+RCO
//! (Teo D.4): codebooks Stiefel por nível, seleção GSQ (gradiente exato),
//! nested dropout (Teo A.2), commitment (Teo A.0) e reconstrução por CE
//! token-level no vocabulário.
//!
//! Recursão residual (n-02 §0):
//!   z_1 = Interp_{n_1}(E(x));  r_ℓ = Quant_{C_ℓ}(z_ℓ);
//!   z_{ℓ+1} = Interp_{n_{ℓ+1}}(z_ℓ) − Interp_{n_{ℓ+1}}(r_ℓ);
//!   ẑ_ℓ = Σ_{j≤ℓ} Interp_{n_ℓ}(r_j)   (reconstrução cumulativa)

use candle_core::{Device, Result, Tensor, Var, D};
use crate::atencao::BlocoAdba;
use crate::camadas::{Embedding, Linear, RmsNorm};
use crate::config::ConfigKhtst;
use crate::mascaras::{interp_para, matriz_interp};
use crate::quantizador::{ModoQuant, QuantizadorNivel};
use rand::prelude::*;

pub struct TokenizadorConceitos {
    // Encoder (treino apenas; atenção plena)
    pub enc_emb: Embedding,
    pub enc_blocos: Vec<BlocoAdba>,
    pub enc_norm: RmsNorm,
    // Projeções por nível d_modelo → d_lat + quantizadores Stiefel/GSQ
    pub proj_niveis: Vec<Linear>,
    pub quants: Vec<QuantizadorNivel>,
    // Decoder AR condicionado em ẑ (B.1: p_ψ(x|c) = Π_i p(x_i|x_{<i}, ẑ))
    pub dec_proj: Linear,
    pub dec_emb: Embedding,
    pub dec_blocos: Vec<BlocoAdba>,
    pub dec_norm: RmsNorm,
    pub dec_head: Linear,
    // máscaras pré-computadas
    pub mascara_causal_dec: Tensor,
    pub niveis: Vec<usize>,
    pub d_lat: usize,
    pub vocab: usize,
    pub cfg_ncm_beta_commit: f64,
    pub modo: ModoQuant,
}

/// Saída da codificação: códigos por nível + reconstruções cumulativas.
pub struct Codigos {
    pub ids: Vec<Tensor>,     // u32 (b, n_ℓ)
    pub zhats: Vec<Tensor>,   // (b, n_ℓ, d_lat) — ẑ_ℓ na granularidade do nível
}

impl TokenizadorConceitos {
    pub fn novo(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> {
        let m = &cfg.modelo;
        let n = &cfg.ncm;
        let l = n.niveis.len();
        let n_total_dec = m.comprimento_ctx;
        let mascara_causal_dec = causal_aditiva(n_total_dec, dev)?;
        let mut enc_blocos = Vec::new();
        let mut dec_blocos = Vec::new();
        for _ in 0..2 {
            enc_blocos.push(BlocoAdba::novo(m.d_modelo, m.n_cabecas, m.d_ff, l, false, false, false, dev)?);
            dec_blocos.push(BlocoAdba::novo(m.d_modelo, m.n_cabecas, m.d_ff, l, false, false, false, dev)?);
        }
        let modo = ModoQuant::de_str(&n.modo_quantizacao);
        let mut quants = Vec::with_capacity(l);
        for _ in 0..l {
            quants.push(QuantizadorNivel::novo(n.k_por_nivel, n.d_lat, modo, n.ema_codebook, &n.escala_r, dev)?);
        }
        let mut proj = Vec::with_capacity(l);
        for _ in 0..l {
            proj.push(Linear::sem_vies(m.d_modelo, n.d_lat, 0.02, dev)?);
        }
        Ok(TokenizadorConceitos {
            enc_emb: Embedding::novo(m.vocab, m.d_modelo, dev)?,
            enc_blocos,
            enc_norm: RmsNorm::novo(m.d_modelo, 1e-5, dev)?,
            proj_niveis: proj,
            quants,
            dec_proj: Linear::sem_vies(n.d_lat, m.d_modelo, 0.02, dev)?,
            dec_emb: Embedding::novo(m.vocab, m.d_modelo, dev)?,
            dec_blocos,
            dec_norm: RmsNorm::novo(m.d_modelo, 1e-5, dev)?,
            dec_head: Linear::sem_vies(m.d_modelo, m.vocab, 0.02, dev)?,
            mascara_causal_dec,
            niveis: n.niveis.clone(),
            d_lat: n.d_lat,
            vocab: m.vocab,
            cfg_ncm_beta_commit: n.beta_commit,
            modo,
        })
    }

    /// Encoder: x (b, N) → z (b, N, d_modelo). Atenção plena (só treino).
    pub fn encoder(&self, ids: &Tensor) -> Result<Tensor> {
        let mut x = self.enc_emb.fwd(ids)?;
        for bloco in &self.enc_blocos {
            let mut cache = Default::default();
            x = bloco.fwd(&x, None, &mut cache, None, None)?;
        }
        self.enc_norm.fwd(&x)
    }

    /// Codificação de TREINO (fase A): relaxamento GSQ suave diferenciável.
    /// Devolve (zhats_suaves, zs_por_nivel) para a perda e telemetria.
    fn codificar_treino(&mut self, ids: &Tensor) -> Result<(Vec<Tensor>, Vec<Tensor>)> {
        let dev = ids.device();
        let z = self.encoder(ids)?;
        let mut zs = Vec::with_capacity(self.niveis.len());
        let mut zhats = Vec::with_capacity(self.niveis.len());
        // projeção ÚNICA d_modelo → d_lat na granularidade do nível 1 (n-02 §0);
        // a recursão residual vive INTEIRAMENTE em d_lat
        let n1 = self.niveis[0];
        let mut z_corrente = self.proj_niveis[0].fwd(&interp_para(&z, n1, dev)?)?;
        let mut zhat_acum: Option<Tensor> = None;
        for (i, &n_l) in self.niveis.iter().enumerate() {
            let zp = z_corrente.clone();
            let r_suave = self.quants[i].quantizar_treino(&zp)?;
            // telemetria + escalas + uso (sem gradiente)
            let (ids_duros, _valores) = self.quants[i].quantizar_duro(&zp)?;
            self.quants[i].registrar_uso(&ids_duros)?;
            self.quants[i].atualizar_escala(&zp)?;
            // reconstrução cumulativa na granularidade n_l
            let zhat_l = match &zhat_acum {
                Some(a) => (&interp_para(a, n_l, dev)? + &r_suave)?,
                None => r_suave.clone(),
            };
            zhat_acum = Some(zhat_l.clone());
            zhats.push(zhat_l);
            // resíduo na granularidade do próximo nível (Teo A.0/D.4)
            if i + 1 < self.niveis.len() {
                let n_prox = self.niveis[i + 1];
                let res = (&zp - &r_suave)?;
                z_corrente = interp_para(&res, n_prox, dev)?;
            }
            zs.push(zp);
        }
        Ok((zhats, zs))
    }

    /// Codificação CONGELADA (fase B / inferência): códigos DUROS + ẑ dos
    /// códigos (T determinístico — hipótese do Teo B.1). Recursão residual
    /// inteiramente em d_lat (projeção única — n-02 §0).
    pub fn codificar_congelado(&self, ids: &Tensor) -> Result<Codigos> {
        {
            let dev = ids.device();
            let z = self.encoder(ids)?.detach();
            let mut saida_ids = Vec::new();
            let mut zhats = Vec::new();
            let n1 = self.niveis[0];
            let mut z_corrente = self.proj_niveis[0].fwd(&interp_para(&z, n1, dev)?)?;
            let mut zhat_acum: Option<Tensor> = None;
            for (i, &n_l) in self.niveis.iter().enumerate() {
                let zp = z_corrente.clone();
                let (ids_l, valores_l) = self.quants[i].quantizar_duro(&zp)?;
                saida_ids.push(ids_l.reshape((zp.dims()[0], n_l))?);
                let zhat_l = match &zhat_acum {
                    Some(a) => (&interp_para(a, n_l, dev)? + &valores_l)?,
                    None => valores_l.clone(),
                };
                zhat_acum = Some(zhat_l.clone());
                zhats.push(zhat_l);
                if i + 1 < self.niveis.len() {
                    let n_prox = self.niveis[i + 1];
                    let res = (&zp - &valores_l)?;
                    z_corrente = interp_para(&res, n_prox, dev)?;
                }
            }
            Ok(Codigos { ids: saida_ids, zhats })
        }
    }

    /// Reconstrói ẑ a partir de códigos (para corrupção κ e amostragem —
    /// o gerador nunca vê os vetores contínuos, apenas c (C.3)).
    pub fn reconstruir_de_codigos(&self, ids: &[Tensor], ate_nivel: usize, b: usize, dev: &Device) -> Result<Vec<Tensor>> {
        {
            let mut zhats = Vec::new();
            let mut zhat_acum: Option<Tensor> = None;
            for (i, item) in self.niveis.iter().enumerate().take(ate_nivel.min(self.niveis.len())) {
                let plano = ids[i].reshape((b, self.niveis[i]))?;
                let valores = self.quants[i].materializar(&plano.flatten_all()?)?;
                let valores = valores.reshape((b, self.niveis[i], self.d_lat))?;
                let zhat_l = match &zhat_acum {
                    Some(a) => (&interp_para(a, *item, dev)? + &valores)?,
                    None => valores.clone(),
                };
                zhat_acum = Some(zhat_l.clone());
                zhats.push(zhat_l);
            }
            Ok(zhats)
        }
    }

    /// Decoder AR condicionado: ẑ (b, n_ℓ, d_lat) → logits (b, N, vocab).
    /// Entrada = proj(upsample_N(ẑ)) + emb(tokens deslocados).
    pub fn decoder(&self, zhat: &Tensor, alvo: Option<&Tensor>) -> Result<Tensor> {
        let dev = zhat.device();
        let b = zhat.dim(0)?;
        let n_z = zhat.dim(1)?;
        let n_total = self.mascara_causal_dec.dim(0)?;
        let up = interp_para(zhat, n_total, dev)?;
        let cond = self.dec_proj.fwd(&up)?;
        let entrada_tokens: Tensor = match alvo {
            Some(t) => {
                // desloca à direita com BOS
                let bos = Tensor::full(crate::tokens::BOS as f32, (b, 1), dev)?.to_dtype(candle_core::DType::U32)?;
                let corpo = t.narrow(1, 0, t.dim(1)? - 1)?;
                Tensor::cat(&[&bos, &corpo], 1)?
            }
            None => Tensor::full(crate::tokens::BOS as f32, (b, n_total), dev)?.to_dtype(candle_core::DType::U32)?,
        };
        let mut x = self.dec_emb.fwd(&entrada_tokens)?;
        x = (&x + &cond)?;
        for bloco in &self.dec_blocos {
            let mut cache = Default::default();
            x = bloco.fwd(&x, Some(&self.mascara_causal_dec), &mut cache, None, None)?;
        }
        let x = self.dec_norm.fwd(&x)?;
        self.dec_head.fwd(&x)
    }

    /// Perda da fase A (n-04 §2.3): nested dropout + commitment.
    /// p(ℓ) uniforme (Teo A.2); RNG DEDICADO (pareamento — Teo 31.7).
    /// Devolve (tensor da perda total, métricas).
    pub fn perda_fase_a(&mut self, ids: &Tensor, rng: &mut StdRng) -> Result<(Tensor, PerdaFaseA)> {
        let (zhats, zs) = self.codificar_treino(ids)?;
        let l = zhats.len();
        // sorteio do nível de parada (nested dropout)
        let nivel_parada = rng.gen_range(1..=l);
        let logits = self.decoder(&zhats[nivel_parada - 1], Some(ids))?;
        let (b, n_t, _) = (logits.dim(0)?, logits.dim(1)?, logits.dim(2)?);
        let plano = logits.reshape((b * n_t, self.vocab))?;
        let alvos = ids.reshape(b * n_t)?;
        let ce = candle_nn::loss::cross_entropy(&plano, &alvos)?;
        // commitment: β Σ_ℓ ‖sg[r_ℓ] − z_ℓ‖² (média por elemento)
        let mut commit_t: Option<Tensor> = None;
        let mut commit = 0f64;
        for (i, zp) in zs.iter().enumerate() {
            let (_ids, valores) = self.quants[i].quantizar_duro(zp)?;
            let dif = (&valores - zp)?;
            let termo = dif.sqr()?.mean_all()?;
            commit += termo.to_scalar::<f32>()? as f64;
            commit_t = Some(match commit_t { Some(s) => (&s + &termo)?, None => termo });
        }
        let commit_t = commit_t.unwrap_or_else(|| ce.zeros_like().unwrap());
        let total = (&ce + &commit_t.affine(self.cfg_ncm_beta_commit, 0.0)?)?;
        // precisão de reconstrução (telemetria)
        let preds = plano.argmax(D::Minus1)?;
        let acertos = preds.eq(&alvos)?.to_dtype(candle_core::DType::F32)?.sum_all()?.to_scalar::<f32>()? as f64;
        let acc = acertos / (b * n_t) as f64;
        // δ_ort por nível (V9)
        let deltas: Result<Vec<f64>> = self.quants.iter().map(|q| q.delta_ort()).collect();
        let deltas = deltas?;
        for (i, d) in deltas.iter().enumerate() {
            self.quants[i].delta_ort_ultimo = *d;
        }
        Ok((total, PerdaFaseA {
            rec: ce.to_scalar::<f32>()? as f64,
            commit,
            nivel_parada,
            acc_recon: acc,
            deltas_ort: deltas,
        }))
    }

    /// Decodificação LIVRE (amostragem): ẑ (b, n_ℓ, d_lat) → ids de token AR
    /// greedy. Teo A.2/B.4(b): qualquer prefixo decodifica (anytime).
    /// O texto final é reconstruído pelo BPE no lado Python.
    pub fn decodificar_livre(&self, zhat: &Tensor, max_tokens: usize) -> Result<Vec<Vec<u32>>> {
        {
            let dev = zhat.device();
            let b = zhat.dim(0)?;
            let n_total = self.mascara_causal_dec.dim(0)?;
            let up = interp_para(zhat, n_total, dev)?;
            let cond = self.dec_proj.fwd(&up)?;
            let mut tokens: Vec<Vec<u32>> = (0..b).map(|_| vec![crate::tokens::BOS]).collect();
            for passo in 1..max_tokens.min(n_total) {
                let planos: Vec<u32> = tokens.iter().flat_map(|t| t.iter().copied()).collect();
                let entrada = Tensor::from_vec(planos, (b, passo), dev)?;
                let mut x = self.dec_emb.fwd(&entrada)?;
                x = (&x + &cond.narrow(1, 0, passo)?)?;
                for bloco in &self.dec_blocos {
                    let mut cache = Default::default();
                    let m = self.mascara_causal_dec.narrow(0, 0, passo)?.narrow(1, 0, passo)?.contiguous()?;
                    x = bloco.fwd(&x, Some(&m), &mut cache, None, None)?;
                }
                let x = self.dec_norm.fwd(&x)?;
                let logits = self.dec_head.fwd(&x)?;
                let ultimo = logits.narrow(1, passo - 1, 1)?.reshape((b, self.vocab))?;
                let preds = ultimo.argmax(D::Minus1)?;
                let v = preds.flatten_all()?.to_vec1::<u32>()?;
                for (i, t) in v.iter().enumerate() {
                    tokens[i].push(*t);
                }
            }
            Ok(tokens)
        }
    }

    /// CE por prefixo (V3: monotonia do risco — Teo A.2) em modo congelado.
    pub fn ce_por_prefixo(&self, ids: &Tensor) -> Result<Vec<f64>> {
        {
            let cod = self.codificar_congelado(ids)?;
            let mut saida = Vec::with_capacity(cod.zhats.len());
            for zhat in &cod.zhats {
                let logits = self.decoder(zhat, Some(ids))?;
                let (b, n_t, _) = (logits.dim(0)?, logits.dim(1)?, logits.dim(2)?);
                let plano = logits.reshape((b * n_t, self.vocab))?;
                let alvos = ids.reshape(b * n_t)?;
                let ce = candle_nn::loss::cross_entropy(&plano, &alvos)?;
                saida.push(ce.to_scalar::<f32>()? as f64);
            }
            Ok(saida)
        }
    }

    /// Agenda τ do GSQ por passo (Teo 30.6).
    pub fn agendar_tau(&mut self, passo: usize, tau0: f64, tau_min: f64, t_alvo: usize) {
        for q in &mut self.quants {
            q.agendar_tau(passo, tau0, tau_min, t_alvo);
        }
    }

    /// Atualização EMA do codebook (modo ste_ema apenas — no gsq_rco o
    /// gradiente riemanniano já treina W).
    pub fn atualizar_ema_codebook(&mut self, ids_lote: &[Tensor]) -> Result<()> {
        if self.modo != ModoQuant::SteEma {
            return Ok(());
        }
        // no modo STE, os codebooks seguem por EMA dos agrupamentos (A.0).
        // Implementação: média dos vetores por código → polar → w.
        for (i, q) in self.quants.iter_mut().enumerate() {
            let ids_l = &ids_lote[i];
            let k = q.uso.len();
            let b = ids_l.dim(0)?;
            let plano = ids_l.reshape(b * self.niveis[i])?;
            let ids_v = plano.flatten_all()?.to_vec1::<u32>()?;
            let mut soma = vec![vec![0f32; self.d_lat]; k];
            let mut conta = vec![0usize; k];
            // usa o codebook corrente para os vetores (aproximação EMA dos
            // agrupamentos — suficiente para o modo STE de compatibilidade)
            let c = q.codebook()?.detach();
            let c_v = c.flatten_all()?.to_vec1::<f32>()?;
            for &idc in &ids_v {
                let idc = idc as usize;
                if idc < k {
                    for t in 0..self.d_lat {
                        soma[idc][t] += c_v[idc * self.d_lat + t];
                    }
                    conta[idc] += 1;
                }
            }
            let ema = q.ema;
            let w_v = q.w.as_tensor().flatten_all()?.to_vec1::<f32>()?;
            let mut novo = w_v.clone();
            for idc in 0..k {
                if conta[idc] > 0 {
                    for t in 0..self.d_lat {
                        let alvo = soma[idc][t] / conta[idc] as f32;
                        novo[idc * self.d_lat + t] =
                            (ema as f32 * w_v[idc * self.d_lat + t] + (1.0 - ema as f32) * alvo) as f32;
                    }
                }
            }
            let dev = q.w.as_tensor().device();
            let nt = Tensor::from_vec(novo, (k, self.d_lat), dev)?;
            q.w.set(&nt)?;
        }
        Ok(())
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v: Vec<Var> = self.enc_emb.vars();
        for b in &self.enc_blocos { v.extend(b.vars()); }
        v.extend(self.enc_norm.vars());
        for p in &self.proj_niveis { v.extend(p.vars()); }
        for q in &self.quants { v.push(q.w.clone()); }
        v.extend(self.dec_proj.vars());
        v.extend(self.dec_emb.vars());
        for b in &self.dec_blocos { v.extend(b.vars()); }
        v.extend(self.dec_norm.vars());
        v.extend(self.dec_head.vars());
        v
    }

    /// Salva estado (checkpoint atômico — Teo 35.3): índices = ordem de vars().
    pub fn salvar(&self, caminho: &str) -> Result<()> {
        let mut obj = serde_json::Map::new();
        for (i, v) in self.vars().iter().enumerate() {
            let dados = v.as_tensor().flatten_all()?.to_vec1::<f32>()?;
            obj.insert(format!("v{i}"), serde_json::json!({
                "forma": v.as_tensor().dims(),
                "dados": dados,
            }));
        }
        let txt = serde_json::to_string(&serde_json::Value::Object(obj))
            .map_err(|e| candle_core::Error::Msg(format!("serde: {e}")))?;
        let tmp = format!("{caminho}.tmp");
        std::fs::write(&tmp, txt).map_err(|e| candle_core::Error::Msg(format!("io: {e}")))?;
        std::fs::rename(&tmp, caminho).map_err(|e| candle_core::Error::Msg(format!("rename: {e}")))?;
        Ok(())
    }
}

pub struct PerdaFaseA {
    pub rec: f64,
    pub commit: f64,
    pub nivel_parada: usize,
    pub acc_recon: f64,
    pub deltas_ort: Vec<f64>,
}

/// Máscara causal aditiva (n, n): 0 permitido, MIN bloqueado.
pub fn causal_aditiva(n: usize, dev: &Device) -> Result<Tensor> {
    let mut m = vec![f32::MIN; n * n];
    for i in 0..n {
        for j in 0..=i {
            m[i * n + j] = 0.0;
        }
    }
    Tensor::from_vec(m, (n, n), dev)
}

fn tensor_para_json(t: &Tensor) -> Result<serde_json::Value> {
    let dados = t.flatten_all()?.to_vec1::<f32>()?;
    Ok(serde_json::json!({
        "forma": t.shape().dims(),
        "dados": dados,
    }))
}

/// Matriz de interpolação cacheada (evita recriar por passo).
pub struct CacheInterp {
    pub n_de: usize,
    pub n_para: usize,
    pub matriz: Tensor,
}

impl CacheInterp {
    pub fn novo(n_de: usize, n_para: usize, dev: &Device) -> Result<Self> {
        Ok(CacheInterp { n_de, n_para, matriz: matriz_interp(n_de, n_para, dev)? })
    }
}