//! `jepa::variantes` — 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::*;


// ---------------------------------------------------------------------------
// Parte 4 — variantes e treino unificado
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum VarianteJepa {
    ImageJepa,
    VJepa2,
    McJepa,
    AudioJepa,
    PointJepa,
    Jepa3D,
    BrainJepa,
    TJepa,
    LeJepa,
    BJepa,
}

impl VarianteJepa {
    pub fn nome(&self) -> &'static str {
        match self {
            Self::ImageJepa => "Image-JEPA",
            Self::VJepa2 => "V-JEPA 2",
            Self::McJepa => "MC-JEPA",
            Self::AudioJepa => "Audio-JEPA",
            Self::PointJepa => "Point-JEPA",
            Self::Jepa3D => "3D-JEPA",
            Self::BrainJepa => "Brain-JEPA",
            Self::TJepa => "T-JEPA",
            Self::LeJepa => "LeJEPA",
            Self::BJepa => "BJEPA",
        }
    }
}

#[derive(Debug, Clone)]
pub struct CfgJepa {
    pub d_emb: usize,
    pub epocas: usize,
    pub lr: f64,
    pub m: f64,          // EMA
    pub lambda_vc: f64,
    pub gamma_vc: f64,
    pub frac_mascara: f64,
    pub semente: u64,
    /// tamanho do bloco-alvo (nº de patches por bloco).
    pub tam_bloco: usize,
}

impl Default for CfgJepa {
    fn default() -> Self {
        Self {
            d_emb: 32,
            epocas: 60,
            lr: 1.5e-3,
            m: 0.99,
            lambda_vc: 0.5,
            gamma_vc: 1.0,
            frac_mascara: 0.35,
            semente: 42,
            tam_bloco: 2,
        }
    }
}

/// Relatório de treino de uma variante.
#[derive(Debug, Clone)]
pub struct RelatorioJepa {
    pub variante: String,
    pub perda_inicial: f64,
    pub perda_final: f64,
    /// perda PREDITIVA isolada (sem o regularizador VC) — critério de
    /// aprendizado honesto (a VC pode crescer legitimamente ao empurrar o
    /// std das incorporações para cima).
    pub perda_pred_inicial: f64,
    pub perda_pred_final: f64,
    /// desvio-padrão médio por dimensão das incorporações (monitor de
    /// colapso: → 0 = colapso).
    pub std_emb_final: f64,
    pub std_emb_inicial: f64,
    pub n_patches: usize,
    pub d_patch: usize,
    pub epocas: usize,
}

