//! KHTST_RUST — refatoração Rust/PyO3 do KHTST com NCM (Next-Concept
//! Modeling), provas integradas (memoria-ncm-khtst n-02/n-03), atenção ADBA
//! (r-04) e ponte duplex (r-05). Documentação: docs/memorias/.
//!
//! Convenções do projeto (n-00 §4): nascimento neutro, leitura honesta,
//! Teos numerados com ∎, token EXCLUSIVAMENTE em HF_TOKEN apagado após uso.

use pyo3::prelude::*;

pub mod amostrador;
pub mod arvore_pesos;
pub mod atencao;
pub mod acoplador;
pub mod camadas;
pub mod config;
pub mod controlador;
pub mod diagnostico;
pub mod fnh;
pub mod gerador;
pub mod mascaras;
pub mod memoria;
pub mod rng;
pub mod jepa;
pub mod moe_redirecionavel;
pub mod perdas;
pub mod ponte;
pub mod nlg;
pub mod nlp;
pub mod pensamento;
pub mod quantizacao;
pub mod quantizador;
pub mod som4d;
pub mod ssd;
pub mod stiefel;
pub mod telemetria;
pub mod tokens;
pub mod tokenizador;
pub mod treino;
pub mod tronco;
pub mod verificador;

pub mod build_info {
    include!(concat!(env!("OUT_DIR"), "/built.rs"));
}

#[pyfunction]
fn versao() -> &'static str {
    "0.7.0 (KHTST_RUST — FNH com gate de punição + árvore de pesos indexada c/ arestas duplas + quebra-cabeça do Pensamento + NLP/NLG + decodificador especulativo + auto-ajuste SOM limitado + correções v070-1/2/3)"
}

#[pyfunction]
fn info_build() -> String {
    format!(
        "rustc={}; perfil={}; features={}; alvo={}",
        build_info::RUSTC_VERSION,
        build_info::PROFILE,
        build_info::FEATURES_STR,
        build_info::TARGET,
    )
}

/// Lint de config com sugestões por Levenshtein (diagnóstico de digitação).
#[pyfunction]
fn lint_config(py: Python, caminho: &str) -> PyResult<PyObject> {
    let txt = std::fs::read_to_string(caminho)
        .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{caminho}: {e}")))?;
    let v: serde_json::Value = serde_json::from_str(&txt)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
    let achatado = diagnostico::achatar(&v);
    let avisos = diagnostico::lint_chaves(&achatado);
    let saida: Vec<serde_json::Value> = avisos
        .iter()
        .map(|a| serde_json::json!({
            "chave": a.chave, "sugestao": a.sugestao, "causa": a.causa, "correcao": a.correcao,
        }))
        .collect();
    Ok(pythonize::pythonize(py, &serde_json::json!(saida))?.into())
}

#[pyclass(name = "Bpe")]
struct BpePy {
    bpe: tokens::Bpe,
}

#[pymethods]
impl BpePy {
    #[staticmethod]
    fn treinar(textos: Vec<String>, vocab_alvo: usize) -> PyResult<Self> {
        Ok(BpePy { bpe: tokens::Bpe::treinar(&textos, vocab_alvo) })
    }

    #[staticmethod]
    fn carregar(caminho: &str) -> PyResult<Self> {
        let txt = std::fs::read_to_string(caminho)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
        let v: serde_json::Value = serde_json::from_str(&txt)
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
        let bpe = tokens::Bpe::de_json(&v)
            .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("formato BPE inválido"))?;
        Ok(BpePy { bpe })
    }

    fn codificar(&mut self, texto: &str) -> Vec<u32> {
        self.bpe.codificar(texto)
    }

    fn decodificar(&self, ids: Vec<u32>) -> String {
        self.bpe.decodificar(&ids)
    }

    fn salvar(&self, caminho: &str) -> PyResult<()> {
        std::fs::write(caminho, serde_json::to_string(&self.bpe.para_json()).unwrap())
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{e}")))?;
        Ok(())
    }

    #[getter]
    fn n_vocab(&self) -> usize {
        self.bpe.n_vocab()
    }
}

#[pyclass(name = "Config")]
struct ConfigPy {
    cfg: config::ConfigKhtst,
    caminho: Option<String>,
}

#[pymethods]
impl ConfigPy {
    #[new]
    #[allow(never_type_fallback)]
    #[pyo3(signature = (caminho=None))]
    fn new(caminho: Option<String>) -> PyResult<Self> {
        let cfg = match &caminho {
            Some(c) => {
                let txt = std::fs::read_to_string(c)
                    .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{c}: {e}")))?;
                let v: serde_json::Value = serde_json::from_str(&txt)
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("JSON inválido: {e}")))?;
                serde_json::from_value(v)
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("config: {e}")))?
            }
            None => config::ConfigKhtst::default(),
        };
        cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)?;
        Ok(ConfigPy { cfg, caminho })
    }

    fn validar(&self) -> PyResult<()> {
        self.cfg.validar().map_err(pyo3::exceptions::PyValueError::new_err)
    }

    fn checar_orcamento_cache(&self) -> PyResult<usize> {
        self.cfg.checar_orcamento_cache().map_err(pyo3::exceptions::PyValueError::new_err)
    }

    fn resumo(&self) -> String {
        format!(
            "modelo(vocab={}, d={}, camadas={}, ctx={}) ncm(L={}, niveis={:?}, K={}, d_lat={}, modo={}, ativo={})",
            self.cfg.modelo.vocab, self.cfg.modelo.d_modelo, self.cfg.modelo.n_camadas,
            self.cfg.modelo.comprimento_ctx, self.cfg.l_niveis(), self.cfg.ncm.niveis,
            self.cfg.ncm.k_por_nivel, self.cfg.ncm.d_lat, self.cfg.ncm.modo_quantizacao, self.cfg.ncm.ativo,
        )
    }

    #[getter]
    fn l_niveis(&self) -> usize {
        self.cfg.l_niveis()
    }
}

#[pyclass(name = "PonteDupla")]
struct PonteDuplaPy {
    ponte: ponte::PonteDupla,
}

#[pymethods]
impl PonteDuplaPy {
    #[new]
    #[pyo3(signature = (capacidade=64, timeout_ms=2000))]
    fn new(capacidade: usize, timeout_ms: u64) -> Self {
        PonteDuplaPy { ponte: ponte::PonteDupla::nova(capacidade, timeout_ms) }
    }

    fn rodada(&self, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<(Vec<f32>, Vec<f32>)> {
        let (ma, mb) = self.ponte.rodada(carga_a, carga_b, nivel)
            .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
        Ok((ma.carga, mb.carga))
    }

    fn demo_professor_aluno(&self, rodadas: usize, dim: usize) -> PyResult<PyObject> {
        let inicio: Vec<f32> = (0..dim).map(|i| (i as f32) * 0.01).collect();
        let stats = ponte::demo_professor_aluno(&self.ponte, rodadas, inicio, &|a, _| {
            a.iter().map(|x| x * 1.05).collect()
        })
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
        Python::with_gil(|py| {
            Ok(pythonize::pythonize(py, &serde_json::to_value(&stats).unwrap())?.into())
        })
    }

    fn estatisticas(&self) -> PyResult<PyObject> {
        let s = self.ponte.estatisticas();
        Python::with_gil(|py| {
            Ok(pythonize::pythonize(py, &serde_json::to_value(&s).unwrap())?.into())
        })
    }
}

#[cfg(feature = "ponte_async")]
#[pyclass(name = "PonteAsync")]
struct PonteAsyncPy {
    interna: std::sync::Arc<ponte::assincrona::PonteAsync>,
}

#[cfg(feature = "ponte_async")]
#[pymethods]
impl PonteAsyncPy {
    #[new]
    #[pyo3(signature = (capacidade=64, timeout_ms=2000))]
    fn new(capacidade: usize, timeout_ms: u64) -> Self {
        PonteAsyncPy { interna: std::sync::Arc::new(ponte::assincrona::PonteAsync::nova(capacidade, timeout_ms)) }
    }

    /// Corrotina: `await ponte.trocar(a, b, nivel)` — duplex simultânea.
    fn trocar<'p>(&self, py: Python<'p>, carga_a: Vec<f32>, carga_b: Vec<f32>, nivel: usize) -> PyResult<Bound<'p, PyAny>> {
        let p = self.interna.clone();
        pyo3_async_runtimes::tokio::future_into_py(py, async move {
            let (ma, mb) = p.rodada_async(carga_a, carga_b, nivel)
                .await
                .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
            Ok::<_, PyErr>((ma.carga, mb.carga))
        })
    }
}

/// Treina a fase A (tokenizador). Devolve dict com resumo + telemetria.
#[pyfunction]
#[pyo3(signature = (cfg, sequencias, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
fn treinar_fase_a(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_telemetria: Option<String>,
    caminho_checkpoint: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb: treino::CallbackPy = callback.map(|c| c.unbind());
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    let mut resumo = treino::treinar_fase_a(&cfg.cfg, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase A: {e}")))?;
    let _ = &mut resumo;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
                "perda_final": resumo.perda_final, "acc_recon": resumo.acc_recon,
                "delta_ort_max": resumo.delta_ort_max,
                "perplexidade_codebook": resumo.perplexidade_codebook,
                "monotonia_prefixo": resumo.monotonia_prefixo,
                "escore_aprendizado": resumo.escore_aprendizado,
            },
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}

/// Treina a fase B (gerador) com tokenizador congelado carregado de checkpoint.
#[pyfunction]
#[pyo3(signature = (cfg, sequencias, caminho_tokenizador, caminho_telemetria=None, caminho_checkpoint=None, callback=None, teto_seg=None))]
fn treinar_fase_b(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_tokenizador: &str,
    caminho_telemetria: Option<String>,
    caminho_checkpoint: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb: treino::CallbackPy = callback.map(|c| c.unbind());
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    // carrega o tokenizador congelado
    let dev = candle_core::Device::Cpu;
    let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("construção: {e}")))?;
    carregar_tokenizador(&mut tok, caminho_tokenizador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let resumo = treino::treinar_fase_b(&cfg.cfg, &tok, &dados, &mut hub, &cb, caminho_checkpoint.as_deref(), teto_seg)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase B: {e}")))?;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos, "perda_inicial": resumo.perda_inicial,
                "perda_final": resumo.perda_final, "acc_conceitos": resumo.acc_recon,
                "escore_aprendizado": resumo.escore_aprendizado,
            },
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}

