//! `python::fnh_arvore` — extraído do monólito na v0.10.0
//! (divisão SRP, crítica Qwen3.8 lote 3 nº 1; r-35 §2).

use super::*;

use super::nucleo::ConfigPy;
use super::treino_fases::{carregar_gerador, carregar_tokenizador};



// ---------------------------------------------------------------------------
// v0.7.0 — bindings dos módulos novos (FNH, árvore, quebra-cabeça, NLP/NLG)
// ---------------------------------------------------------------------------

/// Árvore de pesos indexada (arestas duplas) exposta ao Python.
#[pyclass(name = "ArvorePesos")]
pub(crate) struct ArvorePesosPy {
    arvore: arvore_pesos::ArvorePesos,
}

#[pymethods]
impl ArvorePesosPy {
    #[staticmethod]
    #[pyo3(signature = (contagens=None))]
    fn nova(contagens: Option<std::collections::HashMap<String, usize>>) -> Self {
        let c = contagens.unwrap_or_default();
        ArvorePesosPy { arvore: arvore_pesos::ArvorePesos::nova(&c) }
    }

    fn invariantes_ok(&self) -> bool {
        self.arvore.invariantes_ok().is_ok()
    }

    fn agregar(&mut self, metrica: &str) -> f64 {
        self.arvore.agregar_fundo_topo(metrica)
    }

    fn aplicar_ajuste(
        &mut self,
        no: &str,
        metrica: &str,
        delta: f64,
        delta_max: f64,
    ) -> PyResult<f64> {
        let idx = self
            .arvore
            .no(no)
            .ok_or_else(|| pyo3::exceptions::PyKeyError::new_err(format!("nó '{no}' inexistente")))?;
        Ok(self.arvore.aplicar_ajuste_indexado(idx, metrica, delta, delta_max))
    }

    fn definir_metrica(&mut self, no: &str, metrica: &str, valor: f64) -> PyResult<()> {
        let idx = self
            .arvore
            .no(no)
            .ok_or_else(|| pyo3::exceptions::PyKeyError::new_err(format!("nó '{no}' inexistente")))?;
        self.arvore.nos[idx].metricas.insert(metrica.to_string(), valor);
        Ok(())
    }

    fn metrica(&self, no: &str, metrica: &str) -> PyResult<f64> {
        let idx = self
            .arvore
            .no(no)
            .ok_or_else(|| pyo3::exceptions::PyKeyError::new_err(format!("nó '{no}' inexistente")))?;
        Ok(self.arvore.nos[idx].metricas.get(metrica).copied().unwrap_or(0.0))
    }

    fn regular(&mut self, origem: &str, delta: f64, lambda: f64) -> PyResult<()> {
        let idx = self
            .arvore
            .no(origem)
            .ok_or_else(|| pyo3::exceptions::PyKeyError::new_err(format!("nó '{origem}' inexistente")))?;
        self.arvore.regular_topo_fundo(idx, delta, lambda);
        Ok(())
    }

    fn dump(&self) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            let nos: Vec<serde_json::Value> = self
                .arvore
                .nos
                .iter()
                .map(|no| {
                    serde_json::json!({
                        "nome": no.nome, "ini": no.ini, "fim": no.fim,
                        "profundidade": no.profundidade,
                        "metricas": no.metricas,
                    })
                })
                .collect();
            let d = serde_json::json!({
                "n_nos": self.arvore.n_nos(),
                "n_params": self.arvore.n_params,
                "invariantes_ok": self.arvore.invariantes_ok().is_ok(),
                "nos": nos,
            });
            Ok(pythonize::pythonize(py, &d)?.into())
        })
    }
}