/// Treina uma variante JEPA sobre patches já extraídos `(n, d_patch)` com
/// um esquema de máscara fornecido (padrão: intervalos 1D).
///
/// Máscara=true ⇒ patch é ALVO (bloco); false ⇒ CONTEXTO.
pub fn treinar_variante(
    patches: &Tensor,
    mascara: &dyn Fn(usize, &mut RngDeterministico) -> Vec<bool>,
    cfg: &CfgJepa,
    variante: VarianteJepa,
    cond: Option<&Tensor>,
    dev: &Device,
) -> Result<RelatorioJepa> {
    let n = patches.dim(0)?;
    let d_patch = patches.dim(1)?;
    let mut rng = RngDeterministico::novo(cfg.semente);
    let mut nucleo = NucleoJepa::novo(
        d_patch,
        cfg.d_emb,
        if cond.is_some() { cond.unwrap().dim(1)? } else { 0 },
        n,
        cfg.m,
        cfg.lambda_vc,
        cfg.gamma_vc,
        dev,
    )?;
    let otimo_vars = nucleo.vars();
    let mut otimizador = crate::treino::AdamWManual::novo(otimo_vars, cfg.lr, 0.01);

    // métricas iniciais
    let z0 = nucleo.codificar(patches)?;
    let std0 = std_por_dim(&z0)?;

    let mut perda_inicial = f64::NAN;
    let mut perda_final = f64::NAN;
    let mut pred_inicial = f64::NAN;
    let mut pred_final = f64::NAN;
    let mut vc_final = f64::NAN;
    for epoca in 0..cfg.epocas {
        // máscara fresca a cada época (com proteção contra máscara degenerada:
        // se contexto ou alvo ficar vazio, repete com fração menor — a época
        // NUNCA é pulada, a perda sempre é computada)
        let (ctx_idx, blocos) = loop {
            let mut masc = mascara(n, &mut rng);
            // defensivo: máscara menor que n => o resto é contexto
            if masc.len() < n {
                masc.resize(n, false);
            }
            masc.truncate(n);
            let ctx: Vec<u32> = (0..n).filter(|&i| !masc[i]).map(|i| i as u32).collect();
            let alvos: Vec<usize> = (0..n).filter(|&i| masc[i]).collect();
            if !ctx.is_empty() && !alvos.is_empty() {
                break (ctx, alvos);
            }
        };
        // agrupa alvos em blocos de tam_bloco
        let grupos: Vec<Vec<usize>> = blocos
            .chunks(cfg.tam_bloco.max(1))
            .map(|c| c.to_vec())
            .collect();
        if ctx_idx.is_empty() || grupos.is_empty() {
            continue;
        }
        let dev2 = patches.device();
        let n_ctx = ctx_idx.len();
        let idx_t = Tensor::from_vec(ctx_idx, (n_ctx,), &dev2)?;
        let ctx_patches = patches.index_select(&idx_t, 0)?;
        let z_ctx_todos = nucleo.codificar(&ctx_patches)?; // (n_ctx, d_emb)
        let z_ctx = z_ctx_todos.mean_keepdim(0)?; // (1, d_emb)
        // alvos por grupo
        let mut alvos: Vec<Vec<(usize, Tensor)>> = Vec::new();
        for g in &grupos {
            let idx: Vec<u32> = g.iter().map(|&i| i as u32).collect();
            let n_g = idx.len();
            let it = Tensor::from_vec(idx, (n_g,), &dev2)?;
            alvos.push(vec![(g[0], patches.index_select(&it, 0)?)]);
        }
        // condição por amostra — broadcast (o lote é 1 aqui: o conjunto todo)
        let (perda, lp, lv) = nucleo.perda_lote(&z_ctx, &[alvos_concat(&alvos)], cond, &z_ctx_todos)?;
        let p = perda.to_scalar::<f32>()? as f64;
        if p.is_finite() {
            if epoca == 0 {
                perda_inicial = p;
                pred_inicial = lp;
            }
            perda_final = p;
            pred_final = lp;
            vc_final = lv;
        }
        // decaimento cosseno do lr + clipping de gradiente (estabilidade)
        let frac = epoca as f64 / cfg.epocas.max(1) as f64;
        let lr_t = cfg.lr * 0.5 * (1.0 + (std::f64::consts::PI * frac).cos());
        otimizador.set_lr(lr_t);
        let mut grads = crate::treino::grads_achatados(&perda, &nucleo.vars())?;
        crate::treino::AdamWManual::clipar(&mut grads, 1.0);
        otimizador.passo(grads);
        nucleo.ema_passo()?;
    }
    let z1 = nucleo.codificar(patches)?;
    let std1 = std_por_dim(&z1)?;
    let _ = (pred_inicial, vc_final);
    Ok(RelatorioJepa {
        variante: variante.nome().to_string(),
        perda_inicial,
        perda_final,
        perda_pred_inicial: pred_inicial,
        perda_pred_final: pred_final,
        std_emb_final: std1,
        std_emb_inicial: std0,
        n_patches: n,
        d_patch,
        epocas: cfg.epocas,
    })
}

fn alvos_concat(alvos: &[Vec<(usize, Tensor)>]) -> Vec<(usize, Tensor)> {
    alvos.iter().flatten().cloned().collect()
}

fn std_por_dim(z: &Tensor) -> Result<f64> {
    let n = z.dim(0)? as f64;
    let media = z.mean_keepdim(0)?;
    let centr = z.broadcast_sub(&media)?;
    let var = (centr.sqr()?.sum(0)? / n)?;
    Ok(var.sqrt()?.mean_all()?.to_scalar::<f32>()? as f64)
}


// ---------------------------------------------------------------------------
// Variantes especiais: LeJEPA (edição-condicionado), MC-JEPA (cruz-modal),
// BJEPA (bidirecional com consistência de ciclo)
// ---------------------------------------------------------------------------

