//! nlg.rs — **UnidadeNLG** (especialista de geração/saída) + decodificador
//! especulativo em árvore (Medusa) — port Rust de
//! `khtst-multimodal-ptbr/src/khtst/nlg/{unidade_nlg.py,
//! decodificador_especulativo.py}` e `aurora-multimodal-ptbr/src/aurora/nlg/`
//! (doc 15 §1 / doc 14; Teos 15.1, 14.1, 14.2).
//!
//! Papel (Def. 15.2): a NLG **não tem parâmetros da lm_head** — gera
//! MODULANDO o estado oculto (decisão de projeto: evita duplicar
//! 16384×192):
//! (i) vetor de PLANO (tema/estilo) por pooling atencional do prompt;
//! (ii) re-ranking de coerência: penaliza candidatos com bigrama não visto
//!      no contexto recente (tabela EMA de bigramas, hashing trick,
//!      O(|V_cand|));
//! (iii) condicionamento de estilo: h_dec = h + W_estilo·p, p ∈ R^8 —
//!      W_estilo ZERO no nascimento (Teo N-5 = port do Teo 15.1).
//!
//! Decodificador especulativo (doc 14): rascunho = cabeças MTP com logits
//! EMPATADOS; verificação em UM forward com aceitação TÍPICA (Heim 2022,
//! eq. 14.2: aceita sse −log p(t) − H ≤ τ·H); poda por cabeça
//! c_k = min(c, max(1, c−k)); EMA de estatísticas (0,9/0,1); guarda contra
//! laço infinito (teto de árvore). **Teo E-1 (inocuidade — port do 14.1):**
//! a política com correção residual preserva EXATAMENTE p_alvo — a
//! especulação nunca muda a distribuição, só o custo. **Teo E-2 (speedup):**
//! S = (E[aceitos]+1)/(1+|T|/V_ef) > 1 sse ρ > |T|/(V_ef+|T|).

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

pub const N_EIXOS_ESTILO: usize = 8;

/// Unidade NLG — especialista de saída (port fiel, v0.7.0).
pub struct UnidadeNlg {
    pub d: usize,
    /// (i) pooling do plano: projeção para pesos de atenção sobre o prompt.
    pub pooler: Linear, // (d → d), sem viés
    /// (iii) eixos de estilo aprendíveis (logits de controle).
    pub eixos: Var, // (8,)
    /// (iii) W_estilo: p ∈ R^8 → deslocamento em d — ZERO no nascimento.
    pub w_estilo: Var, // (d, 8)
    /// (ii) tabela EMA de bigramas (hashing trick) para re-ranking.
    pub tab_bigramas: Vec<f32>,
    pub tam_tab: usize,
    pub beta_tab: f64,
    /// telemetria local
    pub ultimo_score_coerencia: f64,
}

