//! nlp.rs — **UnidadeNLP** (especialista de processamento/entrada) — port
//! Rust de `khtst-multimodal-ptbr/src/khtst/nlp/unidade_nlp.py` (doc 15 §1).
//!
//! Papel (Def. 15.1): enriquece a representação textual com
//! (i) traços morfo-sintáticos leves (conv 1-D de n-gramas, k=3 e k=5,
//!     grupos=4 — proxy de bigramas de sufixo);
//! (ii) intenção/tarefa estimada por cabeça própria (13 classes canônicas);
//! (iii) extensões de entidade por apontador (score por posição);
//! e injeta por PORTA ADITIVA `h⁺ = h + g ⊙ Δ(h)`, `g = σ(W_g h)` —
//! **nunca substitui o fluxo do tronco** (monotonia de capacidade).
//!
//! **Teo N-1 (não-interferência no nascimento — port do Teo 15.1):** com
//! `W_Δ` inicializado em ZERO, Δ(h) = 0 bit-a-bit para qualquer h ⇒ a
//! unidade acoplada a um modelo pré-treinado NÃO muda sua saída inicial
//! (extensão segura); com `b_g = −2` (logit), g ≈ 0,12 no nascimento e a
//! contribuição cresce suavemente com o treino. ∎ (espelho do Teo 15.1;
//! gate G18 verifica bit-a-bit.)

use crate::camadas::Linear;
use candle_core::{DType, Device, Result, Tensor, Var};
use std::collections::HashMap;

/// 13 intenções: 12 tarefas canônicas da tabela do KHTST + "nlp-outro"
/// (fonte única — a tabela canônica vive AQUI; pigeonhole: cada tarefa tem
/// exatamente UMA intenção — injetividade sem divergência possível).
pub const INTENCOES: &[&str] = &[
    "nlp-outro",
    "classificacao",
    "resumo",
    "tradução",
    "correcao",
    "extracao",
    "dialogo",
    "raciocinio",
    "codigo",
    "matematica",
    "narrativa",
    "consulta",
    "instrucao",
];

pub const N_EIXOS_ESTILO: usize = 8;

/// Unidade NLP — especialista de entrada (port fiel, v0.7.0).
pub struct UnidadeNlp {
    pub d: usize,
    /// (i) conv de n-gramas k=3 (grupos=4)
    pub conv_ngram: Var, // (d, d/4, 3)
    /// (i) conv de n-gramas k=5 (grupos=4)
    pub conv_ngram2: Var, // (d, d/4, 5)
    /// (ii) cabeça de intenção (13 classes)
    pub cabeca_intencao: Linear,
    /// (iii) apontador de entidades (score por posição)
    pub pontuador: Linear,
    /// fusão aditiva com porta: h⁺ = h + g ⊙ Δ(h)
    pub w_delta: Linear, // (3d → d) — ZERO no nascimento (Teo N-1)
    pub w_gate: Linear,  // (d → d) — bias −2 ⇒ g ≈ 0,12 no nascimento
    /// telemetria local
    pub gate_medio: f64,
}