/// **LeJEPA** — predição latente condicionada por EDIÇÃO: contexto = E(x),
/// condição = instrução de edição `e` (vetor de baixa dimensão), alvo =
/// `E_ξ(edit_e(x))`. A condição é a INSTRUÇÃO (não o resultado — sem
/// vazamento: `e` não contém `edit(x)`).
pub fn treinar_lejepa(
    dados: &Tensor,      // (n, d_patch) — originais
    edits: &Tensor,      // (n, d_cond) — instruções de edição
    editados: &Tensor,   // (n, d_patch) — resultados das edições
    cfg: &CfgJepa,
    dev: &Device,
) -> Result<RelatorioJepa> {
    let n = dados.dim(0)?;
    let d_patch = dados.dim(1)?;
    let d_cond = edits.dim(1)?;
    let mut rng = RngDeterministico::novo(cfg.semente);
    let mut nucleo = NucleoJepa::novo(
        d_patch, cfg.d_emb, d_cond, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev,
    )?;
    let mut otim = crate::treino::AdamWManual::novo(nucleo.vars(), cfg.lr, 0.01);
    let mut perda_inicial = f64::NAN;
    let mut perda_final = f64::NAN;
    let z0 = nucleo.codificar(dados)?;
    let std0 = std_por_dim(&z0)?;
    for epoca in 0..cfg.epocas {
        // subamostra do lote
        let idx = rng.amostra_k_de_n(n.min(64), n);
        let sel: Vec<u32> = idx.iter().map(|&i| i as u32).collect();
        let it = Tensor::from_vec(sel.clone(), (sel.len(),), dev)?;
        let xb = dados.index_select(&it, 0)?;
        let eb = edits.index_select(&it, 0)?;
        let yb = editados.index_select(&it, 0)?;
        let z_ctx = nucleo.codificar(&xb)?; // (m, d_emb)
        let mut perdas: Vec<Tensor> = Vec::new();
        for i in 0..xb.dim(0)? {
            let zc = z_ctx.get(i)?.broadcast_as((1, cfg.d_emb))?;
            let e_i = eb.get(i)?.broadcast_as((1, d_cond))?;
            let yt = yb.get(i)?.broadcast_as((1, d_patch))?;
            let zt = nucleo
                .alvo_codificar(&yt)?
                .mean(0)?
                .broadcast_as((1, cfg.d_emb))?;
            let pred = nucleo.predizer(&zc, 0, Some(&e_i))?;
            let dif = (&pred - &zt)?;
            perdas.push((dif.sqr()?.sum_all()? / cfg.d_emb as f64)?);
        }
        let l_pred = (Tensor::stack(&perdas, 0)?.sum_all()? / perdas.len() as f64)?;
        let l_vc = nucleo.perda_vc(&z_ctx)?;
        let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(cfg.lambda_vc, 0.0)?)?;
        let p = total.to_scalar::<f32>()? as f64;
        if p.is_finite() {
            if epoca == 0 {
                perda_inicial = p;
            }
            perda_final = p;
        }
        let mut grads = crate::treino::grads_achatados(&total, &nucleo.vars())?;
        crate::treino::AdamWManual::clipar(&mut grads, 1.0);
        otim.passo(grads);
        nucleo.ema_passo()?;
    }
    let z1 = nucleo.codificar(dados)?;
    Ok(RelatorioJepa {
        variante: "LeJEPA".into(),
        perda_inicial,
        perda_final,
        perda_pred_inicial: perda_inicial,
        perda_pred_final: perda_final,
        std_emb_final: std_por_dim(&z1)?,
        std_emb_inicial: std0,
        n_patches: n,
        d_patch,
        epocas: cfg.epocas,
    })
}

