//! Gerador de conceitos (Estágio 2 / fase B) — port de
//! `conceitos/gerador_conceitos.py` (n-04 §2.4). Fatoração (Teo B.1/B.2):
//! p_φ(c) = Π_ℓ p_φ(c_ℓ|c_{<ℓ}) com predição PARALELA intra-nível;
//! condicionamento = reconstrução cumulativa interpolada ẑ_{ℓ−1}^{(n_ℓ)} + e_ℓ
//! (função determinística de c_{<ℓ} — Teo D.1). Corrupção κ (Teo C.3) e
//! dropout CFG (Teo D.2(iii)). Cabeças MTP auxiliares (Teo D.6/T7).

use candle_core::{Device, Result, Tensor, Var, D};
use crate::camadas::{Linear, RmsNorm};
use crate::config::ConfigKhtst;
use crate::mascaras::{interp_para, mascara_block_causal};
use crate::tokenizador::TokenizadorConceitos;
use crate::tronco::{CabecaConceitos, EmbNiveis, Tronco};
use rand::prelude::*;

pub const N_PREFIXO: usize = 8;

pub struct GeradorConceitos {
    pub tronco: Tronco,
    pub emb_niveis: EmbNiveis,
    pub proj_entrada: Linear, // d_lat → d_modelo
    pub cabeca: CabecaConceitos,
    pub cabecas_mtp: Vec<Linear>, // auxiliares — Teo D.6
    pub norm_entrada: RmsNorm,
    pub mascara: Tensor,
    pub niveis: Vec<usize>,
    pub d_lat: usize,
    pub corrupcao: f64,
    pub dropout_cfg: f64,
    pub k: usize,
    pub mtp_ativo: bool,
    pub beta_mtp: f64,
}

pub struct SaidaTreino {
    pub perda: Tensor,
    pub ce_por_nivel: Vec<f64>,
    pub ce_mtp: f64,
    pub acuracia_media: f64,
    pub uncond: bool,
}

impl GeradorConceitos {
    pub fn novo(cfg: &ConfigKhtst, dev: &Device) -> Result<Self> {
        let m = &cfg.modelo;
        let n = &cfg.ncm;
        let tronco = Tronco::novo(cfg, dev)?;
        let mascara = mascara_block_causal(&n.niveis, N_PREFIXO, dev)?;
        let mut mtp = Vec::new();
        if m.mtp.ativo {
            for _ in 0..n.niveis.len().saturating_sub(1) {
                mtp.push(Linear::sem_vies(m.d_modelo, n.k_por_nivel, 0.02, dev)?);
            }
        }
        Ok(GeradorConceitos {
            tronco,
            emb_niveis: EmbNiveis::nova(cfg, dev)?,
            proj_entrada: Linear::sem_vies(n.d_lat, m.d_modelo, 0.02, dev)?,
            cabeca: CabecaConceitos::nova(cfg, dev)?,
            cabecas_mtp: mtp,
            norm_entrada: RmsNorm::novo(m.d_modelo, 1e-5, dev)?,
            mascara,
            niveis: n.niveis.clone(),
            d_lat: n.d_lat,
            corrupcao: n.corrupcao_cond,
            dropout_cfg: n.dropout_cfg,
            k: n.k_por_nivel,
            mtp_ativo: m.mtp.ativo,
            beta_mtp: m.mtp.beta_entropia,
        })
    }

    pub fn entrada_nivel_pub(
        &self,
        zhat_anterior: Option<&Tensor>,
        nivel: usize,
        b: usize,
        dev: &Device,
    ) -> Result<Tensor> {
        self.entrada_nivel(zhat_anterior, nivel, b, dev)
    }

    fn entrada_nivel(
        &self,
        zhat_anterior: Option<&Tensor>,
        nivel: usize,
        b: usize,
        dev: &Device,
    ) -> Result<Tensor> {
        let n_l = self.niveis[nivel];
        // dims()[0] = d_modelo (SAÍDA da projeção d_lat→d_modelo)
        let d = self.proj_entrada.w.as_tensor().dims()[0];
        let base = match zhat_anterior {
            Some(z) => {
                let interp = interp_para(z, n_l, dev)?;
                self.proj_entrada.fwd(&interp)?
            }
            None => Tensor::zeros((b, n_l, d), candle_core::DType::F32, dev)?,
        };
        let ids_nivel = Tensor::from_vec(vec![nivel as u32 + 1; b * n_l], (b, n_l), dev)?;
        let e_l = self.emb_niveis.tabela.fwd(&ids_nivel)?;
        self.norm_entrada.fwd(&(&base + &e_l)?)
    }

