Blame for src/tronco.rs

//! Tronco híbrido 5:1 (port do `mamba3/tronco.py`): 5 camadas ADBA (atenção
//! block-causal com melhorias) + 1 camada SSD (slot Mamba-3 causal) — Teo D.1
//! com a restrição honesta H_D1 (bidirecionalidade intra-nível nas camadas de
//! atenção; a camada recorrente é causal).
//!
//! Sequência do gerador: s = [cond; lev_1; …; lev_L] com entradas
//! ẑ_{ℓ−1}^{(n_ℓ)} + e_ℓ (reconstrução cumulativa interpolada + embedding de
//! nível — Teo D.1). Treino: 1 passada com teacher forcing (D.1(iv));
//! inferência: bloco a bloco com KV-cache (D.1(iii), T2).

use candle_core::{Device, Result, Tensor, Var};
use crate::atencao::{AncoraNivel, BlocoAdba};
pub use crate::atencao::CacheKv;
use crate::camadas::{Embedding, Linear, RmsNorm};
use crate::config::ConfigKhtst;
use crate::ssd::BlocoSsd;

pub struct Tronco {
    pub blocos_atencao: Vec<BlocoAdba>,
    pub bloco_ssd: BlocoSsd,
    pub norm_final: RmsNorm,
    pub ancora: Option<AncoraNivel>,
    /// gamma FiLM por nível (2*L+2: [escala, deslocamento] × nível) — ADBA.
    pub film_gamma: Option<Var>,
    pub film_beta: Option<Var>,
    /// W_f do FiLM (r-04 §1: x̃ = x⊙(1+γ·tanh(W_f·h)) + β·(W_f·h)) — ZERO no
    /// nascimento: tanh(0)=0 e β·0=0 ⇒ identidade bit-a-bit (Teo R.2).
    /// v070-1 (auditoria): a v0.6.0 injetava o resumo do âncora DIRETO (projeção
    /// aleatória) — nascimento não-neutro nos níveis ≥ 2; correção restaura o
    /// enunciado R.2 independentemente de γ=1.
    pub film_w: Option<crate::camadas::Linear>,
    pub n_niveis: usize,
}