fn carregar_tokenizador(tok: &mut tokenizador::TokenizadorConceitos, caminho: &str) -> Result<(), String> {
    let txt = std::fs::read_to_string(caminho).map_err(|e| format!("{e}"))?;
    let v: serde_json::Value = serde_json::from_str(&txt).map_err(|e| format!("{e}"))?;
    let obj = v.as_object().ok_or("formato de checkpoint inválido")?;
    let vars = tok.vars();
    for (i, var) in vars.iter().enumerate() {
        if let Some(entrada) = obj.get(&format!("v{i}")) {
            let forma: Vec<usize> = entrada["forma"].as_array().map(|a| a.iter().filter_map(|x| x.as_u64().map(|y| y as usize)).collect()).unwrap_or_default();
            let dados: Vec<f32> = entrada["dados"].as_array().map(|a| a.iter().filter_map(|x| x.as_f64().map(|y| y as f32)).collect()).unwrap_or_default();
            if forma.is_empty() || dados.is_empty() {
                continue;
            }
            let dev = candle_core::Device::Cpu;
            let t2 = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
                .and_then(|t| t.reshape(forma))
                .map_err(|e| format!("{e}"))?;
            var.set(&t2).map_err(|e| format!("{e}"))?;
        }
    }
    Ok(())
}

/// Gera amostras com NFE = L e devolve textos decodificados (modo conceitos).
#[pyfunction]
#[pyo3(signature = (cfg, caminho_tokenizador, caminho_gerador, n_amostras=2, parar_no_nivel=None, greedy=true))]
fn gerar_amostras(
    cfg: &ConfigPy,
    caminho_tokenizador: &str,
    caminho_gerador: &str,
    n_amostras: usize,
    parar_no_nivel: Option<usize>,
    greedy: bool,
) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{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!("{e}")))?;
    carregar_gerador(&mut ger, caminho_gerador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let mut amostrador = amostrador::AmostradorConceitos::novo(
        &mut ger, &tok, cfg.cfg.ncm.cfg_scale, cfg.cfg.ncm.temperaturas.limites, cfg.cfg.ncm.temperaturas.entropia_alvo,
    );
    let mut rng = <rand::prelude::StdRng as rand::SeedableRng>::seed_from_u64(4242);
    let inicio = std::time::Instant::now();
    let (codigos, nfe) = amostrador.gerar(n_amostras.max(1), &dev, &mut rng, parar_no_nivel, greedy)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("amostragem: {e}")))?;
    let tempo_ms = inicio.elapsed().as_secs_f64() * 1000.0;
    // decodifica: ẑ final → decoder AR greedy → tokens
    let nivel_final = codigos.len();
    let zhats = tok.reconstruir_de_codigos(&codigos, nivel_final, n_amostras.max(1), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let textos_ids = decodificar_textos(&tok, zhats.last(), n_amostras.max(1))
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    Python::with_gil(|py| {
        let d = serde_json::json!({
            "nfe": nfe, "l": cfg.cfg.l_niveis(), "tempo_ms": tempo_ms,
            "niveis_usados": nivel_final, "ids_tokens": textos_ids,
            "temperaturas": amostrador.temperaturas,
        });
        Ok(pythonize::pythonize(py, &d)?.into())
    })
}

fn decodificar_textos(tok: &tokenizador::TokenizadorConceitos, zhat: Option<&candle_core::Tensor>, b: usize) -> Result<Vec<Vec<u32>>, String> {
    let Some(z) = zhat else { return Ok(vec![vec![]; b]) };
    tok.decodificar_livre(z, 48).map_err(|e| format!("{e}"))
}

fn carregar_gerador(ger: &mut gerador::GeradorConceitos, caminho: &str) -> Result<(), String> {
    let txt = std::fs::read_to_string(caminho).map_err(|e| format!("{e}"))?;
    let v: serde_json::Value = serde_json::from_str(&txt).map_err(|e| format!("{e}"))?;
    let obj = v.as_object().ok_or("formato inválido")?;
    let vars = ger.vars();
    for (i, var) in vars.iter().enumerate() {
        if let Some(entrada) = obj.get(&format!("v{i}")) {
            let forma: Vec<usize> = entrada["forma"].as_array().map(|a| a.iter().filter_map(|x| x.as_u64().map(|y| y as usize)).collect()).unwrap_or_default();
            let dados: Vec<f32> = entrada["dados"].as_array().map(|a| a.iter().filter_map(|x| x.as_f64().map(|y| y as f32)).collect()).unwrap_or_default();
            if forma.is_empty() || dados.is_empty() {
                continue;
            }
            let dev = candle_core::Device::Cpu;
            let t = candle_core::Tensor::from_vec(dados, forma.clone(), &dev)
                .and_then(|t| t.reshape(forma))
                .map_err(|e| format!("{e}"))?;
            var.set(&t).map_err(|e| format!("{e}"))?;
        }
    }
    Ok(())
}