    /// Aplica a corrupção κ (por posição — Teo C.3) e devolve os códigos de uso.
    fn corromper(&self, codigos: &[Tensor], b: usize, rng: &mut StdRng, dev: &Device) -> Result<Vec<Tensor>> {
        let mut saida = Vec::with_capacity(codigos.len());
        for (i, cid) in codigos.iter().enumerate() {
            let mut c = cid.clone();
            if self.corrupcao > 0.0 {
                let plano = c.reshape((b * self.niveis[i],))?;
                let mut v = plano.flatten_all()?.to_vec1::<u32>()?;
                for x in v.iter_mut() {
                    if rng.gen::<f64>() < self.corrupcao {
                        *x = rng.gen_range(0..self.k as u32);
                    }
                }
                c = Tensor::from_vec(v, (b, self.niveis[i]), dev)?;
            }
            saida.push(c);
        }
        Ok(saida)
    }

    /// Forward de TREINO com perda (1 passada — Teo D.1(iv)): teacher forcing
    /// com códigos VERDADEIROS no alvo; condicionamento corrompido por κ;
    /// dropout CFG por exemplo (ramo não-condicionado = códigos aleatórios).
    pub fn perda_treino(
        &self,
        tok: &TokenizadorConceitos,
        codigos_verdadeiros: &[Tensor],
        rng: &mut StdRng,
    ) -> Result<SaidaTreino> {
        let dev = self.mascara.device();
        let b = codigos_verdadeiros[0].dim(0)?;
        let codigos_uso = self.corromper(codigos_verdadeiros, b, rng, dev)?;
        let uncond = rng.gen::<f64>() < self.dropout_cfg;
        let zhats = if uncond {
            let aleatorios: Vec<Tensor> = (0..codigos_uso.len())
                .map(|i| {
                    let n_l = self.niveis[i];
                    let v: Vec<u32> = (0..b * n_l).map(|_| rng.gen_range(0..self.k as u32)).collect();
                    Tensor::from_vec(v, (b, n_l), dev).unwrap()
                })
                .collect();
            tok.reconstruir_de_codigos(&aleatorios, aleatorios.len(), b, dev)?
        } else {
            tok.reconstruir_de_codigos(&codigos_uso, codigos_uso.len(), b, dev)?
        };
        // monta [prefixo; lev_1; …; lev_L]
        let mut blocos = Vec::with_capacity(self.niveis.len());
        for nivel in 0..self.niveis.len() {
            let zhat_ant = if nivel == 0 { None } else { Some(&zhats[nivel - 1]) };
            blocos.push(self.entrada_nivel(zhat_ant, nivel, b, dev)?);
        }
        let prefixo_ids = Tensor::zeros((b, N_PREFIXO), candle_core::DType::U32, dev)?;
        let prefixo = self.emb_niveis.tabela.fwd(&prefixo_ids)?;
        let mut seqs = vec![prefixo];
        seqs.extend(blocos.iter().cloned());
        let s = Tensor::cat(&seqs, 1)?;
        let mut niveis_pos = vec![0usize; N_PREFIXO];
        for (i, &n_l) in self.niveis.iter().enumerate() {
            niveis_pos.extend(std::iter::repeat(i + 1).take(n_l));
        }
        let mut resumos = Vec::new();
        let h = self.tronco.forward_treino(&s, &self.mascara, &niveis_pos, &mut resumos)?;

        let mut ces = Vec::with_capacity(self.niveis.len());
        let mut accs = Vec::new();
        let mut inicio = N_PREFIXO;
        for (i, &n_l) in self.niveis.iter().enumerate() {
            let trecho = h.narrow(1, inicio, n_l)?;
            let logits = self.cabeca.cabecas[i].fwd(&trecho)?;
            let plano = logits.reshape((b * n_l, self.k))?;
            let alvo = codigos_verdadeiros[i].reshape(b * n_l)?.to_dtype(candle_core::DType::U32)?;
            let ce = candle_nn::loss::cross_entropy(&plano, &alvo)?;
            let preds = plano.argmax(D::Minus1)?;
            let ac = preds.eq(&alvo)?.to_dtype(candle_core::DType::F32)?.sum_all()?.to_scalar::<f32>()? as f64
                / (b * n_l) as f64;
            accs.push(ac);
            ces.push(ce);
            inicio += n_l;
        }
        let l = ces.len() as f64;
        let mut perda = ces[0].affine(1.0 / l, 0.0)?;
        for ce in &ces[1..] {
            perda = (&perda + &ce.affine(1.0 / l, 0.0)?)?;
        }
        // MTP auxiliar (Teo D.6): cabeça no nível ℓ prevê códigos do nível ℓ+1
        let mut ce_mtp = 0f64;
        if self.mtp_ativo && !self.cabecas_mtp.is_empty() {
            let dev = self.mascara.device();
            let mut inicio = N_PREFIXO;
            let mut soma = None;
            for i in 0..self.cabecas_mtp.len() {
                let n_l = self.niveis[i];
                let n_prox = self.niveis[i + 1];
                let trecho = h.narrow(1, inicio, n_l)?;
                // reamostra para o tamanho do nível alvo (n_ℓ ≠ n_{ℓ+1})
                let trecho_mtp = if n_l != n_prox {
                    interp_para(&trecho, n_prox, dev)?
                } else {
                    trecho
                };
                let logits = self.cabecas_mtp[i].fwd(&trecho_mtp)?;
                let plano = logits.reshape((b * n_prox, self.k))?;
                let alvo = codigos_verdadeiros[i + 1].reshape(b * n_prox)?
                    .to_dtype(candle_core::DType::U32)?;
                let ce = candle_nn::loss::cross_entropy(&plano, &alvo)?;
                soma = Some(match soma { Some(s) => (&s + &ce)?, None => ce });
                inicio += n_l;
            }
            if let Some(s) = soma {
                let m = self.cabecas_mtp.len() as f64;
                ce_mtp = (s.affine(1.0 / m, 0.0)?).to_scalar::<f32>()? as f64;
                perda = (&perda + &s.affine(self.beta_mtp / m, 0.0)?)?;
            }
        }
        Ok(SaidaTreino {
            perda,
            ce_por_nivel: ces.iter().map(|t| t.to_scalar::<f32>().unwrap() as f64).collect(),
            ce_mtp,
            acuracia_media: accs.iter().sum::<f64>() / accs.len() as f64,
            uncond,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = self.tronco.vars();
        v.extend(self.emb_niveis.vars());
        v.extend(self.proj_entrada.vars());
        v.extend(self.cabeca.vars());
        for m in &self.cabecas_mtp { v.extend(m.vars()); }
        v.extend(self.norm_entrada.vars());
        v
    }

    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().shape().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(())
    }

    /// Carrega um checkpoint salvo por `salvar` (mesma ordem de construção
    /// de `novo` ⇒ índices de variáveis casam 1-a-1 — ajuste fino honesto).
    pub fn carregar(&mut self, caminho: &str) -> Result<()> {
        let txt = std::fs::read_to_string(caminho)
            .map_err(|e| candle_core::Error::Msg(format!("io: {e}")))?;
        let obj: serde_json::Map<String, serde_json::Value> = serde_json::from_str(&txt)
            .map_err(|e| candle_core::Error::Msg(format!("serde: {e}")))?;
        let vars = self.vars();
        for (i, v) in vars.iter().enumerate() {
            let e = obj.get(&format!("v{i}")).ok_or_else(|| {
                candle_core::Error::Msg(format!("checkpoint sem v{i}"))
            })?;
            let forma: Vec<usize> = serde_json::from_value(e["forma"].clone())
                .map_err(|e| candle_core::Error::Msg(format!("forma: {e}")))?;
            let dados: Vec<f32> = serde_json::from_value(e["dados"].clone())
                .map_err(|e| candle_core::Error::Msg(format!("dados: {e}")))?;
            if dados.len() != forma.iter().product::<usize>() {
                candle_core::bail!("forma {:?} incompatível com {} valores", forma, dados.len());
            }
            let t = Tensor::from_vec(dados, forma, v.device())?;
            v.set(&t)?;
        }
        Ok(())
    }
}