//! `python::jepa_treinos` — 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::*;


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


/// Extrai patches 2D de uma imagem achatada (C, H, W) — adaptador Image-JEPA.
#[pyfunction]
pub(crate) 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]
pub(crate) 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]
pub(crate) 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())
}