/// Utilitários de prova expostos ao Python (V-série da suíte).
#[pyfunction]
fn verificar_mascara(cfg: &ConfigPy) -> PyResult<usize> {
    let dev = candle_core::Device::Cpu;
    let m = mascaras::mascara_block_causal(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    mascaras::contar_vazamentos(&cfg.cfg.ncm.niveis, gerador::N_PREFIXO, &m)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

#[pyfunction]
fn verificar_stiefel(k: usize, d: usize) -> PyResult<f64> {
    let dev = candle_core::Device::Cpu;
    let c = stiefel::init_stiefel(k, d, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    stiefel::erro_ortogonalidade(&c.transpose(0, 1).unwrap()).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

// ===========================================================================
// v0.5.0 — suíte SOM 4D, suíte JEPA, Pensamento, telemetria robusta
// ===========================================================================

fn tensor_de_2d(dados: Vec<Vec<f64>>, dev: &candle_core::Device) -> candle_core::Result<candle_core::Tensor> {
    let d = dados.first().map(|r| r.len()).unwrap_or(1).max(1);
    let n = dados.len();
    let flat: Vec<f32> = dados.into_iter().flat_map(|r| r.into_iter().map(|x| x as f32).collect::<Vec<_>>()).collect();
    candle_core::Tensor::from_vec(flat, (n, d), dev)
}

/// Treina um 4D-SOM e devolve telemetria completa (EQ/TE/massa/cobertura).
#[pyfunction]
fn treinar_som4d(py: Python, dados: Vec<Vec<f64>>, cfg: Option<String>) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let t = tensor_de_2d(dados, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut c = som4d::CfgSom4D::default();
    if let Some(js) = &cfg {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(js) {
            if let Some(o) = v.as_object() {
                for (k, val) in o {
                    match k.as_str() {
                        "nx" => c.nx = val.as_u64().unwrap_or(4) as usize,
                        "ny" => c.ny = val.as_u64().unwrap_or(4) as usize,
                        "nz" => c.nz = val.as_u64().unwrap_or(4) as usize,
                        "nt" => c.nt = val.as_u64().unwrap_or(4) as usize,
                        "epocas" => c.epocas = val.as_u64().unwrap_or(50) as usize,
                        "sigma0" => c.sigma0 = val.as_f64().unwrap_or(2.0),
                        "sigma1" => c.sigma1 = val.as_f64().unwrap_or(0.5),
                        "eta0" => c.eta0 = val.as_f64().unwrap_or(0.5),
                        "eta1" => c.eta1 = val.as_f64().unwrap_or(0.05),
                        "semente" => c.semente = val.as_u64().unwrap_or(42),
                        "espalhada" => {
                            c.inicializacao = if val.as_bool().unwrap_or(false) {
                                som4d::Inicializacao4D::Espalhada
                            } else {
                                som4d::Inicializacao4D::Dados
                            }
                        }
                        _ => {}
                    }
                }
            }
        }
    }
    let som = som4d::treinar_som4d(&t, &c, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let d = serde_json::json!({
        "versao": som4d::versao(),
        "celulas": som.n_celulas(),
        "grade": [som.nx, som.ny, som.nz, som.nt],
        "eq_inicial": som.eq_inicial,
        "eq_final": som.eq_final,
        "te_final": som.te_final,
        "historico_eq_te": som.historia,
        "massa_soma": som.massa.iter().sum::<f64>(),
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Treina um GNG e devolve telemetria (vértices, arestas, componentes, EQ).
#[pyfunction]
fn treinar_gng(py: Python, dados: Vec<Vec<f64>>, cfg: Option<String>) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let t = tensor_de_2d(dados, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut c = som4d::CfgGng::default();
    if let Some(js) = &cfg {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(js) {
            if let Some(o) = v.as_object() {
                for (k, val) in o {
                    match k.as_str() {
                        "passos" => c.passos = val.as_u64().unwrap_or(4000) as usize,
                        "max_vertices" => c.max_vertices = val.as_u64().unwrap_or(60) as usize,
                        "lambda" => c.lambda = val.as_u64().unwrap_or(50) as usize,
                        "eps_b" => c.eps_b = val.as_f64().unwrap_or(0.05),
                        "eps_n" => c.eps_n = val.as_f64().unwrap_or(0.006),
                        "idade_max" => c.idade_max = val.as_u64().unwrap_or(50) as u32,
                        "semente" => c.semente = val.as_u64().unwrap_or(42),
                        _ => {}
                    }
                }
            }
        }
    }
    let g = som4d::treinar_gng(&t, &c)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let d = serde_json::json!({
        "versao": som4d::versao(),
        "vertices": g.vertices.len(),
        "arestas": g.arestas.len(),
        "componentes_conexas": g.n_componentes_conexas(),
        "insercoes": g.insercoes,
        "eq_medio": g.eq_medio,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Treina um H-SOM e devolve a decomposição hierárquica (Teo H-1).
#[pyfunction]
fn treinar_hsom(py: Python, dados: Vec<Vec<f64>>, cfg: Option<String>) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let t = tensor_de_2d(dados, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut c = som4d::CfgHSom::default();
    if let Some(js) = &cfg {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(js) {
            if let Some(o) = v.as_object() {
                if let Some(x) = o.get("celulas_finas") {
                    c.celulas_finas = x.as_u64().unwrap_or(9) as usize;
                }
                if let Some(x) = o.get("epocas_finas") {
                    c.epocas_finas = x.as_u64().unwrap_or(20) as usize;
                }
            }
        }
    }
    let h = som4d::treinar_hsom(&t, &c, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let cota = h.eq_grosso + 2.0 * h.r_fino + 2.0 * (h.eq_grosso * h.r_fino).sqrt();
    let d = serde_json::json!({
        "versao": som4d::versao(),
        "eq_hierarquico": h.eq_hierarquico,
        "eq_grosso": h.eq_grosso,
        "r_fino": h.r_fino,
        "cota_teo_h1": cota,
        "cota_respeitada": h.eq_hierarquico <= cota + 1e-9,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Demo T-SOM: streaming com deriva controlada — rastreio (Teo T-1).
#[pyfunction]
fn demo_tsom(py: Python, dados_iniciais: Vec<Vec<f64>>, janela: Vec<Vec<f64>>, passos: usize, gamma: f64) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let t0 = tensor_de_2d(dados_iniciais, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let cfg = som4d::CfgTSom {
        cfg_grade: som4d::CfgSom4D { nx: 3, ny: 3, nz: 2, nt: 2, epocas: 15, ..Default::default() },
        gamma,
        janela: janela.len(),
    };
    let mut ts = som4d::TSom4D::novo(&t0, &cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut historico = Vec::new();
    for _ in 0..passos {
        let eq = ts.passo(&janela);
        historico.push(eq);
    }
    let d = serde_json::json!({
        "versao": som4d::versao(),
        "gamma": gamma,
        "eq_temporal_historico": historico,
        "t_passos": ts.t,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Treina uma variante JEPA (7 variantes por máscara; lejepa/mcjepa/bjepa
/// têm funções próprias) sobre patches (n, d) com grade dada.
#[pyfunction]
fn treinar_jepa(py: Python, dados: Vec<Vec<f64>>, grid: Vec<usize>, variante: String, cfg: Option<String>) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let t = tensor_de_2d(dados, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut c = jepa::CfgJepa::default();
    if let Some(js) = &cfg {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(js) {
            if let Some(o) = v.as_object() {
                for (k, val) in o {
                    match k.as_str() {
                        "d_emb" => c.d_emb = val.as_u64().unwrap_or(32) as usize,
                        "epocas" => c.epocas = val.as_u64().unwrap_or(60) as usize,
                        "lr" => c.lr = val.as_f64().unwrap_or(3e-3),
                        "m" => c.m = val.as_f64().unwrap_or(0.99),
                        "lambda_vc" => c.lambda_vc = val.as_f64().unwrap_or(0.5),
                        "gamma_vc" => c.gamma_vc = val.as_f64().unwrap_or(1.0),
                        "frac_mascara" => c.frac_mascara = val.as_f64().unwrap_or(0.35),
                        "semente" => c.semente = val.as_u64().unwrap_or(42),
                        _ => {}
                    }
                }
            }
        }
    }
    let var = match variante.as_str() {
        "image" => jepa::VarianteJepa::ImageJepa,
        "vjepa2" => jepa::VarianteJepa::VJepa2,
        "audio" => jepa::VarianteJepa::AudioJepa,
        "point" => jepa::VarianteJepa::PointJepa,
        "jepa3d" => jepa::VarianteJepa::Jepa3D,
        "brain" => jepa::VarianteJepa::BrainJepa,
        "tjepa" => jepa::VarianteJepa::TJepa,
        _ => jepa::VarianteJepa::TJepa,
    };
    let grid2 = if grid.is_empty() { vec![t.dim(0).unwrap_or(1)] } else { grid };
    let rel = jepa::treinar_variante(
        &t,
        &|_n, r| jepa::mascara_por_variante(var, &grid2, c.frac_mascara, r),
        &c,
        var,
        None,
        &dev,
    )
    .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let d = serde_json::json!({
        "versao": jepa::versao(),
        "variante": rel.variante,
        "perda_inicial": rel.perda_inicial,
        "perda_final": rel.perda_final,
        "perda_pred_inicial": rel.perda_pred_inicial,
        "perda_pred_final": rel.perda_pred_final,
        "queda_pct": if rel.perda_pred_inicial.is_finite() && rel.perda_pred_inicial > 0.0 {
            100.0 * (1.0 - rel.perda_pred_final / rel.perda_pred_inicial)
        } else { 0.0 },
        "std_emb_inicial": rel.std_emb_inicial,
        "std_emb_final": rel.std_emb_final,
        "colapso_evitado": rel.std_emb_final > 0.05,
        "n_patches": rel.n_patches,
        "epocas": rel.epocas,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Treina LeJEPA (edição-condicionado): dados + instruções + editados.
#[pyfunction]
fn treinar_lejepa(py: Python, dados: Vec<Vec<f64>>, edits: Vec<Vec<f64>>, editados: Vec<Vec<f64>>, cfg: Option<String>) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let t = tensor_de_2d(dados, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let e = tensor_de_2d(edits, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let y = tensor_de_2d(editados, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut c = jepa::CfgJepa::default();
    if let Some(js) = &cfg {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(js) {
            if let Some(o) = v.as_object() {
                if let Some(x) = o.get("epocas") { c.epocas = x.as_u64().unwrap_or(60) as usize; }
                if let Some(x) = o.get("d_emb") { c.d_emb = x.as_u64().unwrap_or(32) as usize; }
                if let Some(x) = o.get("lr") { c.lr = x.as_f64().unwrap_or(3e-3); }
                if let Some(x) = o.get("semente") { c.semente = x.as_u64().unwrap_or(42); }
            }
        }
    }
    let rel = jepa::treinar_lejepa(&t, &e, &y, &c, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let d = serde_json::json!({
        "variante": rel.variante, "perda_inicial": rel.perda_inicial,
        "perda_final": rel.perda_final, "std_emb_final": rel.std_emb_final,
        "colapso_evitado": rel.std_emb_final > 0.05,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Treina MC-JEPA (cruz-modal): duas modalidades pareadas.
#[pyfunction]
fn treinar_mcjepa(py: Python, mod_a: Vec<Vec<f64>>, mod_b: Vec<Vec<f64>>, cfg: Option<String>) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let a = tensor_de_2d(mod_a, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let b = tensor_de_2d(mod_b, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut c = jepa::CfgJepa::default();
    if let Some(js) = &cfg {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(js) {
            if let Some(o) = v.as_object() {
                if let Some(x) = o.get("epocas") { c.epocas = x.as_u64().unwrap_or(60) as usize; }
                if let Some(x) = o.get("d_emb") { c.d_emb = x.as_u64().unwrap_or(32) as usize; }
                if let Some(x) = o.get("lr") { c.lr = x.as_f64().unwrap_or(3e-3); }
            }
        }
    }
    let rel = jepa::treinar_mcjepa(&a, &b, &c, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let d = serde_json::json!({
        "variante": rel.variante, "perda_inicial": rel.perda_inicial,
        "perda_final": rel.perda_final, "std_emb_final": rel.std_emb_final,
        "colapso_evitado": rel.std_emb_final > 0.05,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Treina BJEPA (bidirecional com consistência de ciclo) sobre sequência.
#[pyfunction]
fn treinar_bjepa(py: Python, serie: Vec<Vec<f64>>, cfg: Option<String>) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let t = tensor_de_2d(serie, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut c = jepa::CfgJepa::default();
    if let Some(js) = &cfg {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(js) {
            if let Some(o) = v.as_object() {
                if let Some(x) = o.get("epocas") { c.epocas = x.as_u64().unwrap_or(60) as usize; }
                if let Some(x) = o.get("d_emb") { c.d_emb = x.as_u64().unwrap_or(32) as usize; }
                if let Some(x) = o.get("lr") { c.lr = x.as_f64().unwrap_or(3e-3); }
            }
        }
    }
    let rel = jepa::treinar_bjepa(&t, &c, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let d = serde_json::json!({
        "variante": rel.variante, "perda_inicial": rel.perda_inicial,
        "perda_final": rel.perda_final, "std_emb_final": rel.std_emb_final,
        "colapso_evitado": rel.std_emb_final > 0.05,
    });
    Ok(pythonize::pythonize(py, &d)?.into())
}

/// Demo SonicMoE: roteamento top-k com capacidade e perda de balanceamento.
#[pyfunction]
fn demo_sonic_moe(py: Python, n: usize, d: usize, n_expertos: usize, top_k: usize, capacidade: usize) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let moe = jepa::SonicMoe::novo(d, n_expertos, top_k, capacidade, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let x = candle_core::Tensor::randn(0f32, 1.0, (n, d), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let s = moe.fwd(&x).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let aux = s.perda_aux.to_scalar::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))? as f64;
    let d_ = serde_json::json!({
        "versao": jepa::versao(),
        "carga_por_especialista": s.carga,
        "descartados": s.descartados,
        "perda_aux": aux,
        "capacidade_respeitada": s.carga.iter().all(|&c| c <= capacidade),
        "saida_forma": s.saida.dims().to_vec(),
    });
    Ok(pythonize::pythonize(py, &d_)?.into())
}

/// Demo ReplaySSM: estabilidade + replay em cenário de troca de tarefa.
#[pyfunction]
fn demo_replay_ssm(py: Python, passos: usize, alpha: f64) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let mut ssm = jepa::ReplaySsm::novo(4, 6, alpha, 8, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut otim = crate::treino::AdamWManual::novo(ssm.vars(), 1e-2, 0.0);
    let x_antigo = candle_core::Tensor::from_vec(vec![1.0f32; 16 * 4], (16, 4), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut hist = Vec::new();
    for _ in 0..passos {
        let h = ssm.scan(&x_antigo).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let alvo = h.detach();
        let (pt, pn, pr) = ssm.passo_treino(&x_antigo, &alvo, &mut otim)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        hist.push(serde_json::json!({"total": pt, "novo": pn, "replay": pr}));
    }
    let x_novo = candle_core::Tensor::from_vec(vec![-1.0f32; 16 * 4], (16, 4), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    for _ in 0..passos {
        let h = ssm.scan(&x_novo).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let alvo = h.detach();
        let _ = ssm.passo_treino(&x_novo, &alvo, &mut otim)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    }
    let h_fin = ssm.scan(&x_antigo).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mag = h_fin.sqr().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?
        .max_all().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?
        .to_scalar::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))? as f64;
    let d_ = serde_json::json!({
        "versao": jepa::versao(),
        "alpha_replay": alpha,
        "historico": hist,
        "magnitudes_finitas": mag.is_finite(),
        "magnitude_max_pos_troca": mag,
    });
    Ok(pythonize::pythonize(py, &d_)?.into())
}

/// Classificação de erro por taxonomia (diagnóstico de refatoração).
#[pyfunction]
fn classificar_erro(msg: String) -> Vec<(String, String, String)> {
    diagnostico::classificar_erro(&msg)
        .into_iter()
        .map(|(a, b, c)| (a.to_string(), b.to_string(), c.to_string()))
        .collect()
}

/// Confiança de uma distribuição (T-P1): `1 − H(softmax(z))/log K`.
#[pyfunction]
fn confianca_de_logits(logits: Vec<f32>) -> PyResult<f64> {
    let dev = candle_core::Device::Cpu;
    let k = logits.len();
    let t = candle_core::Tensor::from_vec(logits, (1, k), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    pensamento::confianca_de_logits(&t).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

/// Projeção ridge (T-P3 — montador algébrico): `r* = (WᵀW+εI)⁻¹Wᵀr`.
#[pyfunction]
fn projetar_ridge(w: Vec<Vec<f64>>, r_alvo: Vec<Vec<f64>>, eps: f64) -> PyResult<Vec<Vec<f64>>> {
    let dev = candle_core::Device::Cpu;
    let wt = tensor_de_2d(w, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let rt = tensor_de_2d(r_alvo, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let r = pensamento::projetar_ridge(&wt, &rt, eps)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let out = r.to_vec2::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    Ok(out.into_iter().map(|l| l.into_iter().map(|v| v as f64).collect()).collect())
}

/// Demo completa da tecnologia Pensamento: treino cooperativo das N redes
/// (item a) + ciclo d com verificador externo simulado + telemetria.
#[pyfunction]
fn demo_pensamento(
    py: Python,
    d_cond: usize,
    d_saida: usize,
    n_tentativas: usize,
    passos_treino: usize,
    limiar_conf: f64,
    orcamento: usize,
    semente: u64,
) -> PyResult<PyObject> {
    use candle_core::Tensor;
    let dev = candle_core::Device::Cpu;
    crate::rng::definir_semente(semente);
    let mut motor = pensamento::MotorPensamento::novo(d_cond, d_saida, n_tentativas, limiar_conf, orcamento, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut otim = crate::treino::AdamWManual::novo(motor.vars(), 5e-3, 0.0);
    let mut rng = crate::rng::RngDeterministico::novo(semente);

    // dados sintéticos: alvo = função determinística não-linear de cond (d=a+b+c)
    let n_dados = 256usize;
    let conds: Vec<Vec<f32>> = (0..n_dados)
        .map(|_| (0..d_cond).map(|_| rng.normal(0.0, 1.0) as f32).collect())
        .collect();
    let alvos: Vec<Vec<f32>> = conds
        .iter()
        .map(|c| {
            (0..d_saida)
                .map(|j| {
                    let cj = c[j % d_cond];
                    let ck = c[(j + 1) % d_cond];
                    0.5 * cj + 0.3 * cj * ck + 0.1
                })
                .collect()
        })
        .collect();
    let cond_t = Tensor::from_vec(conds.into_iter().flatten().collect(), (n_dados, d_cond), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let alvo_t = Tensor::from_vec(alvos.into_iter().flatten().collect(), (n_dados, d_saida), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;

    // item a: treino cooperativo
    let mut perda_hist = Vec::new();
    for passo in 0..passos_treino {
        let perda = motor.perda_treino(&cond_t, &alvo_t)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        perda_hist.push(perda.to_scalar::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))? as f64);
        otim.passo(crate::treino::grads_achatados(&perda, &motor.vars())
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?);
        let _ = passo;
    }

    // item c+d: consulta fora do padrão (OOD) — confiança baixa ⇒ Pensamento
    let consulta: Vec<f32> = (0..d_cond).map(|j| 3.0 + (j as f32) * 0.1).collect();
    let consulta_t = Tensor::from_vec(consulta, (d_cond,), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    // verificador externo: negativo do MSE contra o alvo verdadeiro da consulta
    let alvo_verdadeiro: Vec<f32> = (0..d_saida)
        .map(|j| {
            let cj = 3.0 + (j % d_cond) as f32 * 0.1;
            let ck = 3.0 + ((j + 1) % d_cond) as f32 * 0.1;
            0.5 * cj + 0.3 * cj * ck + 0.1
        })
        .collect();
    let alvo_verd = Tensor::from_vec(alvo_verdadeiro.clone(), (d_saida,), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let verifica = |resp: &Tensor| -> f64 {
        // confiança = 1/(1+10·MSE) — medível e causaL (só da resposta)
        let dif = (resp.broadcast_as((1, d_saida)).unwrap() - &alvo_verd.broadcast_as((1, d_saida)).unwrap()).unwrap();
        let mse = dif.sqr().unwrap().mean_all().unwrap().to_scalar::<f32>().unwrap() as f64;
        1.0 / (1.0 + 10.0 * mse)
    };
    let mut rng2 = crate::rng::RngDeterministico::novo(semente + 1);
    let res = motor.raciocinar(&consulta_t, &verifica, &mut rng2)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;

    // linha de base SEM pensamento: uma tentativa isolada, sem agregação
    let base = motor.tentativas[0].fwd(&consulta_t.broadcast_as((1, d_cond)).unwrap())
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let conf_base = verifica(&base);
    let conf_pensamento = res.confianca;

    let d_ = serde_json::json!({
        "versao": pensamento::versao(),
        "n_tentativas": n_tentativas,
        "perda_treino_inicial": perda_hist.first(),
        "perda_treino_final": perda_hist.last(),
        "treino_aprendeu": perda_hist.last().copied().unwrap_or(f64::INFINITY) < perda_hist.first().copied().map(|p| p * 0.5).unwrap_or(f64::INFINITY),
        "iteracoes_ciclo_d": res.iteracoes,
        "confianca_base_1_rede": conf_base,
        "confianca_pensamento": conf_pensamento,
        "ganho_confianca": conf_pensamento - conf_base,
        "orcamento_esgotado": res.telemetria.get("orcamento_esgotado").copied().unwrap_or(0.0),
        "telemetria": res.telemetria,
        "eventos_ciclo": res.trajetoria.iter().map(|e| match e {
            pensamento::EventoCiclo::Planejamento { iteracao, temperatura } =>
                serde_json::json!({"tipo": "planejamento", "iteracao": iteracao, "temperatura": temperatura}),
            pensamento::EventoCiclo::DistribuicaoCompetencias { por_tentativa } =>
                serde_json::json!({"tipo": "distribuicao", "pesos": por_tentativa}),
            pensamento::EventoCiclo::RePlanejamento { motivo } =>
                serde_json::json!({"tipo": "replanejamento", "motivo": motivo}),
            pensamento::EventoCiclo::Resposta { confianca, iteracoes } =>
                serde_json::json!({"tipo": "resposta", "confianca": confianca, "iteracoes": iteracoes}),
        }).collect::<Vec<_>>(),
    });
    Ok(pythonize::pythonize(py, &d_)?.into())
}

/// Extrai patches 2D de uma imagem achatada (C, H, W) — adaptador Image-JEPA.
#[pyfunction]
fn patches_2d(img: Vec<f32>, c: usize, h: usize, w: usize, ph: usize, pw: usize) -> PyResult<Vec<Vec<f32>>> {
    let dev = candle_core::Device::Cpu;
    let t = candle_core::Tensor::from_vec(img, (c, h, w), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let p = jepa::em_patches_2d(&t, ph, pw)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    p.to_vec2::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

/// Extrai patches canal-independentes de série temporal (C, T) — T-JEPA.
#[pyfunction]
fn patches_serie(ts: Vec<f32>, c: usize, t: usize, pt: usize) -> PyResult<Vec<Vec<f32>>> {
    let dev = candle_core::Device::Cpu;
    let tt = candle_core::Tensor::from_vec(ts, (c, t), &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let p = jepa::em_patches_serie(&tt, pt)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    p.to_vec2::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))
}

/// Extrai grupos FPS de nuvem de pontos (n, d) — Point-JEPA.
#[pyfunction]
fn grupos_pontos(pc: Vec<Vec<f64>>, k: usize, m: usize, semente: u64) -> PyResult<Vec<Vec<f64>>> {
    let dev = candle_core::Device::Cpu;
    let t = tensor_de_2d(pc, &dev).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let mut rng = crate::rng::RngDeterministico::novo(semente);
    let g = jepa::em_grupos_pontos(&t, k, m, &mut rng)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    let out = g.to_vec2::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
    Ok(out.into_iter().map(|l| l.into_iter().map(|v| v as f64).collect()).collect())
}

// ---------------------------------------------------------------------------
// v0.6.0 — verificador NCM (sem oráculo), fase C, MoE redirecionável,
// acoplador duplex e autoajuste
// ---------------------------------------------------------------------------

/// Converte matriz f64 → Tensor (n, d) F32.
fn tensor_f64_2d(v: &[Vec<f64>], dev: &candle_core::Device) -> candle_core::Result<candle_core::Tensor> {
    let (n, d) = (v.len(), v.first().map(|l| l.len()).unwrap_or(0));
    candle_core::Tensor::from_vec(v.iter().flatten().cloned().collect(), (n, d), dev)
}

/// **V-1..V-3:** calibra o verificador NCM em logits retidos e devolve
/// temperatura, ECE antes/depois, NLL antes/depois e o τ_c por quantil DKW.
#[pyfunction]
fn calibrar_verificador(
    logits: Vec<Vec<f32>>,
    alvos: Vec<usize>,
    epsilon: f64,
) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let dev = candle_core::Device::Cpu;
        let (n, k) = (logits.len(), logits.first().map(|l| l.len()).unwrap_or(1));
        let flat: Vec<f32> = logits.into_iter().flatten().collect();
        let lt = candle_core::Tensor::from_vec(flat, (n, k), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let at = candle_core::Tensor::from_vec(
            alvos.iter().map(|&a| a as u32).collect::<Vec<u32>>(), (n,), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut v = verificador::VerificadorNcm::novo();
        let (e_antes, e_depois) = v.ajustar(&lt, &at)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let confs: Vec<f64> = verificador::confianca_por_linha(&lt)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?
            .flatten_all().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?
            .to_vec1::<f32>().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?
            .iter().map(|x| *x as f64).collect();
        let tau_c = verificador::limiar_por_erro(&confs, epsilon);
        let d_ = serde_json::json!({
            "temperatura": v.temperatura, "ece_antes": e_antes, "ece_depois": e_depois,
            "nll_antes": v.n_ll_antes, "nll_depois": v.n_ll_depois,
            "tau_c": tau_c, "epsilon": epsilon, "n": n,
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}

/// **Demo do verificador NCM substituindo o oráculo** (r-21 §3.1): treina
/// um classificador pequeno, calibra T* em dados retidos, mede ECE, τ_c
/// (DKW) e a soundness empírica `P(erro|aceito) ≤ 1−τ_c+ECE`; compara com
/// o ORÁCULO (que usa o alvo — impossível em produção) no mesmo problema.
#[pyfunction]
fn demo_verificador_ncm(n: usize, k: usize, semente: u64) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let dev = candle_core::Device::Cpu;
        crate::rng::definir_semente(semente);
        let mut rng = crate::rng::RngDeterministico::novo(semente);
        // dados sintéticos com DRIFT (treino ≠ consulta — força OOD leve)
        let gerar = |rng: &mut crate::rng::RngDeterministico, m: usize, desloc: f64|
            -> (candle_core::Tensor, Vec<usize>) {
            let mut l = Vec::with_capacity(m * k);
            let mut a = Vec::with_capacity(m);
            for _ in 0..m {
                let alvo = rng.inteiro(k);
                a.push(alvo);
                for j in 0..k {
                    let base = if j == alvo { 2.5 } else { 0.0 };
                    l.push((base + rng.normal(0.0, 1.2) as f32 + desloc as f32) as f32);
                }
            }
            let lt = candle_core::Tensor::from_vec(l, (m, k), &dev).unwrap();
            (lt, a)
        };
        let (lt_tre, alvos_tre) = gerar(&mut rng, n, 0.0);
        let (lt_val, alvos_val) = gerar(&mut rng, n / 2, 0.3);
        // mini-classificador: matriz W (k, k) treinada por softmax-CE
        let w = candle_core::Tensor::randn(0f32, 0.1, (k, k), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let w_var = candle_core::Var::from_tensor(&w)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let at_t = candle_core::Tensor::from_vec(
            alvos_tre.iter().map(|&a| a as u32).collect::<Vec<u32>>(), (n,), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut otim = treino::AdamWManual::novo(vec![w_var.clone()], 0.05, 0.0);
        for _ in 0..120 {
            let logits = lt_tre.matmul(&w_var.as_tensor().transpose(0, 1).unwrap()).unwrap();
            let ce = candle_nn::loss::cross_entropy(&logits, &at_t).unwrap();
            let g = treino::grads_achatados(&ce, &[w_var.clone()]).unwrap();
            otim.passo(g);
        }
        // logits de validação com o modelo treinado
        let logits_val = lt_val.matmul(&w_var.as_tensor().transpose(0, 1).unwrap()).unwrap();
        let at_v = candle_core::Tensor::from_vec(
            alvos_val.iter().map(|&a| a as u32).collect::<Vec<u32>>(), (n / 2,), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        // calibração (V-1) + soundness (V-2) + quantil (V-3)
        let mut v = verificador::VerificadorNcm::novo();
        let (ece_antes, ece_depois) = v.ajustar(&logits_val, &at_v)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let confs: Vec<f64> = v.confianca(&logits_val).unwrap()
            .flatten_all().unwrap().to_vec1::<f32>().unwrap().iter().map(|x| *x as f64).collect();
        let corretos: Vec<bool> = logits_val.argmax(candle_core::D::Minus1).unwrap()
            .eq(&at_v).unwrap().to_dtype(candle_core::DType::U32).unwrap()
            .flatten_all().unwrap().to_vec1::<u32>().unwrap()
            .iter().map(|x| *x == 1).collect();
        let tau_c = verificador::limiar_por_erro(&confs, 0.25);
        let (mut erro_aceito, mut n_aceito) = (0usize, 0usize);
        for (c, ok) in confs.iter().zip(corretos.iter()) {
            if *c >= tau_c {
                n_aceito += 1;
                if !ok {
                    erro_aceito += 1;
                }
            }
        }
        let p_erro = erro_aceito as f64 / n_aceito.max(1) as f64;
        let cota = (1.0 - tau_c) + ece_depois + 0.06;
        // ORÁCULO no mesmo problema (usa o alvo — só para contraste)
        let acerto_oraculo = corretos.iter().filter(|c| **c).count() as f64 / corretos.len() as f64;
        let d_ = serde_json::json!({
            "versao": verificador::versao(), "n": n, "k": k,
            "temperatura_calibrada": v.temperatura,
            "ece_antes": ece_antes, "ece_depois": ece_depois,
            "nll_antes": v.n_ll_antes, "nll_depois": v.n_ll_depois,
            "tau_c_dkw": tau_c,
            "p_erro_aceito": p_erro, "cota_soundness": cota,
            "soundness_ok": p_erro <= cota,
            "calibracao_melhorou": ece_depois < ece_antes,
            "acerto_oraculo_referencia": acerto_oraculo,
            "verificador_sem_oraculo": true,
            "gates": {
                "gate_v1_calibracao": ece_depois < ece_antes,
                "gate_v2_soundness": p_erro <= cota,
                "gate_v3_tau_positivo": tau_c > 0.0 && tau_c < 1.0,
            },
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}

/// **FASE C — treino conjunto Pensamento↔backbone** (dados reais).
#[pyfunction]
fn treinar_fase_c(
    cfg: &ConfigPy,
    sequencias: Vec<Vec<u32>>,
    caminho_tokenizador: &str,
    passos: usize,
    caminho_gerador: Option<String>,
    caminho_telemetria: Option<String>,
    callback: Option<Bound<'_, PyAny>>,
    teto_seg: Option<f64>,
) -> PyResult<PyObject> {
    let cb: treino::CallbackPy = callback.map(|c| c.unbind());
    let mut hub = telemetria::TelemetryHub::novo(caminho_telemetria);
    let dados = treino::DadosTreino::novo(sequencias, cfg.cfg.modelo.comprimento_ctx);
    let dev = candle_core::Device::Cpu;
    let mut tok = tokenizador::TokenizadorConceitos::novo(&cfg.cfg, &dev)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("construção: {e}")))?;
    carregar_tokenizador(&mut tok, caminho_tokenizador)
        .map_err(pyo3::exceptions::PyRuntimeError::new_err)?;
    let resumo = treino::treinar_fase_c(&cfg.cfg, &tok, &dados, &mut hub, &cb, teto_seg, passos, caminho_gerador.as_deref())
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("fase C: {e}")))?;
    Python::with_gil(|py| {
        let dados = serde_json::json!({
            "resumo": {
                "passos": resumo.passos,
                "ce_inicial": resumo.ce_inicial,
                "ce_final": resumo.ce_final,
                "ce_gatilhada_sem_correcao": resumo.ce_gatilhada_sem_correcao,
                "ce_gatilhada_com_correcao": resumo.ce_gatilhada_com_correcao,
                "confianca_media_inicial": resumo.confianca_media_inicial,
                "confianca_media_final": resumo.confianca_media_final,
                "confianca_pos_correcao": resumo.confianca_pos_correcao,
                "taxa_gatilho": resumo.taxa_gatilho,
                "ece_antes": resumo.ece_antes,
                "ece_depois": resumo.ece_depois,
                "temperatura_calibrada": resumo.temperatura_calibrada,
                "tau_c": resumo.tau_c,
                "taxa_conflito_pcgrad": resumo.taxa_conflito_pcgrad,
                "decaimento_delta": resumo.decaimento_delta,
                "acordos_medios": resumo.acordos_medios,
                "escore_aprendizado": resumo.escore_aprendizado,
                "rodadas_acoplador": resumo.rodadas_acoplador,
                "timeouts_acoplador": resumo.timeouts_acoplador,
                "maior_sobreposicao_us": resumo.maior_sobreposicao_us,
            },
            "painel_v060": hub.painel_v060(),
            "telemetria": hub.dump(),
        });
        Ok(pythonize::pythonize(py, &dados)?.into())
    })
}

/// **MoE redirecionável — demo completa (M-R1..M-R5):** neutralidade da
/// atenção própria (M-R2), invariância de capacidade/perda auxiliar sob
/// redirecionamento (M-R1), paralelismo medido (M-R3), water-filling por
/// QE vs uniforme (M-R5) e fechamento da classe (M-R4 — via permutação).
#[pyfunction]
fn demo_moe_redirecionavel(d: usize, n: usize, semente: u64) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let dev = candle_core::Device::Cpu;
        crate::rng::definir_semente(semente);
        let mut rng = crate::rng::RngDeterministico::novo(semente);
        let x: Vec<f32> = (0..n * d).map(|_| rng.normal(0.0, 1.0) as f32).collect();
        let x = candle_core::Tensor::from_vec(x, (n, d), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut moe = moe_redirecionavel::MoeRedirecionavel::novo(d, 8, 2, n, &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        // M-R2: neutralidade no nascimento
        let g0 = moe.atencao.gate.as_tensor().to_scalar::<f32>().unwrap();
        let y0 = moe.atencao.fwd(&x).unwrap();
        let dif0 = (&y0 - &x).unwrap().sqr().unwrap().sum_all().unwrap()
            .to_scalar::<f32>().unwrap() as f64;
        // abre o gate (mecanismo ativo)
        moe.atencao.gate.set(&candle_core::Tensor::new(1.0f32, &dev).unwrap()).unwrap();
        // M-R1: invariância sob redirecionamento
        let s_enc = moe.fwd(&x, moe_redirecionavel::Papel::Codificador).unwrap();
        let carga_antes: usize = s_enc.carga.iter().sum();
        let perda_antes = s_enc.perda_aux.to_scalar::<f32>().unwrap() as f64;
        let permuta: Vec<usize> = (0..8).map(|i| (i + 1) % 8).collect();
        moe.redirecionar(&permuta);
        let s_enc2 = moe.fwd(&x, moe_redirecionavel::Papel::Codificador).unwrap();
        let carga_depois: usize = s_enc2.carga.iter().sum();
        let perda_depois = s_enc2.perda_aux.to_scalar::<f32>().unwrap() as f64;
        // M-R3: paralelismo medido (enc ∥ dec em 2 threads)
        let (_e, _dc, t_par, t_seq, speedup) = moe.fwd_ambos_par(&x).unwrap();
        // M-R5: water-filling vs uniforme (cota analítica)
        let qes: Vec<f64> = vec![0.15, 0.4, 1.0, 0.25, 0.6, 0.1, 0.35, 0.5];
        let c_wf = moe_redirecionavel::alocar_por_qe(&qes, 16);
        let c_un = vec![2usize; 8];
        let cota = |c: &[usize]| qes.iter().zip(c).map(|(q, ci)| q * q / *ci as f64).sum::<f64>();
        let (cota_wf, cota_un) = (cota(&c_wf), cota(&c_un));
        let d_ = serde_json::json!({
            "versao": moe_redirecionavel::versao(), "d": d, "n": n,
            "valor_gate_no_nascimento": g0, "neutralidade_dif": dif0,
            "carga_antes": carga_antes, "carga_depois": carga_depois,
            "perda_aux_antes": perda_antes, "perda_aux_depois": perda_depois,
            "t_paralelo_s": t_par, "t_sequencial_s": t_seq, "speedup": speedup,
            "alocacao_water_filling": c_wf, "cota_wf": cota_wf, "cota_uniforme": cota_un,
            "gates": {
                "gate_mr1_capacidade": carga_antes == carga_depois,
                "gate_mr1_perda_aux": (perda_antes - perda_depois).abs() < 2.0 + 1e-6,
                "gate_mr2_neutralidade": dif0 == 0.0 && g0 == 0.0,
                "gate_mr3_paralelismo": t_par < t_seq,
                "gate_mr5_water_filling": cota_wf <= cota_un + 1e-9,
            },
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}

/// **JEPA com MoE redirecionável (ablação de 3 braços):** `modo` ∈
/// {"pequena", "densa", "moe"} — o braço "densa" tem a capacidade TOTAL dos
/// especialistas SEM roteamento (isola o ganho do roteamento em si).
#[pyfunction]
fn jepa_com_moe(
    patches: Vec<Vec<f32>>,
    modo: String,
    passos: usize,
    frac_mascara: f64,
    semente: u64,
) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let dev = candle_core::Device::Cpu;
        let t = tensor_f64_2d(&patches.into_iter().map(|l| l.into_iter().map(|x| x as f64).collect()).collect::<Vec<Vec<f64>>>(), &dev)
            .and_then(|t| t.to_dtype(candle_core::DType::F32))
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let modo: &'static str = match modo.as_str() {
            "pequena" => "pequena",
            "densa" => "densa",
            _ => "moe",
        };
        let r = moe_redirecionavel::treinar_jepa_com_moe(&t, modo, passos, frac_mascara, semente)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let d_ = serde_json::json!({
            "modo": r.modo, "perda_pred_inicial": r.perda_pred_inicial,
            "perda_pred_final": r.perda_pred_final, "perda_vc_final": r.perda_vc_final,
            "std_emb": r.std_emb, "carga_max": r.carga_max, "descartados": r.descartados,
            "queda_abs": r.perda_pred_inicial - r.perda_pred_final,
            "aprendeu": r.perda_pred_final < r.perda_pred_inicial,
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}

/// **SOM refinado por MoE com water-filling (M-R5) + redirecionamento
/// para decodificador:** treina um 4D-SOM, refina com MoE (QE′ < QE),
/// compara com alocação uniforme e mede a decodificação pós-redirect.
#[pyfunction]
fn som_com_moe(
    dados: Vec<Vec<f64>>,
    lado: usize,
    passos: usize,
    semente: u64,
) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let dev = candle_core::Device::Cpu;
        let t = tensor_f64_2d(&dados, &dev)
            .and_then(|t| t.to_dtype(candle_core::DType::F32))
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        // treina o 4D-SOM (r-19: S4-1/S4-2/S4-3) — som4d opera em F32
        let cfg_som = som4d::CfgSom4D { nx: lado, ny: lado, nz: lado, nt: lado, ..Default::default() };
        let som = som4d::treinar_som4d(&t, &cfg_som, &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let centroides: Vec<Vec<f64>> = som.prototipos.clone();
        // refinamento MoE (mesmo tensor F32; centroides seguem em f64)
        let r = moe_redirecionavel::refinar_som_com_moe(&t, &centroides, passos, semente)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let d_ = serde_json::json!({
            "qe_antes": r.qe_antes, "qe_depois_water_filling": r.qe_depois,
            "qe_depois_uniforme": r.qe_alocacao_uniforme,
            "carga_por_celula": r.carga_por_celula,
            "tempo_treino_s": r.tempo_treino_s,
            "mse_decodificador_pos_redirect": r.mse_decodificador,
            "reducao_qe_pct": 100.0 * (r.qe_antes - r.qe_depois) / r.qe_antes.max(1e-12),
            "gates": {
                "gate_moe_refina": r.qe_depois < r.qe_antes,
                "gate_wf_melhor_uniforme": r.qe_depois <= r.qe_alocacao_uniforme * 1.02,
                "gate_redirect_decodifica": r.mse_decodificador < r.qe_antes.powi(2),
            },
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}

/// **AcopladorDuplo — demo do protocolo lag-1 síncrono** (AC-1/AC-2):
/// rodadas duplex Pensamento↔backbone com correções decaindo (co-adaptação)
/// e confiança subindo; zero deadlocks/timeouts (P.1/P.2 herdados).
#[pyfunction]
fn demo_acoplador_duplo(rodadas: usize, timeout_ms: u64) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let mut ac = acoplador::AcopladorDuplo::nova(16, timeout_ms);
        let mut rng = crate::rng::RngDeterministico::novo(77);
        let mut conf = 0.05f64;
        let mut traj = Vec::new();
        for k in 0..rodadas {
            let estado = acoplador::EstadoModelo {
                a: vec![0.1 * (k % 7) as f32; 16],
                b: vec![rng.normal(0.0, 0.3) as f32; 16],
                c: vec![0.05; 8],
                confianca: conf,
                nivel: k % 7,
                passo: k,
            };
            let _ = ac.rodada(estado).map_err(|e|
                pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
            // correção fresca: decaindo com a confiança (co-adaptação)
            let amplitude = (1.0 - conf).max(0.05);
            let delta: Vec<f32> = (0..16).map(|i| amplitude as f32 * ((i % 5) as f32 - 2.0) * 0.1).collect();
            ac.registrar_correcao(acoplador::Correcao {
                delta, acordo: 0.6 + 0.3 * conf, confianca_agregada: conf,
                iteracao: k, pesos: vec![0.25; 4],
            });
            traj.push(conf);
            conf = (conf + 0.04 + rng.uniforme() * 0.02).min(0.99);
        }
        let t = &ac.telemetria;
        let p = ac.estatisticas();
        let d_ = serde_json::json!({
            "versao": acoplador::versao(), "rodadas": t.rodadas,
            "timeouts": t.timeouts, "tempo_total_ms": t.tempo_total_ms,
            "soma_norma_delta": t.soma_norma_delta,
            "conf_final": traj.last(), "conf_trajetoria_len": traj.len(),
            "envios_a_para_b": p.envios_a_para_b, "envios_b_para_a": p.envios_b_para_a,
            "maior_sobreposicao_us": t.maior_sobreposicao_us,
            "gates": {
                "gate_terminou": t.rodadas == rodadas as u64,
                "gate_sem_timeout": t.timeouts == 0,
                "gate_duplex_simultaneo": p.envios_a_para_b == p.envios_b_para_a && p.envios_a_para_b == rodadas as u64,
            },
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}

/// **Ciclo d do Pensamento com o VERIFICADOR NCM** (produção, sem oráculo):
/// backbone CABEÇA TREINADA (alvo → classes) + verificador calibrado em
/// dados retidos; N tentativas agregadas; ganho medido em consultas retidas
/// EM DISTRIBUIÇÃO (onde a calibração vale — V-2) e OOD reportado à parte.
#[pyfunction]
fn demo_pensamento_ncm(
    d_cond: usize,
    d_saida: usize,
    n_tentativas: usize,
    passos_treino: usize,
    limiar_conf: f64,
    orcamento: usize,
    semente: u64,
) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        use candle_core::Tensor;
        let dev = candle_core::Device::Cpu;
        crate::rng::definir_semente(semente);
        let mut motor = pensamento::MotorPensamento::novo(
            d_cond, d_saida, n_tentativas, limiar_conf, orcamento, &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut otim = treino::AdamWManual::novo(motor.vars(), 5e-3, 0.0);
        let mut rng = crate::rng::RngDeterministico::novo(semente);
        // ---- dados sintéticos (item a: treino cooperativo) ----
        let n_dados = 320usize;
        let conds: Vec<Vec<f32>> = (0..n_dados)
            .map(|_| (0..d_cond).map(|_| rng.normal(0.0, 1.0) as f32).collect())
            .collect();
        let alvos: Vec<Vec<f32>> = conds.iter().map(|c| {
            (0..d_saida).map(|j| {
                let cj = c[j % d_cond];
                let ck = c[(j + 1) % d_cond];
                0.5 * cj + 0.3 * cj * ck + 0.1
            }).collect()
        }).collect();
        let mk = |c: &[f32], j: usize| -> f32 {
            let cj = c[j % d_cond];
            let ck = c[(j + 1) % d_cond];
            0.5 * cj + 0.3 * cj * ck + 0.1
        };
        let cond_t = Tensor::from_vec(conds.iter().flatten().cloned().collect(), (n_dados, d_cond), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let alvo_t = Tensor::from_vec(alvos.iter().flatten().cloned().collect(), (n_dados, d_saida), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        // ---- BACKBONE (cabeça treinada): alvo → classes K ----
        // classe = faixa do 1º componente do alvo (aprendível, mensurável)
        let k_classes = 12usize;
        let classes: Vec<u32> = alvos
            .iter()
            .map(|a| {
                let v = a[0];
                let f = (((v + 2.0) / 4.0) * (k_classes as f32)).floor();
                (f.clamp(0.0, (k_classes - 1) as f32)) as u32
            })
            .collect();
        let cabeca = crate::jepa::MlpJepa::novo(d_saida, 4 * d_saida, k_classes, &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let classe_t = Tensor::from_vec(classes.clone(), (n_dados,), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut otim_cab = treino::AdamWManual::novo(cabeca.vars(), 1e-2, 0.0);
        for _ in 0..200 {
            let logits = cabeca.fwd(&alvo_t)
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
            let ce = candle_nn::loss::cross_entropy(&logits, &classe_t)
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
            otim_cab.passo(treino::grads_achatados(&ce, &cabeca.vars())
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?);
        }
        // ---- VERIFICADOR calibrado em RETIDOS (metade final) ----
        let n_ret = n_dados / 2;
        let (alvo_ret, clas_ret) = (
            alvo_t.narrow(0, n_ret, n_dados - n_ret).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?,
            classe_t.narrow(0, n_ret, n_dados - n_ret).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?,
        );
        let logits_ret = cabeca.fwd(&alvo_ret).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut verif = verificador::VerificadorNcm::novo();
        let (ece_antes, ece_depois) = verif.ajustar(&logits_ret, &clas_ret)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        // logits_fn: resposta (1, d_saida) → logits do backbone (1, K)
        let logits_fn = |resp: &Tensor| -> candle_core::Result<Tensor> {
            cabeca.fwd(resp)
        };
        // ---- treino cooperativo do motor (item a) ----
        let cond_tr = cond_t.narrow(0, 0, n_ret).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let alvo_tr = alvo_t.narrow(0, 0, n_ret).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut hist = Vec::new();
        for _ in 0..passos_treino {
            let perda = motor.perda_treino(&cond_tr, &alvo_tr)
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
            let pv = perda.to_scalar::<f32>()
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))? as f64;
            hist.push(pv);
            otim.passo(treino::grads_achatados(&perda, &motor.vars())
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?);
        }
        // ---- consultas RETIDAS em distribuição (ganho medido AQUI) ----
        let n_consultas = 8usize;
        let mut confs_base = Vec::new();
        let mut confs_pens = Vec::new();
        let mut iteracoes = Vec::new();
        for qi in 0..n_consultas {
            let consulta: Vec<f32> = conds[n_ret + qi].clone();
            let consulta_t = Tensor::from_vec(consulta, (d_cond,), &dev)
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
            let mut rng2 = crate::rng::RngDeterministico::novo(semente + 1 + qi as u64);
            let res = motor.raciocinar_ncm(&consulta_t, &verif, &logits_fn, &mut rng2)
                .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
            // base: 1 tentativa isolada, mesma verificação NCM
            let base = motor.tentativas[0]
                .fwd(&consulta_t.broadcast_as((1, d_cond)).unwrap()).unwrap();
            let conf_base = verif.confianca_1d(&logits_fn(&base).unwrap()).unwrap();
            confs_base.push(conf_base);
            confs_pens.push(res.confianca);
            iteracoes.push(res.iteracoes);
        }
        let media = |v: &[f64]| v.iter().sum::<f64>() / v.len().max(1) as f64;
        let (conf_base_m, conf_pens_m) = (media(&confs_base), media(&confs_pens));
        let iteracoes_media = iteracoes.iter().sum::<usize>() as f64 / iteracoes.len().max(1) as f64;
        // ---- consulta OOD (reporte honesto — sem gate de ganho) ----
        let consulta_ood: Vec<f32> = (0..d_cond).map(|j| 3.0 + (j as f32) * 0.1).collect();
        let consulta_ood_t = Tensor::from_vec(consulta_ood, (d_cond,), &dev)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let mut rng3 = crate::rng::RngDeterministico::novo(semente + 99);
        let res_ood = motor.raciocinar_ncm(&consulta_ood_t, &verif, &logits_fn, &mut rng3)
            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("{e}")))?;
        let base_ood = motor.tentativas[0]
            .fwd(&consulta_ood_t.broadcast_as((1, d_cond)).unwrap()).unwrap();
        let conf_base_ood = verif.confianca_1d(&logits_fn(&base_ood).unwrap()).unwrap();
        let _ = mk;
        let d_ = serde_json::json!({
            "versao": "pensamento-0.6.0-ncm",
            "n_tentativas": n_tentativas,
            "perda_treino_inicial": hist.first(), "perda_treino_final": hist.last(),
            "treino_aprendeu": hist.last().copied() < hist.first().map(|p| p * 0.5),
            "ece_verificador_antes": ece_antes, "ece_verificador_depois": ece_depois,
            "temperatura_calibrada": verif.temperatura,
            "iteracoes_ciclo_d_media": iteracoes_media,
            "confianca_base_1_rede": conf_base_m,
            "confianca_pensamento": conf_pens_m,
            "ganho_confianca": conf_pens_m - conf_base_m,
            "confianca_ood_base": conf_base_ood,
            "confianca_ood_pensamento": res_ood.confianca,
            "iteracoes_ood": res_ood.iteracoes,
            "orcamento_esgotado_ood": res_ood.telemetria.get("orcamento_esgotado").copied().unwrap_or(0.0),
            "verificador": "ncm-calibrado (sem oraculo)",
            "gates": {
                "gate_verificador_sem_oraculo": true,
                "gate_calibrado": ece_depois <= ece_antes,
                "gate_ganho_nao_negativo": conf_pens_m >= conf_base_m - 1e-3,
                "gate_terminou": iteracoes_media <= orcamento as f64,
            },
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}


/// **Autoajuste do Pensamento por telemetria (A-1..A-5) — demo:** alimenta
/// o controlador com sinais sintéticos medidos e verifica cada regra.
#[pyfunction]
fn demo_autoajuste(semente: u64) -> PyResult<PyObject> {
    Python::with_gil(|py| {
        let mut hub = telemetria::TelemetryHub::novo(None);
        let mut rng = crate::rng::RngDeterministico::novo(semente);
        let mut auto = controlador::AutoAjustePensamento::novo(0.5, 8, 2);
        // A-1: confianças com mistura (60% altas ~0.85–0.95, 40% baixas ~0.25–0.40)
        // ⇒ quantil 0.75 cai no grupo alto (posições 0.6–1.0) ≈ 0.85+0.375·0.1 ≈ 0.89
        for _ in 0..200 {
            let c = if rng.uniforme() < 0.6 { 0.85 + rng.uniforme() * 0.1 } else { 0.25 + rng.uniforme() * 0.15 };
            auto.confs_retidas.push(c);
        }
        let tau = auto.ajustar_tau_c(0.25, &mut hub);
        // A-3: acordos altos (0.88 ⇒ p̂ = (1+√0.76)/2 ≈ 0.936 ⇒ N = ln(20)/(2·0.436²) ≈ 7.9 ⇒ 8)
        for _ in 0..32 {
            auto.acordos.push(0.88);
        }
        let n_ajustado = auto.ajustar_n(0.05, 16, &mut hub);
        // A-5: ganhos estáveis positivos (n=64 ⇒ 3º termo Bernstein ≈ 0.14 < 0.30) ⇒ gate sobe
        for _ in 0..64 {
            auto.ajustar_gate(0.30, 0.05, &mut hub);
        }
        let g_final = auto.gate_g;
        // A-2: ganhos decaindo geometricamente
        for k in 0..12 {
            auto.ganhos_por_iteracao.push(0.25 * 0.6f64.powi(k));
        }
        let b_ajustado = auto.ajustar_orcamento(0.02, &mut hub);
        // A-4: UCB1 com 3 braços
        let mut u = controlador::Ucb1::novo(vec!["n1".into(), "n4_moe".into(), "n4".into()]);
        for _ in 0..200 {
            let braco = u.escolher();
            let media = match braco { 0 => 0.35, 1 => 0.80, _ => 0.55 };
            let r = (media + rng.normal(0.0, 0.06) as f64).clamp(0.0, 1.0);
            u.registrar(braco, r);
        }
        let medias = u.medias();
        let melhor = medias.iter().enumerate().max_by(|a, b| a.1.partial_cmp(b.1).unwrap()).map(|(i, _)| i).unwrap_or(0);
        let d_ = serde_json::json!({
            "versao": "autoajuste-0.6.0",
            "tau_c_final": tau, "n_tentativas_final": n_ajustado,
            "gate_final": g_final, "orcamento_b_final": b_ajustado,
            "ajustes_auditados": hub.eventos_ajuste.len(),
            "ucb1_medias": medias, "ucb1_melhor_braco": u.nomes[melhor],
            "ucb1_arrependimento_medio": u.arrependimento_empirico() / 200.0,
            "gates": {
                "gate_a1_tau_quantil": (tau - 0.89).abs() < 0.08,
                "gate_a3_n_reduzido": n_ajustado <= 8,
                "gate_a5_gate_subiu": g_final > 0.0,
                "gate_a2_b_finito": b_ajustado >= 1 && b_ajustado <= 16,
                "gate_a4_ucb1_achou": u.nomes[melhor] == "n4_moe",
            },
        });
        Ok(pythonize::pythonize(py, &d_)?.into())
    })
}


// ---------------------------------------------------------------------------
// 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")]
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))]
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())
}

fn py_err(e: candle_core::Error) -> pyo3::PyErr {
    pyo3::exceptions::PyRuntimeError::new_err(format!("{e}"))
}

/// **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))]
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))]
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())
}

/// Demo/gate das unidades NLP/NLG sobre estados ocultos sintéticos
/// (nascimento neutro + gates + coerência de bigramas + especulativo).
#[pyfunction]
#[pyo3(signature = (d=32, t=6, seed=11))]
fn demo_nlp_nlg(py: Python, d: usize, t: usize, seed: u64) -> PyResult<PyObject> {
    let dev = candle_core::Device::Cpu;
    let mut rng = rng::RngDeterministico::novo(seed);
    let v: Vec<f32> = (0..2 * t * d).map(|_| rng.normal_padrao() as f32).collect();
    let h = candle_core::Tensor::from_vec(v, (2, t, d), &dev).map_err(py_err)?;
    // NLP: enriquecimento aditivo — neutro no nascimento
    let mut nlp = nlp::UnidadeNlp::nova(d, &dev).map_err(|e| {
        pyo3::exceptions::PyRuntimeError::new_err(format!("nlp (d%4): {e}"))
    })?;
    let h_enr = nlp.enriquecer(&h).map_err(py_err)?;
    let dif_nlp = (&h).sub(&h_enr).map_err(py_err)?.abs().map_err(py_err)?.max_all().map_err(py_err)?.to_scalar::<f32>().map_err(py_err)? as f64;
    // intenção
    let (intencoes, confs) = nlp.intencao_argmax(&h).map_err(py_err)?;
    // NLG: plano + condicionamento neutro + coerência
    let mut nlg = nlg::UnidadeNlg::nova(d, 4096, &dev).map_err(py_err)?;
    let plano = nlg.plano(&h).map_err(py_err)?;
    let p8 = candle_core::Tensor::randn(0f32, 1f32, (nlg::N_EIXOS_ESTILO,), &dev).map_err(py_err)?;
    let p01 = candle_nn::ops::sigmoid(&p8).map_err(py_err)?;
    let h_cond = nlg.condicionar(&h, Some(&p01)).map_err(py_err)?;
    let dif_nlg = (&h).sub(&h_cond).map_err(py_err)?.abs().map_err(py_err)?.max_all().map_err(py_err)?.to_scalar::<f32>().map_err(py_err)? as f64;
    // coerência: contexto observado vs candidato novo
    let ctx = vec![5u32, 6, 7, 8, 9];
    nlg.observar_contexto(&ctx);
    let s_mesmo = nlg.score_coerencia(&ctx);
    let s_novo = nlg.score_coerencia(&[50u32, 51, 52, 53]);
    // especulativo: aceitação típica + stats
    let mut espec = nlg::DecodificadorEspeculativo::novo(4, 0.95, 10);
    let logits = candle_core::Tensor::from_vec(vec![3.0f32, 0.2, 0.1, 0.05], (4,), &dev)
        .map_err(py_err)?;
    let mut aceites = 0usize;
    for _ in 0..30 {
        let (_, ok) = espec.verificar_e_corrigir(&logits, &[0, 1]).map_err(py_err)?;
        if ok {
            aceites += 1;
        }
    }
    let d_ = serde_json::json!({
        "versao": "nlp-nlg-0.7.0",
        "nlp_gate_medio": nlp.gate_medio,
        "nlp_dif_nascimento": dif_nlp,
        "nlp_intencoes": intencoes,
        "nlp_confiancas": confs,
        "nlg_plano_forma": plano.dims(),
        "nlg_dif_nascimento": dif_nlg,
        "nlg_coerencia_mesmo": s_mesmo,
        "nlg_coerencia_novo": s_novo,
        "espec_aceitacao_media": espec.stats.aceitacao_media,
        "espec_speedup": espec.stats.speedup_estimado,
        "espec_aceites_30": aceites,
        "gates": {
            "nlp_neutro": dif_nlp == 0.0,
            "nlg_neutro": dif_nlg == 0.0,
            "gate_medio_ok": (nlp.gate_medio - 0.119).abs() < 0.05,
            "coerencia_ok": s_mesmo > 0.9 && s_novo < 0.1,
            "espec_ok": espec.stats.aceitacao_media > 0.8,
        },
    });
    Ok(pythonize::pythonize(py, &d_)?.into())
}

/// Demo/gate do auto-ajuste SOM com tetos (S-AB/S-SL).
#[pyfunction]
fn demo_autoajuste_som(py: Python) -> PyResult<PyObject> {
    let mut a = controlador::AutoAjusteSom::novo(0.5, 1.0);
    let (al_qe, _) = a.passo(1.0, 0.0, 0.0);
    let (al_cross, sg_cross) = a.passo(0.0, 0.0, 1.0);
    let (al_max, sg_max) = a.passo(1.0, 1.0, 1.0);
    let mut parou_5chunks = false;
    for _ in 0..5 {
        parou_5chunks = a.avaliar_parada(0.95, 0.1, 0.1, false, true, false, false);
    }
    let d_ = serde_json::json!({
        "versao": "som-autoadjust-0.7.0",
        "alfa_qe_alta": al_qe, "alfa_crossing": al_cross, "sigma_crossing": sg_cross,
        "alfa_max": al_max, "sigma_max": sg_max,
        "parou_5chunks": parou_5chunks,
        "motivo": a.motivo_parada,
        "gates": {
            "tetos_ok": al_max <= 1.0 + 1e-9 && sg_max <= 2.0 + 1e-9,
            "piso_ok": al_cross >= 0.5 - 1e-9,
            "qe_sobe": al_qe > 0.5,
            "stop_dispara": parou_5chunks,
        },
    });
    Ok(pythonize::pythonize(py, &d_)?.into())
}

#[pymodule]
fn nucleo_rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
    pyo3_log::init();
    m.add_function(wrap_pyfunction!(versao, m)?)?;
    m.add_function(wrap_pyfunction!(info_build, m)?)?;
    m.add_function(wrap_pyfunction!(lint_config, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_fase_a, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_fase_b, m)?)?;
    m.add_function(wrap_pyfunction!(gerar_amostras, m)?)?;
    m.add_function(wrap_pyfunction!(verificar_mascara, m)?)?;
    m.add_function(wrap_pyfunction!(verificar_stiefel, m)?)?;
    // v0.5.0 — SOM 4D / JEPA / Pensamento
    m.add_function(wrap_pyfunction!(treinar_som4d, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_gng, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_hsom, m)?)?;
    m.add_function(wrap_pyfunction!(demo_tsom, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_jepa, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_lejepa, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_mcjepa, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_bjepa, m)?)?;
    m.add_function(wrap_pyfunction!(demo_sonic_moe, m)?)?;
    m.add_function(wrap_pyfunction!(demo_replay_ssm, m)?)?;
    m.add_function(wrap_pyfunction!(classificar_erro, m)?)?;
    m.add_function(wrap_pyfunction!(confianca_de_logits, m)?)?;
    m.add_function(wrap_pyfunction!(projetar_ridge, m)?)?;
    m.add_function(wrap_pyfunction!(demo_pensamento, m)?)?;
    m.add_function(wrap_pyfunction!(patches_2d, m)?)?;
    m.add_function(wrap_pyfunction!(patches_serie, m)?)?;
    m.add_function(wrap_pyfunction!(grupos_pontos, m)?)?;
    // v0.6.0
    m.add_function(wrap_pyfunction!(calibrar_verificador, m)?)?;
    m.add_function(wrap_pyfunction!(demo_verificador_ncm, m)?)?;
    m.add_function(wrap_pyfunction!(treinar_fase_c, m)?)?;
    m.add_function(wrap_pyfunction!(demo_moe_redirecionavel, m)?)?;
    m.add_function(wrap_pyfunction!(jepa_com_moe, m)?)?;
    m.add_function(wrap_pyfunction!(som_com_moe, m)?)?;
    m.add_function(wrap_pyfunction!(demo_acoplador_duplo, m)?)?;
    m.add_function(wrap_pyfunction!(demo_pensamento_ncm, m)?)?;
    m.add_function(wrap_pyfunction!(demo_autoajuste, m)?)?;
    m.add_function(wrap_pyfunction!(demo_fnh, m)?)?;
    m.add_function(wrap_pyfunction!(demo_quebra_cabeca, m)?)?;
    m.add_function(wrap_pyfunction!(quebra_cabeca_real, m)?)?;
    m.add_function(wrap_pyfunction!(demo_nlp_nlg, m)?)?;
    m.add_function(wrap_pyfunction!(demo_autoajuste_som, m)?)?;
    m.add_class::<BpePy>()?;
    m.add_class::<ConfigPy>()?;
    m.add_class::<PonteDuplaPy>()?;
    #[cfg(feature = "ponte_async")]
    m.add_class::<PonteAsyncPy>()?;
    m.add_class::<ArvorePesosPy>()?;
    Ok(())
}