/// Demo/gate da camada FNH sobre o backbone: mede o efeito do gate de
/// punição nas rotações aprendíveis (ablação: π baixa vs π alta).
#[pyfunction]
#[pyo3(signature = (d_in=8, d_out=6, passos=60, seed=42))]
pub(crate) fn demo_fnh(py: Python, d_in: usize, d_out: usize, passos: usize, seed: u64) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let mut rng = rng::RngDeterministico::novo(seed);
    let mut camada = fnh::CamadaFNH::nova(d_in, d_out, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fnh: {e}")))?;
    let x: Vec<f32> = (0..4 * d_in).map(|_| rng.normal_padrao() as f32).collect();
    let xt = candle_core::Tensor::from_vec(x, (4, d_in), &dev)
        .map_err(py_err)?;
    camada
        .suavizador
        .calibrar(camada.base.w.as_tensor(), &xt)
        .map_err(py_err)?;
    // treino curto com punição ALTA (gate abre) vs BAIXA (gate fechado)
    let mut resumo = serde_json::Map::new();
    for (rotulo, punicao) in [("baixa", 0.1f64), ("alta", 0.95)] {
        let mut cam = fnh::CamadaFNH::nova(d_in, d_out, &dev).map_err(py_err)?;
        // mesma semente de pesos: copia os valores
        let _ = &mut cam;
        let mut hub = telemetria::TelemetryHub::novo(None);
        let mut perdas = Vec::new();
        for p in 0..passos {
            // ruído no x a cada passo (regime não-estacionário leve)
            let ruido: Vec<f32> = (0..4 * d_in)
                .map(|_| (rng.normal(0.0, 0.05)) as f32)
                .collect();
            let xt2 = (&xt
                + &candle_core::Tensor::from_vec(ruido, (4, d_in), &dev).map_err(py_err)?)
            .map_err(py_err)?;
            let y = cam.forward(&xt2, punicao).map_err(py_err)?;
            // alvo: saída da camada base (auto-referência: regular para W)
            let z_ref = cam.linear_sem_bias(&xt2, cam.base.w.as_tensor()).map_err(py_err)?;
            let perda = (&y).sub(&z_ref).map_err(py_err)?.sqr().map_err(py_err)?.sum_all().map_err(py_err)?;
            perdas.push(perda.to_scalar::<f32>().map_err(py_err)? as f64);
            let grads = perda.backward().map_err(py_err)?;
            // gradiente descendente manual simples nos θ das rotações
            if let Some(g) = grads.get(&cam.r_wp.theta) {
                let gv = g.flatten_all().map_err(py_err)?.to_vec1::<f32>().map_err(py_err)?;
                let th = cam.r_wp.theta.as_tensor().flatten_all().map_err(py_err)?.to_vec1::<f32>().map_err(py_err)?;
                let novo: Vec<f32> = th.iter().zip(gv.iter()).map(|(t, g_)| t - 0.05 * g_).collect();
                cam.r_wp
                    .theta
                    .set(&candle_core::Tensor::from_vec(novo, th.len(), &dev).map_err(py_err)?)
                    .map_err(py_err)?;
            }
            hub.registrar("fnh_gate", cam.gate_ultimo);
            hub.registrar("fnh_punicao", punicao);
            hub.registrar("perda", perdas[p]);
            let _ = p;
        }
        let theta_norma: f64 = cam
            .r_wp
            .theta
            .as_tensor()
            .flatten_all()
            .map_err(py_err)?
            .to_vec1::<f32>()
            .map_err(py_err)?
            .iter()
            .map(|v| (*v as f64).powi(2))
            .sum::<f64>()
            .sqrt();
        let delta_ort = cam.r_wp.erro_ortogonalidade().map_err(py_err)?;
        resumo.insert(
            format!("punicao_{rotulo}"),
            serde_json::json!({
                "gate_final": cam.gate_ultimo,
                "theta_norma": theta_norma,
                "delta_ort": delta_ort,
                "perda_inicial": perdas[0],
                "perda_final": perdas[perdas.len() - 1],
            }),
        );
    }
    let _ = &mut camada;
    let d = serde_json::json!({
        "versao": "fnh-0.7.0",
        "resumo": resumo,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}


/// **Quebra-cabeça sobre o conhecimento TREINADO REAL** (script 63):
/// carrega o tokenizador (codebook do nível `nivel` — o conhecimento
/// aprendido dos dados reais) e o gerador (a cabeça real como logits_fn
/// do verificador); executa o protocolo completo de r-30 Q-AD.
#[pyfunction]
#[pyo3(signature = (cfg, caminho_tokenizador, caminho_gerador, n_pecas=4, n_rodadas=12, tau_c=0.4, nivel=0, seed=777))]
pub(crate) fn quebra_cabeca_real(
    cfg: &ConfigPy,
    caminho_tokenizador: &str,
    caminho_gerador: &str,
    n_pecas: usize,
    n_rodadas: usize,
    tau_c: f64,
    nivel: usize,
    seed: u64,
) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let dev = candle_core::Device::Cpu;
        crate::rng::definir_semente(seed);
        let mut rng = rng::RngDeterministico::novo(seed);
        let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("tok: {e}")))?;
        carregar_tokenizador(&mut tok, caminho_tokenizador)
            .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
        let mut ger = gerador::GeradorConceitos::novo(&cfg.cfg, &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("ger: {e}")))?;
        ger.carregar(caminho_gerador)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("ger carregar: {e}")))?;
        let nivel = nivel.min(tok.quants.len().saturating_sub(1));
        // conhecimento treinado REAL: codebook Stiefel (K, d_lat)
        let codebook = tok
            .quants[nivel]
            .codebook()
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("codebook: {e}")))?;
        let d_lat = codebook.dims()[1];
        let k = ger.k;
        // δ_ort REAL do codebook treinado (rotação Stiefel em treino — A2)
        let delta_ort_codebook = tok.quants[nivel]
            .delta_ort()
            .unwrap_or(1.0);
        // logits_fn REAL: agregado (d_lat,) → proj_entrada → norm → cabeça
        let proj = &ger.proj_entrada;
        let norm = &ger.norm_entrada;
        let cabeca = &ger.cabeca.cabecas[nivel];
        let logits_fn = move |x: &candle_core::Tensor| -> candle_core::Result<candle_core::Tensor> {
            let h = proj.fwd(&x.reshape((1, d_lat))?)?;
            let h = norm.fwd(&h)?;
            let lg = cabeca.fwd(&h)?;
            lg.reshape((k,))
        };
        // verificador com temperatura calibrada pela fase C quando
        // disponível; default 1,0 (a cota V-2 aplica-se com ECE medido)
        let verif = verificador::VerificadorNcm::novo();
        let mut arvore = arvore_pesos::ArvorePesos::nova(&std::collections::HashMap::new());
        let cfg_qc = pensamento::ConfigQuebraCabeca::novo(n_pecas, n_rodadas, tau_c);
        let r = pensamento::quebra_cabeca_inovar(&codebook, &verif, &logits_fn, &mut arvore, &cfg_qc, &mut rng)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("qc: {e}")))?;
        let rodadas: Vec<serde_json::Value> = r
            .rodadas
            .iter()
            .map(|x| {
                serde_json::json!({
                    "rodada": x.rodada, "confianca": x.confianca,
                    "aceito": x.aceito, "delta": x.delta_aplicado,
                })
            })
            .collect();
        let d_ = serde_json::json!({
            "versao": "quebra-cabeca-real-0.7.0",
            "nivel": nivel,
            "delta_ort_codebook": delta_ort_codebook,
            "n_conceitos": codebook.dims()[0],
            "d_lat": d_lat,
            "aceitos": r.aceitos,
            "rejeitados": r.rejeitados,
            "confianca_base": r.confianca_base,
            "confianca_melhor": r.confianca_melhor,
            "telemetria": r.telemetria,
            "rodadas": rodadas,
            "arvore_invariantes": arvore.invariantes_ok().is_ok(),
            "punicao_no_pensamento": arvore
                .no("pensamento")
                .and_then(|i| arvore.nos[i].metricas.get("punicao").copied())
                .unwrap_or(0.0),
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}


/// Demo/gate do quebra-cabeça sobre conceitos SINTÉTICOS estruturados
/// (validação de mecanismo; a avaliação com dados reais vive no script 63).
#[pyfunction]
#[pyo3(signature = (m=24, d=8, n_pecas=4, n_rodadas=12, tau_c=0.35, seed=777))]
pub(crate) fn demo_quebra_cabeca(
    py: Python,
    m: usize,
    d: usize,
    n_pecas: usize,
    n_rodadas: usize,
    tau_c: f64,
    seed: u64,
) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let mut rng = rng::RngDeterministico::novo(seed);
    // conceitos estruturados: cluster dominante +1 (média direcional) com
    // 25% de outliers fracos — o agregado TEM direção (sem cancelamento)
    let mut v = Vec::with_capacity(m * d);
    for i in 0..m {
        let centro = if i < m / 4 { 0.2 } else { 1.0 };
        for j in 0..d {
            v.push((centro + rng.normal(0.0, 0.25)) as f32);
        }
    }
    let conceitos = candle_core::Tensor::from_vec(v, (m, d), &dev).map_err(py_err)?;
    // logits_fn ASSIMÉTRICO: projeção fixa W_proj (rotação-sensível — a
    // agregação simétrica seria invariante sob permutação/rotação, e o
    // backbone real é exatamente o avaliador NÃO-invariente da prática).
    let mut wv = Vec::with_capacity(d * 4);
    for _ in 0..d * 4 {
        wv.push((rng.normal_padrao() as f32) * 2.5);
    }
    let w_proj = candle_core::Tensor::from_vec(wv, (d, 4), &dev).map_err(py_err)?;
    // verificador com temperatura padrão
    let verif = verificador::VerificadorNcm::novo();
    let logits_fn = |x: &candle_core::Tensor| -> candle_core::Result<candle_core::Tensor> {
        // logits = xᵀ·W_proj (4,) — sensível à direção do agregado
        x.reshape((1, d))?.matmul(&w_proj)?.reshape((4,))
    };
    let mut arvore = arvore_pesos::ArvorePesos::nova(&std::collections::HashMap::new());
    let cfg = pensamento::ConfigQuebraCabeca::novo(n_pecas, n_rodadas, tau_c);
    let r = pensamento::quebra_cabeca_inovar(&conceitos, &verif, &logits_fn, &mut arvore, &cfg, &mut rng)
        .map_err(py_err)?;
    let rodadas: Vec<serde_json::Value> = r
        .rodadas
        .iter()
        .map(|x| {
            serde_json::json!({
                "rodada": x.rodada, "confianca": x.confianca,
                "aceito": x.aceito, "delta": x.delta_aplicado,
            })
        })
        .collect();
    let d_ = serde_json::json!({
        "versao": "quebra-cabeca-0.7.0",
        "aceitos": r.aceitos,
        "rejeitados": r.rejeitados,
        "confianca_base": r.confianca_base,
        "confianca_melhor": r.confianca_melhor,
        "telemetria": r.telemetria,
        "rodadas": rodadas,
        "arvore_invariantes": arvore.invariantes_ok().is_ok(),
    });
    Ok(pythonize::pythonize(py, &d_)?.into())
}