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

use super::*;


/// **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]
pub(crate) 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.forward(&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.forward(&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]
pub(crate) 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]
pub(crate) 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]
pub(crate) 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())
    })
}