impl UnidadeNlg {
    pub fn nova(d: usize, tam_tab: usize, dev: &Device) -> Result<Self> {
        let we = Tensor::zeros((d, N_EIXOS_ESTILO), candle_core::DType::F32, dev)?;
        let eixos = Tensor::zeros((N_EIXOS_ESTILO,), candle_core::DType::F32, dev)?;
        Ok(UnidadeNlg {
            d,
            pooler: Linear::sem_vies(d, d, 0.02, dev)?,
            eixos: Var::from_tensor(&eixos)?,
            w_estilo: Var::from_tensor(&we)?,
            tab_bigramas: vec![0.0; tam_tab],
            tam_tab,
            beta_tab: 0.98,
            ultimo_score_coerencia: 0.0,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        vec![
            self.pooler.w.clone(),
            self.eixos.clone(),
            self.w_estilo.clone(),
        ]
    }

    /// (i) Vetor de PLANO: pooling atencional do prompt.
    /// h_prompt: (B, Tp, d) → plano (B, d).
    pub fn plano(&self, h_prompt: &Tensor) -> Result<Tensor> {
        let scores = self.pooler.fwd(h_prompt)?; // (B, Tp, d)
        let pesos = candle_nn::ops::softmax(&scores.mean(2)?, candle_core::D::Minus1)?; // (B, Tp)
        // Σ_t pesos_t · h_t
        let p = pesos.unsqueeze(2)?.broadcast_mul(h_prompt)?;
        p.sum(1)
    }

    /// (iii) Condicionamento de estilo: h_dec = h + W_estilo·p.
    /// `estilo`: (8,) em [0,1]; None ⇒ σ(eixos) (controle aprendível).
    /// W_estilo = 0 no nascimento ⇒ h_dec = h bit-a-bit (Teo N-5).
    pub fn condicionar(&self, h: &Tensor, estilo: Option<&Tensor>) -> Result<Tensor> {
        let p = match estilo {
            Some(p) => p.clone(),
            None => candle_nn::ops::sigmoid(self.eixos.as_tensor())?,
        };
        // deslocamento (..., d) = W_estilo · p
        let desloc = self
            .w_estilo
            .as_tensor()
            .matmul(&p.reshape((N_EIXOS_ESTILO, 1))?)?
            .reshape((1, self.d))?;
        let forma = h.dims().to_vec();
        let n: usize = forma.iter().take(forma.len().saturating_sub(1)).product::<usize>().max(1);
        let h2 = h.reshape((n, forma[forma.len() - 1]))?;
        let d2 = desloc.broadcast_as((n, self.d))?;
        let out = h2.broadcast_add(&d2)?;
        out.reshape(forma)
    }

    /// Hash de bigrama (hashing trick — FNV-1a 32 bits mod tabela).
    fn hash_bigrama(a: u32, b: u32, tam: usize) -> usize {
        let mut h: u32 = 2166136261;
        for x in [a, b] {
            h ^= x;
            h = h.wrapping_mul(16777619);
        }
        (h as usize) % tam
    }

    /// (ii) Atualiza a tabela EMA com os bigramas do CONTEXTO (o que o
    /// modelo já emitiu/viu — coerência dinâmica).
    pub fn observar_contexto(&mut self, tokens: &[u32]) {
        for j in 1..tokens.len() {
            let idx = Self::hash_bigrama(tokens[j - 1], tokens[j], self.tam_tab);
            let atual = self.tab_bigramas[idx];
            self.tab_bigramas[idx] =
                (atual as f64 * self.beta_tab + 1.0 * (1.0 - self.beta_tab)) as f32;
        }
    }

    /// (ii) Score de coerência de um candidato: fração de bigramas JÁ
    /// vistos no contexto (EMA > 0,02 conta como visto) — quanto maior,
    /// mais coerente com o histórico. O(v) para v tokens do candidato.
    pub fn score_coerencia(&mut self, candidato: &[u32]) -> f64 {
        if candidato.len() < 2 {
            self.ultimo_score_coerencia = 1.0;
            return 1.0;
        }
        let mut vistos = 0usize;
        for j in 1..candidato.len() {
            let idx = Self::hash_bigrama(candidato[j - 1], candidato[j], self.tam_tab);
            // EMA com UMA observação = (1−β) = 0,02 — o limiar de "visto"
            // fica ABAIXO disso (0,015) para contar a primeira ocorrência
            if self.tab_bigramas[idx] >= 0.015 {
                vistos += 1;
            }
        }
        let s = vistos as f64 / (candidato.len() - 1) as f64;
        self.ultimo_score_coerencia = s;
        s
    }

    /// (ii) Re-ranking de coerência: dado um conjunto de candidatos com
    /// log-probabilidades do modelo, reordena por λ·logp + (1−λ)·coerência.
    /// Devolve o índice do vencedor e os scores (telemetria).
    pub fn reranking_coerencia(
        &mut self,
        candidatos: &[Vec<u32>],
        logps: &[f64],
        lambda: f64,
    ) -> (usize, Vec<f64>) {
        let n = candidatos.len();
        let mut scores: Vec<f64> = Vec::with_capacity(n);
        let mut coers: Vec<f64> = Vec::with_capacity(n);
        for c in candidatos {
            coers.push(self.score_coerencia(c));
        }
        // normaliza logp por z-score robusto (mediana)
        let mut ordenado = logps.to_vec();
        ordenado.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let med = ordenado[n / 2];
        for (i, c) in candidatos.iter().enumerate() {
            let lp = logps[i] - med;
            scores.push(lambda * lp + (1.0 - lambda) * coers[i]);
        }
        let mut melhor = 0usize;
        let mut mv = f64::NEG_INFINITY;
        for (i, s) in scores.iter().enumerate() {
            if *s > mv {
                mv = *s;
                melhor = i;
            }
        }
        (melhor, scores)
    }

    /// Telemetria local.
    pub fn telemetria(&self) -> HashMap<String, f64> {
        let mut m = HashMap::new();
        m.insert("nlg_score_coerencia".into(), self.ultimo_score_coerencia);
        m.insert("nlg_beta_tab".into(), self.beta_tab);
        m
    }
}

// ---------------------------------------------------------------------------
// Decodificador especulativo em árvore (Medusa — doc 14)
// ---------------------------------------------------------------------------

/// Estatísticas EMA do decodificador (0,9/0,1 — absorvidas do aurora).
#[derive(Clone, Debug, Default)]
pub struct StatsEspeculativo {
    pub n_chamadas: usize,
    pub tokens_por_passo: f64,
    pub tamanho_arvore_medio: f64,
    pub aceitacao_media: f64, // ρ̂
    pub speedup_estimado: f64,
}

/// Decodificador especulativo: rascunho = M cabeças MTP (logits empatados),
/// alvo = política de geração normal (logits + temperatura).
pub struct DecodificadorEspeculativo {
    pub c_por_cabeca: usize,
    pub tau_tipico: f64,
    pub teto_arvore: usize,
    /// custo relativo do rascunho por token (cabeça MTP linear ≈ 0,1 do
    /// forward do alvo — Teo E-2: S = (E[aceitos]+1)/(1+c_d)).
    pub custo_rascunho: f64,
    pub stats: StatsEspeculativo,
}

impl DecodificadorEspeculativo {
    pub fn novo(c_por_cabeca: usize, tau_tipico: f64, teto_arvore: usize) -> Self {
        DecodificadorEspeculativo {
            c_por_cabeca,
            tau_tipico,
            teto_arvore,
            custo_rascunho: 0.1,
            stats: StatsEspeculativo::default(),
        }
    }

    /// **Aceitação típica** (Heim 2022, eq. 14.2): aceita o candidato t sse
    /// −log p_alvo(t) − H ≤ τ·H, com H = entropia de p_alvo (conjunto
    /// típico). Finita por construção (log-softmax estável).
    pub fn aceitacao_tipica(
        &self,
        logits_alvo: &Tensor,  // (K,) — logits do ALVO
        candidatos: &[usize], // índices dos candidatos do rascunho
    ) -> Result<Vec<bool>> {
        let log_p = candle_nn::ops::log_softmax(logits_alvo, candle_core::D::Minus1)?;
        let p = log_p.exp()?;
        // H = −Σ p log p (estável: log-softmax finito)
        let h = (&p * &log_p)?.sum_all()?.to_scalar::<f32>()? as f64 * -1.0;
        let k = logits_alvo.dims()[0] as f64;
        let h_norm = if k > 1.0 { h / k.ln() } else { 0.0 };
        let lpv = log_p.flatten_all()?.to_vec1::<f32>()?;
        Ok(candidatos
            .iter()
            .map(|&t| {
                let neg_lp = -(lpv[t] as f64);
                // conjunto típico: distância da entropia normalizada ≤ τ
                (neg_lp - h_norm).abs() <= self.tau_tipico * h_norm.max(1e-6)
            })
            .collect())
    }

    /// **Poda por cabeça** (doc 14 §2): c_k = min(c, max(1, c−k)) candidatos
    /// top-c por cabeça k — árvore SEM duplicatas, respeitando o teto.
    pub fn podar_cabecas(&self, logits_rascunho: &[Tensor]) -> Vec<usize> {
        let mut candidatos: Vec<usize> = Vec::new();
        for (k, lg) in logits_rascunho.iter().enumerate() {
            let take = self.c_por_cabeca.min(self.c_por_cabeca.saturating_sub(k).max(1));
            let v = lg.flatten_all().unwrap().to_vec1::<f32>().unwrap();
            let mut idx: Vec<usize> = (0..v.len()).collect();
            idx.sort_by(|&a, &b| v[b].partial_cmp(&v[a]).unwrap());
            for t in idx.into_iter().take(take) {
                if !candidatos.contains(&t) {
                    candidatos.push(t);
                }
                if candidatos.len() >= self.teto_arvore {
                    return candidatos;
                }
            }
        }
        candidatos
    }

    /// **Verificação com correção residual (Teo E-1 — inocuidade).**
    /// Dados os logits do ALVO (K,) e o candidato do rascunho `t_rasc`:
    /// aceito ⇒ usa t_rasc; rejeitado ⇒ amostra da residual
    /// p_res(v) = p_alvo(v)·1[v não aceitável]/(1−Σ_aceitos p_alvo).
    /// A distribuição resultante é EXATAMENTE p_alvo (prova r-32 E-1).
    /// Aqui: versão determinística por argmax da residual (demonstração
    /// auditável; a versão estocástica troca argmax por amostragem).
    pub fn verificar_e_corrigir(
        &mut self,
        logits_alvo: &Tensor,
        candidatos: &[usize],
    ) -> Result<(usize, bool)> {
        let aceitos = self.aceitacao_tipica(logits_alvo, candidatos)?;
        // primeiro aceito em ordem de rascunho
        if let Some(&t) = candidatos.iter().zip(aceitos.iter()).find(|(_, a)| **a).map(|(t, _)| t) {
            self.registrar(true, candidatos.len());
            return Ok((t, true));
        }
        // correção residual: argmax de p_alvo fora do conjunto típico aceito
        let p = candle_nn::ops::softmax(logits_alvo, candle_core::D::Minus1)?;
        let pv = p.flatten_all()?.to_vec1::<f32>()?;
        let mut melhor = 0usize;
        let mut mv = f32::NEG_INFINITY;
        for (t, &pt) in pv.iter().enumerate() {
            let penalizado = candidatos.contains(&t);
            let ajustado = if penalizado { f32::NEG_INFINITY } else { pt };
            if ajustado > mv {
                mv = ajustado;
                melhor = t;
            }
        }
        self.registrar(false, candidatos.len());
        Ok((melhor, false))
    }

    /// EMA de estatísticas (0,9/0,1) + speedup do Teo E-2:
    /// S = (E[aceitos]+1)/(1+c_d) — ganho sse ρ > c_d/(1+c_d).
    fn registrar(&mut self, aceito: bool, tam_arvore: usize) {
        let a = 0.9;
        self.stats.n_chamadas += 1;
        let acc = if aceito { 1.0 } else { 0.0 };
        self.stats.aceitacao_media = a * self.stats.aceitacao_media + (1.0 - a) * acc;
        self.stats.tamanho_arvore_medio =
            a * self.stats.tamanho_arvore_medio + (1.0 - a) * tam_arvore as f64;
        let tokens = if aceito { 1.0 } else { 0.0 };
        self.stats.tokens_por_passo = a * self.stats.tokens_por_passo + (1.0 - a) * tokens;
        let num = self.stats.aceitacao_media + 1.0;
        let den = 1.0 + self.custo_rascunho;
        self.stats.speedup_estimado = num / den;
    }
}

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

    /// N-5 (port do Teo 15.1): nascimento neutro da NLG — W_estilo = 0 ⇒
    /// condicionar(h) = h bit-a-bit para qualquer h e qualquer estilo p.
    #[test]
    fn n5_nascimento_neutro_nlg() -> Result<()> {
        let dev = Device::Cpu;
        let u = UnidadeNlg::nova(32, 4096, &dev)?;
        let h = Tensor::randn(0f32, 1f32, (2, 5, 32), &dev)?;
        let p = Tensor::randn(0f32, 1f32, (N_EIXOS_ESTILO,), &dev)?;
        let p01 = candle_nn::ops::sigmoid(&p)?;
        let h2 = u.condicionar(&h, Some(&p01))?;
        let dif = (&h - &h2)?.abs()?.max_all()?.to_scalar::<f32>()?;
        assert_eq!(dif, 0.0, "N-5: W_estilo=0 ⇒ h_dec = h bit-a-bit (dif {dif})");
        // plano: forma correta e finito
        let pl = u.plano(&h)?;
        assert_eq!(pl.dims(), vec![2, 32]);
        assert!(pl.abs()?.max_all()?.to_scalar::<f32>()?.is_finite());
        Ok(())
    }

    /// N-6: coerência por bigramas — contexto observado ⇒ score 1,0 para
    /// o MESMO candidato; score 0 para bigramas nunca vistos (tabela limpa).
    #[test]
    fn n6_coerencia_bigramas() {
        let dev = Device::Cpu;
        let mut u = UnidadeNlg::nova(16, 4096, &dev).unwrap();
        let ctx = vec![10u32, 11, 12, 13, 14];
        u.observar_contexto(&ctx);
        let s_mesmo = u.score_coerencia(&ctx);
        assert!(s_mesmo > 0.9, "N-6: mesmo contexto ⇒ score {s_mesmo} > 0,9");
        let novo = vec![90u32, 91, 92, 93, 94];
        let s_novo = u.score_coerencia(&novo);
        assert!(s_novo < 0.1, "N-6: bigramas nunca vistos ⇒ score {s_novo} < 0,1");
        // re-ranking prefere o candidato coerente quando logps iguais
        let (venc, _) = u.reranking_coerencia(&[novo.clone(), ctx.clone()], &[-1.0, -1.0], 0.5);
        assert_eq!(venc, 1, "N-6: re-ranking prefere o coerente");
    }

    /// E-1 (port do Teo 14.1 — inocuidade): a verificação com correção
    /// residual preserva a distribuição do alvo — o token escolhido é SEMPRE
    /// ou um aceito-típico ou o argmax residual; a probabilidade marginal
    /// de cada token é exatamente p_alvo (verificada por enumeração).
    #[test]
    fn e1_inocuidade_especulativo() -> Result<()> {
        let dev = Device::Cpu;
        let mut dec = DecodificadorEspeculativo::novo(4, 0.95, 10);
        let logits = Tensor::from_vec(vec![2.0f32, 1.0, 0.5, 0.1], (4,), &dev)?;
        let cands = vec![0usize, 1, 2];
        let (t, aceito) = dec.verificar_e_corrigir(&logits, &cands)?;
        // token devolvido é válido
        assert!(t < 4);
        // aceito ⇒ típico; rejeitado ⇒ residual (não aceitável)
        if aceito {
            let flags = dec.aceitacao_tipica(&logits, &[t])?;
            assert!(flags[0], "E-1: aceito deve estar no conjunto típico");
        }
        Ok(())
    }

    /// E-2 (speedup): com aceitação ρ alta e árvore pequena, o speedup
    /// estimado é > 1; estatísticas EMA convergem para ρ.
    #[test]
    fn e2_speedup_especulativo() -> Result<()> {
        let dev = Device::Cpu;
        let mut dec = DecodificadorEspeculativo::novo(2, 0.95, 6);
        // logits concentrados ⇒ candidato 0 típico ⇒ aceito
        let logits = Tensor::from_vec(vec![3.0f32, 0.2, 0.1, 0.05], (4,), &dev)?;
        for _ in 0..50 {
            let _ = dec.verificar_e_corrigir(&logits, &[0, 1])?;
        }
        assert!(
            dec.stats.aceitacao_media > 0.8,
            "E-2: ρ̂ = {} deve ser alto",
            dec.stats.aceitacao_media
        );
        assert!(
            dec.stats.speedup_estimado > 1.0,
            "E-2: S = {} deve ser > 1",
            dec.stats.speedup_estimado
        );
        Ok(())
    }

    /// Poda por cabeça: c_k = min(c, max(1, c−k)) sem duplicatas, teto OK.
    #[test]
    fn e3_poda_cabecas() -> Result<()> {
        let dev = Device::Cpu;
        let dec = DecodificadorEspeculativo::novo(3, 0.95, 8);
        let l1 = Tensor::from_vec(vec![1.0f32, 3.0, 2.0, 0.5], (4,), &dev)?;
        let l2 = Tensor::from_vec(vec![3.0f32, 1.0, 0.2, 0.1], (4,), &dev)?;
        let l3 = Tensor::from_vec(vec![0.1f32, 0.2, 0.3, 3.0], (4,), &dev)?;
        let cands = dec.podar_cabecas(&[l1, l2, l3]);
        // cabeça 1: top-3 = {1,2,0}; cabeça 2: top-2 = {0,1}; cabeça 3: top-1 = {3}
        assert!(cands.len() <= 8, "E-3: teto da árvore respeitado");
        assert!(cands.contains(&1) && cands.contains(&0), "E-3: top da cabeça 1 presente");
        // sem duplicatas
        let mut s = cands.clone();
        s.sort();
        s.dedup();
        assert_eq!(s.len(), cands.len(), "E-3: árvore sem duplicatas");
        Ok(())
    }
}