impl UnidadeNlp {
    pub fn nova(d: usize, dev: &Device) -> Result<Self> {
        if d % 4 != 0 {
            return Err(candle_core::Error::Msg(
                "UnidadeNlp: d deve ser divisível por 4 (conv em grupos)".into(),
            ));
        }
        let dg = d / 4;
        // conv: peso (out_channels, in_channels/groups, kernel)
        let c1 = Tensor::randn(0f32, 0.02f32, (d, dg, 3), dev)?;
        let c2 = Tensor::randn(0f32, 0.02f32, (d, dg, 5), dev)?;
        // W_Δ ZERO (nascimento neutro — Teo N-1)
        let wd = Tensor::zeros((d, 3 * d), DType::F32, dev)?;
        // W_g com bias −2
        let wg = Tensor::randn(0f32, 0.02f32, (d, d), dev)?;
        let bg = Tensor::full(-2.0f32, (d,), dev)?;
        Ok(UnidadeNlp {
            d,
            conv_ngram: Var::from_tensor(&c1)?,
            conv_ngram2: Var::from_tensor(&c2)?,
            cabeca_intencao: Linear::novo(d, INTENCOES.len(), 0.02, dev)?,
            pontuador: Linear::novo(d, 1, 0.02, dev)?,
            w_delta: Linear { w: Var::from_tensor(&wd)?, b: None },
            w_gate: Linear {
                w: Var::from_tensor(&wg)?,
                b: Some(Var::from_tensor(&bg)?),
            },
            gate_medio: 0.0,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        vec![
            self.conv_ngram.clone(),
            self.conv_ngram2.clone(),
            self.w_delta.w.clone(),
            self.w_gate.w.clone(),
            self.w_gate.b.clone().unwrap(),
        ]
        .into_iter()
        .chain(self.cabeca_intencao.vars())
        .chain(self.pontuador.vars())
        .collect()
    }

    /// Traços morfo-sintáticos: conv1D sobre os embeddings (proxy de
    /// n-gramas de caracteres). h: (B, T, d) → (B, T, d).
    pub fn tracos(&self, h: &Tensor) -> Result<Tensor> {
        let (b, t, d) = (h.dims()[0], h.dims()[1], h.dims()[2]);
        // conv1d espera (B, C, T): transpõe, aplica, transpõe de volta
        let x = h.transpose(1, 2)?.contiguous()?;
        let c1 = x.conv1d(&self.conv_ngram.as_tensor(), 1, 1, 1, 4)?;
        let c2 = x.conv1d(&self.conv_ngram2.as_tensor(), 2, 1, 1, 4)?;
        let soma = c1.add(&c2)?;
        soma.transpose(1, 2)?.contiguous()?.reshape((b, t, d))
    }

    /// Intenção estimada (logits das 13 classes) a partir do pooling médio.
    pub fn intencao(&self, h: &Tensor) -> Result<Tensor> {
        let pool = h.mean(1)?; // (B, d)
        self.cabeca_intencao.fwd(&pool)
    }

    /// Scores de entidade por posição: (B, T, 1).
    pub fn spans(&self, h: &Tensor) -> Result<Tensor> {
        self.pontuador.fwd(h)
    }

    /// **Enriquecimento aditivo com porta** (Teo N-1): h⁺ = h + g ⊙ Δ(h).
    /// Δ = W_Δ · [tracos; int_pool; span_d] (3d → d).
    pub fn enriquecer(&mut self, h: &Tensor) -> Result<Tensor> {
        let (b, t, d) = (h.dims()[0], h.dims()[1], h.dims()[2]);
        let tr = self.tracos(h)?;
        let int = self.intencao(h)?;
        // sinal global da intenção: média dos logits → repetido por posição
        // (a projeção fina vive no w_delta — custo O(1) por posição)
        let int_pool = int.mean(1)?.reshape((b, 1, 1))?.broadcast_as((b, t, d))?;
        let sp = self.spans(h)?;
        let denom = sp.abs()?.max_all()?.to_scalar::<f32>()?.max(1e-6);
        let sp_norm = sp.affine(1.0 / denom as f64, 0.0)?;
        let span_d = sp_norm.reshape((b, t, 1))?.broadcast_as((b, t, d))?;
        // concatena [tracos; int_pool; span_d] → (B, T, 3d) → Δ (B, T, d)
        let concat = Tensor::cat(&[&tr, &int_pool, &span_d], 2)?;
        let delta = self.w_delta.fwd(&concat)?;
        // porta g = σ(W_g h + b_g) — bias −2 ⇒ g ≈ 0,12 no nascimento
        let gate = candle_nn::ops::sigmoid(&self.w_gate.fwd(h)?)?;
        // telemetria: gate médio
        self.gate_medio = gate.mean_all()?.to_scalar::<f32>()? as f64;
        // h⁺ = h + g ⊙ Δ (ADITIVO — nunca substitui o tronco)
        h.broadcast_add(&gate.broadcast_mul(&delta)?)
    }

    /// Intenção predita (índice) + confiança simples.
    pub fn intencao_argmax(&self, h: &Tensor) -> Result<(Vec<usize>, Vec<f64>)> {
        let logits = self.intencao(h)?;
        let (b, k) = (logits.dims()[0], logits.dims()[1]);
        let v = logits.flatten_all()?.to_vec1::<f32>()?;
        let mut idx = Vec::with_capacity(b);
        let mut conf = Vec::with_capacity(b);
        for bi in 0..b {
            let mut melhor = 0usize;
            let mut melhor_v = f32::NEG_INFINITY;
            for ki in 0..k {
                let x = v[bi * k + ki];
                if x > melhor_v {
                    melhor_v = x;
                    melhor = ki;
                }
            }
            // softmax só sobre o melhor (confiança relativa simples)
            let soma_exp: f32 = (0..k).map(|ki| (v[bi * k + ki] - melhor_v).exp()).sum();
            idx.push(melhor);
            conf.push(1.0 / soma_exp as f64);
        }
        Ok((idx, conf))
    }

    /// Telemetria local da unidade.
    pub fn telemetria(&self) -> HashMap<String, f64> {
        let mut m = HashMap::new();
        m.insert("nlp_gate_medio".into(), self.gate_medio);
        m.insert("nlp_n_intencoes".into(), INTENCOES.len() as f64);
        m
    }
}

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

    /// N-1 (port do Teo 15.1): nascimento neutro — W_Δ = 0 ⇒ enriquecer(h)
    /// = h bit-a-bit para qualquer h (a unidade acoplada não muda o tronco).
    #[test]
    fn n1_nascimento_neutro() -> Result<()> {
        let dev = Device::Cpu;
        let mut u = UnidadeNlp::nova(32, &dev)?;
        let h = Tensor::randn(0f32, 1f32, (2, 7, 32), &dev)?;
        let h2 = u.enriquecer(&h)?;
        let dif = (&h - &h2)?.abs()?.max_all()?.to_scalar::<f32>()?;
        assert_eq!(dif, 0.0, "N-1: Δ=0 no nascimento ⇒ enriquecer ≡ h (dif {dif})");
        // gate médio no nascimento ≈ σ(−2) ≈ 0,12
        let g = u.gate_medio;
        assert!(
            (g - 0.119).abs() < 0.03,
            "N-1: gate médio {g} deve ≈ σ(−2) = 0,119"
        );
        Ok(())
    }

    /// N-2: traços preservam a forma (B, T, d) e são finitos.
    #[test]
    fn n2_tracos_forma() -> Result<()> {
        let dev = Device::Cpu;
        let u = UnidadeNlp::nova(16, &dev)?;
        let h = Tensor::randn(0f32, 1f32, (2, 5, 16), &dev)?;
        let tr = u.tracos(&h)?;
        assert_eq!(tr.dims(), vec![2, 5, 16]);
        let mx = tr.abs()?.max_all()?.to_scalar::<f32>()?;
        assert!(mx.is_finite(), "N-2: traços finitos");
        Ok(())
    }

    /// N-3: injetividade da tabela canônica (pigeonhole — sem duplicatas).
    #[test]
    fn n3_intencoes_injetivas() {
        let mut vistas: Vec<&str> = INTENCOES.to_vec();
        vistas.sort();
        let n = vistas.len();
        vistas.dedup();
        assert_eq!(vistas.len(), n, "N-3: intenções canônicas sem duplicatas");
        assert_eq!(n, 13, "N-3: 12 tarefas + nlp-outro");
    }

    /// N-4: após perturbar W_Δ (treino simulado), o enriquecimento muda —
    /// a porta EXISTE e aprende (não é uma constante morta).
    #[test]
    fn n4_gate_apos_treino_simulado() -> Result<()> {
        let dev = Device::Cpu;
        let mut u = UnidadeNlp::nova(32, &dev)?;
        let h = Tensor::randn(0f32, 1f32, (1, 4, 32), &dev)?;
        // treino simulado: W_Δ ← W_Δ + ruído
        let ruido = Tensor::randn(0f32, 0.1f32, (32, 96), &dev)?;
        let novo = u.w_delta.w.as_tensor().add(&ruido)?;
        u.w_delta.w.set(&novo)?;
        let h2 = u.enriquecer(&h)?;
        let dif = (&h - &h2)?.abs()?.max_all()?.to_scalar::<f32>()?;
        assert!(dif > 0.0, "N-4: com W_Δ ≠ 0 o enriquecimento deve agir (dif {dif})");
        assert!(dif < 1.0, "N-4: porta g ≈ 0,12 limita a perturbação (dif {dif})");
        Ok(())
    }
}