/// **MC-JEPA** — cruz-modal: dois codificadores, predição nos DOIS sentidos
/// (A→B e B→A) com preditores separados e alvos EMA dos dois lados.
pub fn treinar_mcjepa(
    mod_a: &Tensor, // (n, d_a)
    mod_b: &Tensor, // (n, d_b)
    cfg: &CfgJepa,
    dev: &Device,
) -> Result<RelatorioJepa> {
    let n = mod_a.dim(0)?;
    let d_a = mod_a.dim(1)?;
    let d_b = mod_b.dim(1)?;
    // núcleo A→B e núcleo B→A (compartilham apenas a estrutura)
    let mut nb = NucleoJepa::novo(d_a, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev)?;
    let mut na = NucleoJepa::novo(d_b, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev)?;
    let mut otim = crate::treino::AdamWManual::novo(
        [nb.vars(), na.vars()].concat(),
        cfg.lr,
        0.01,
    );
    let mut rng = RngDeterministico::novo(cfg.semente);
    let mut perda_inicial = f64::NAN;
    let mut perda_final = f64::NAN;
    for epoca in 0..cfg.epocas {
        let idx = rng.amostra_k_de_n(n.min(64), n);
        let sel: Vec<u32> = idx.iter().map(|&i| i as u32).collect();
        let it = Tensor::from_vec(sel, (idx.len(),), dev)?;
        let xa = mod_a.index_select(&it, 0)?;
        let xb = mod_b.index_select(&it, 0)?;
        // z_a = E_a(x_a) prevê ẑ_b; z_b = E_b(x_b) prevê ẑ_a
        let z_a = nb.codificar(&xa)?; // (m, d_emb)
        let z_b = na.codificar(&xb)?;
        let t_a = na.alvo_codificar(&xa)?; // alvo EMA do lado A
        let t_b = nb.alvo_codificar(&xb)?; // alvo EMA do lado B
        let m_ = xa.dim(0)?;
        let mut perdas: Vec<Tensor> = Vec::new();
        for i in 0..m_ {
            // A → B
            let zc = z_a.get(i)?.broadcast_as((1, cfg.d_emb))?;
            let zt = t_b.get(i)?.broadcast_as((1, cfg.d_emb))?;
            let pred = nb.predizer(&zc, 0, None)?;
            let dif = (&pred - &zt)?;
            perdas.push((dif.sqr()?.sum_all()? / cfg.d_emb as f64)?);
            // B → A
            let zc2 = z_b.get(i)?.broadcast_as((1, cfg.d_emb))?;
            let zt2 = t_a.get(i)?.broadcast_as((1, cfg.d_emb))?;
            let pred2 = na.predizer(&zc2, 0, None)?;
            let dif2 = (&pred2 - &zt2)?;
            perdas.push((dif2.sqr()?.sum_all()? / cfg.d_emb as f64)?);
        }
        let l_pred = (Tensor::stack(&perdas, 0)?.sum_all()? / perdas.len() as f64)?;
        let l_vc = (nb.perda_vc(&z_a)? + na.perda_vc(&z_b)?)?;
        let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(cfg.lambda_vc, 0.0)?)?;
        let p = total.to_scalar::<f32>()? as f64;
        if p.is_finite() {
            if epoca == 0 {
                perda_inicial = p;
            }
            perda_final = p;
        }
        let mut grads = crate::treino::grads_achatados(&total, &[nb.vars(), na.vars()].concat())?;
        crate::treino::AdamWManual::clipar(&mut grads, 1.0);
        otim.passo(grads);
        nb.ema_passo()?;
        na.ema_passo()?;
    }
    let z1 = nb.codificar(mod_a)?;
    Ok(RelatorioJepa {
        variante: "MC-JEPA".into(),
        perda_inicial,
        perda_final,
        perda_pred_inicial: perda_inicial,
        perda_pred_final: perda_final,
        std_emb_final: std_por_dim(&z1)?,
        std_emb_inicial: std_por_dim(&nb.codificar(mod_a)?)?,
        n_patches: n,
        d_patch: d_a.max(d_b),
        epocas: cfg.epocas,
    })
}