impl Tronco {
    pub fn novo(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> {
        let m = &cfg.modelo;
        let l = cfg.l_niveis();
        let adba = &cfg.adba;
        let mut blocos = Vec::with_capacity(m.n_camadas - 1);
        for _ in 0..(m.n_camadas - 1) {
            blocos.push(BlocoAdba::novo(
                m.d_modelo, m.n_cabecas, m.d_ff, l,
                adba.ativo, adba.reponderacao_nivel, adba.temperatura_nivel, dev,
            )?);
        }
        let ancora = if adba.ativo && adba.gamma_ancora > 0.0 {
            Some(AncoraNivel::nova(m.d_modelo, l, dev)?)
        } else {
            None
        };
        let (film_gamma, film_beta, film_w) = if adba.ativo {
            let g = Tensor::ones((l + 1, m.d_modelo), candle_core::DType::F32, dev)?;
            let b = Tensor::zeros((l + 1, m.d_modelo), candle_core::DType::F32, dev)?;
            let w = crate::camadas::Linear::zeros_sem_vies(m.d_modelo, m.d_modelo, dev)?;
            (Some(Var::from_tensor(&g)?), Some(Var::from_tensor(&b)?), Some(w))
        } else {
            (None, None, None)
        };
        Ok(Tronco {
            blocos_atencao: blocos,
            bloco_ssd: BlocoSsd::novo(m.d_modelo, m.d_ff, dev)?,
            norm_final: RmsNorm::novo(m.d_modelo, 1e-5, dev)?,
            ancora,
            film_gamma,
            film_beta,
            film_w,
            n_niveis: l,
        })
    }

    /// Injeção FiLM do resumo-âncora nas posições do nível `nivel`
    /// (ADBA (i): o resumo depende APENAS de c_{<ℓ} — paralelismo preservado).
    /// r-04 §1: x̃ = x⊙(1+γ·tanh(W_f·h)) + β·(W_f·h) com W_f=0 no nascimento
    /// ⇒ identidade bit-a-bit para QUALQUER resumo (Teo R.2, v070-1).
    pub fn aplicar_film(&self, x: &Tensor, resumo: &Tensor, nivel: usize) -> Result<Tensor> {
        match (&self.film_gamma, &self.film_beta, &self.film_w) {
            (Some(g), Some(b), Some(wf)) => {
                let d = x.dims()[2];
                let bsz = resumo.dims()[0].max(1);
                let gv = g.as_tensor().get(nivel)?.reshape((1, 1, d))?;
                let bv = b.as_tensor().get(nivel)?.reshape((1, 1, d))?;
                // r = W_f·resumo — 0 no nascimento (identidade garantida)
                let r = wf.fwd(resumo)?.reshape((bsz, 1, d))?;
                // x ⊙ (1 + γ·tanh(r)) + β·r
                let gate = gv.broadcast_mul(&r.tanh()?)?.affine(1.0, 1.0)?; // 1 + γ·tanh(r)
                let xg = x.broadcast_mul(&gate)?;
                let desloc = bv.broadcast_mul(&r)?;
                xg.broadcast_add(&desloc)
            }
            _ => Ok(x.clone()),
        }
    }

    /// Forward de treino: sequência completa s (b, T, d) com máscara aditiva
    /// (T, T) e mapa de níveis por posição (query e chave). 1 passada — D.1(iv).
    pub fn forward_treino(
        &self,
        s: &Tensor,
        mascara: &Tensor,
        niveis: &[usize],
        resumos_por_nivel: &mut Vec<Tensor>,
    ) -> Result<Tensor> {
        let b = s.dim(0)?;
        let mut x = s.clone();
        for (i, bloco) in self.blocos_atencao.iter().enumerate() {
            let mut cache = CacheKv::default();
            x = bloco.fwd(&x, Some(mascara), &mut cache, Some(niveis), Some(niveis))?;
            // após o 1º bloco: injeta o resumo-âncora por nível (ADBA)
            if i == 0 {
                if let Some(anc) = &self.ancora {
                    resumos_por_nivel.clear();
                    let n_prefixo = x.dims()[1] - self.tamanhos_por_nivel(niveis).iter().sum::<usize>();
                    // prefixo entra inalterado (sem FiLM — nível 0)
                    let mut partes = vec![x.narrow(1, 0, n_prefixo)?];
                    let mut inicio = n_prefixo;
                    for (nivel, &n_n) in self.tamanhos_por_nivel(niveis).iter().enumerate() {
                        let trecho = x.narrow(1, inicio, n_n)?;
                        // resumo do CONTEXTO ANTERIOR (níveis < atual): usa as
                        // posições já processadas (função de c_{<ℓ} apenas)
                        if inicio == n_prefixo {
                            resumos_por_nivel.push(Tensor::zeros((b, s.dims()[2]), candle_core::DType::F32, s.device())?);
                        } else {
                            let ctx = x.narrow(1, 0, inicio)?;
                            let mut r_nivel = Vec::new();
                            for bi in 0..b {
                                r_nivel.push(anc.resumir(&ctx.narrow(0, bi, 1)?, nivel)?);
                            }
                            resumos_por_nivel.push(Tensor::cat(&r_nivel, 0)?);
                        }
                        let r = resumos_por_nivel[nivel].clone();
                        partes.push(self.aplicar_film(&trecho, &r, nivel + 1)?);
                        inicio += n_n;
                    }
                    x = Tensor::cat(&partes, 1)?;
                }
            }
        }
        let mut cache_ssd = None;
        x = self.bloco_ssd.fwd(&x, &mut cache_ssd)?;
        self.norm_final.fwd(&x)
    }

    /// Forward incremental (inferência): bloco novo (b, n_novo, d) + caches.
    /// `niveis_novos`/`niveis_cache` mapeiam posições para níveis.
    /// `film` = Some((resumo, nivel 1-indexado)) aplica o FiLM nas posições novas
    /// após o bloco 0 — PARIDADE com forward_treino (v070-2: a v0.6.0 não
    /// aplicava FiLM na inferência incremental — treino↔geração divergiam).
    /// `saida_b0` = Some(&mut Tensor) recebe os estados pós-bloco-0 das
    /// posições novas (para o âncora acumular contexto nos níveis seguintes).
    pub fn forward_nivel(
        &self,
        bloco_x: &Tensor,
        caches: &mut Vec<CacheKv>,
        mascara_nova: &Tensor,
        niveis_novos: &[usize],
        niveis_totais: &[usize],
        film: Option<(&Tensor, usize)>,
        mut saida_b0: Option<&mut Tensor>,
    ) -> Result<Tensor> {
        let mut x = bloco_x.clone();
        for (i, bloco) in self.blocos_atencao.iter().enumerate() {
            if caches.len() <= i {
                caches.push(CacheKv::default());
            }
            x = bloco.fwd(&x, Some(mascara_nova), &mut caches[i], Some(niveis_novos), Some(niveis_totais))?;
            if i == 0 {
                if let Some(slot) = saida_b0.as_deref_mut() {
                    *slot = x.clone();
                }
                if let Some((resumo, nivel)) = film {
                    x = self.aplicar_film(&x, resumo, nivel)?;
                }
            }
        }
        let mut cache_ssd = None;
        x = self.bloco_ssd.fwd(&x, &mut cache_ssd)?;
        self.norm_final.fwd(&x)
    }

    fn tamanhos_por_nivel(&self, niveis: &[usize]) -> Vec<usize> {
        let max_n = niveis.iter().copied().max().unwrap_or(0);
        let mut out = vec![0usize; max_n];
        for &n in niveis {
            if n >= 1 && n <= max_n {
                out[n - 1] += 1;
            }
        }
        out
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v: Vec<Var> = self.blocos_atencao.iter().flat_map(|b| b.vars()).collect();
        v.extend(self.bloco_ssd.vars());
        v.extend(self.norm_final.vars());
        if let Some(a) = &self.ancora {
            v.extend(a.vars());
        }
        if let Some(g) = &self.film_gamma {
            v.push(g.clone());
        }
        if let Some(b) = &self.film_beta {
            v.push(b.clone());
        }
        if let Some(wf) = &self.film_w {
            v.push(wf.w.clone());
        }
        v
    }
}

/// Cabeças de conceitos por nível (via A de T1): linear d → K_ℓ.
pub struct CabecaConceitos {
    pub cabecas: Vec<Linear>,
}

impl CabecaConceitos {
    pub fn nova(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> {
        let mut c = Vec::with_capacity(cfg.l_niveis());
        for _ in 0..cfg.l_niveis() {
            c.push(Linear::sem_vies(cfg.modelo.d_modelo, cfg.ncm.k_por_nivel, 0.02, dev)?);
        }
        Ok(CabecaConceitos { cabecas: c })
    }

    pub fn vars(&self) -> Vec<Var> {
        self.cabecas.iter().flat_map(|c| c.vars()).collect()
    }
}

/// Embeddings de nível e_ℓ (T1 via A).
pub struct EmbNiveis {
    pub tabela: Embedding,
}

impl EmbNiveis {
    pub fn nova(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> {
        Ok(EmbNiveis { tabela: Embedding::novo(cfg.l_niveis() + 1, cfg.modelo.d_modelo, dev)? })
    }
    pub fn vars(&self) -> Vec<Var> {
        self.tabela.vars()
    }
}