/// **BJEPA** — bidirecional: preditor causal (z_t → z̃_{t+1}), anti-causal
/// (z_{t+1} → z̃_t) e consistência de ciclo
/// `‖P_b(P_f(z_t)) − sg(z_t)‖²` — a volta deve reconstruir o presente.
pub fn treinar_bjepa(patches: &Tensor, cfg: &CfgJepa, dev: &Device) -> Result<RelatorioJepa> {
    let n = patches.dim(0)?;
    let d_patch = patches.dim(1)?;
    let mut nucleo = NucleoJepa::novo(
        d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev,
    )?;
    // preditor anti-causal: núcleo separado
    let mut ninv = NucleoJepa::novo(
        d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, dev,
    )?;
    let mut otim =
        crate::treino::AdamWManual::novo([nucleo.vars(), ninv.vars()].concat(), cfg.lr, 0.01);
    let mut perda_inicial = f64::NAN;
    let mut perda_final = f64::NAN;
    for epoca in 0..cfg.epocas {
        let z = nucleo.codificar(patches)?; // (n, d_emb)
        let zt = nucleo.alvo_codificar(patches)?; // alvos EMA (stop-grad)
        let mut perdas: Vec<Tensor> = Vec::new();
        for t in 0..n.saturating_sub(1) {
            let zc = z.get(t)?.broadcast_as((1, cfg.d_emb))?;
            let z_prox = zt.get(t + 1)?.broadcast_as((1, cfg.d_emb))?;
            // causal: z_t → z̃_{t+1} (pos 0 = "próximo")
            let pf = nucleo.predizer(&zc, 0, None)?;
            let dif = (&pf - &z_prox)?;
            perdas.push((dif.sqr()?.sum_all()? / cfg.d_emb as f64)?);
            // anti-causal: z_{t+1} → z̃_t (pos 1 = "anterior")
            let zb = z.get(t + 1)?.broadcast_as((1, cfg.d_emb))?;
            let z_ant = zt.get(t)?.broadcast_as((1, cfg.d_emb))?;
            let pb = ninv.predizer(&zb, 1, None)?;
            let dif2 = (&pb - &z_ant)?;
            perdas.push((dif2.sqr()?.sum_all()? / cfg.d_emb as f64)?);
            // ciclo: P_b(P_f(z_t)) ≈ sg(z_t)
            let pf_sg = nucleo.predizer(&zc, 0, None)?.detach();
            let pciclo = ninv.predizer(&pf_sg.broadcast_as((1, cfg.d_emb))?, 1, None)?;
            let zt_sg = zt.get(t)?.broadcast_as((1, cfg.d_emb))?;
            let dif3 = (&pciclo - &zt_sg)?;
            perdas.push((dif3.sqr()?.sum_all()? / cfg.d_emb as f64)?);
        }
        let l_pred = (Tensor::stack(&perdas, 0)?.sum_all()? / perdas.len().max(1) as f64)?;
        let l_vc = nucleo.perda_vc(&z)?;
        let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(cfg.lambda_vc, 0.0)?)?;
        let p = total.to_scalar::<f32>()? as f64;
        if p.is_finite() {
            if epoca == 0 {
                perda_inicial = p;
            }
            perda_final = p;
        }
        let mut grads = crate::treino::grads_achatados(&total, &[nucleo.vars(), ninv.vars()].concat())?;
        crate::treino::AdamWManual::clipar(&mut grads, 1.0);
        otim.passo(grads);
        nucleo.ema_passo()?;
        ninv.ema_passo()?;
    }
    let z1 = nucleo.codificar(patches)?;
    Ok(RelatorioJepa {
        variante: "BJEPA".into(),
        perda_inicial,
        perda_final,
        perda_pred_inicial: perda_inicial,
        perda_pred_final: perda_final,
        std_emb_final: std_por_dim(&z1)?,
        std_emb_inicial: std_por_dim(&nucleo.codificar(patches)?)?,
        n_patches: n,
        d_patch,
        epocas: cfg.epocas,
    })
}

/// Esquema de máscara por variante a partir da grade de patches.
pub fn mascara_por_variante(
    variante: VarianteJepa,
    grid: &[usize],
    frac: f64,
    rng: &mut RngDeterministico,
) -> Vec<bool> {
    match variante {
        VarianteJepa::ImageJepa => mascara_blocos_2d(grid[0], grid.get(1).copied().unwrap_or(1), 4, frac, rng),
        VarianteJepa::VJepa2 | VarianteJepa::Jepa3D => mascara_tubos_3d(
            grid[0],
            grid.get(1).copied().unwrap_or(1),
            grid.get(2).copied().unwrap_or(1),
            3,
            frac,
            rng,
        ),
        VarianteJepa::AudioJepa => mascara_bandas(grid[0], grid.get(1).copied().unwrap_or(1), 3, frac, rng),
        VarianteJepa::BrainJepa => mascara_segmentos(grid[0], grid.get(1).copied().unwrap_or(1), frac, rng),
        _ => mascara_intervalos_1d(grid.iter().product::<usize>().max(1), frac, rng),
    }
}

pub fn versao() -> &'static str {
    "jepa-0.5.0"
}