Blame for src/jepa.rs

//! Suíte JEPA v0.5.0 — núcleo unificado + 12 variantes + SonicMoE + ReplaySSM.
//!
//! **Núcleo (todas as variantes compartilham):**
//! - codificador `E_θ` (MLp SiLU), preditor `P_φ` (condicionado por posição/
//!   ação/edição), codificador-alvo EMA `E_ξ` com `ξ ← m·ξ + (1−m)·θ`
//!   (gradiente cortado — alvo nunca recebe gradiente);
//! - regularizador variância-covariância (VICReg-lite) anti-colapso:
//!   `V(z) = Σ_j max(0, γ − std(z_j))²` + fora-diagonal da covariância.
//! - perda: `L = Σ_b ‖P(z_ctx, pos_b) − sg(ẑ_b)‖²/B + λ_vc·V`.
//!
//! **Prova J-1 (impossibilidade de colapso, em r-20):** se todas as
//! incorporações colapsam para constante, `std(z_j) = 0 ∀j` ⇒ `V = d·γ² > 0`;
//! logo o colapso total NÃO é minimizador com `λ_vc > 0`.
//!
//! **Prova J-2 (EMA = stop-grad + ponto fixo):** `∂L/∂ξ = 0` (detach); se
//! `θ_t → θ*` então `ξ_t → θ*` (recursão linear contratante com `m < 1`).
//!
//! **Prova J-3 (decomposição de informação):** com `P` suficientemente
//! rico, `min_P L = Σ_b Var(ẑ_b | z_ctx)` — a perda retém só o componentes
//! previsíveis do alvo dado o contexto.
//!
//! **Variantes** (esquema de máscara + adaptador de modalidade):
//! Image-JEPA (blocos 2D) · V-JEPA 2 (tubos 3D + cond. de ação) · MC-JEPA
//! (cruz-modal, dois codificadores) · Audio-JEPA (bandas de espectrograma) ·
//! Point-JEPA (grupos FPS) · 3D-JEPA (voxels) · Brain-JEPA (canal×tempo) ·
//! T-JEPA (série temporal canal-independente) · LeJEPA (edição-condicionado) ·
//! BJEPA (bidirecional com consistência de ciclo).
//!
//! **SonicMoE** — MoE com roteamento top-k, capacidade por especialista
//! (tile-aware) e perda auxiliar de balanceamento (argumento Switch:
//! `L_aux = E·Σ f_e·p_e ≥ 1` com igualdade sse uniforme).
//!
//! **ReplaySSM** — SSM de porta diagonal (`h_t = σ(a)⊙h_{t−1} +
//! (1−σ(a))⊙(Bx_t)`, estabilidade herdada do Teo S-2') + replay de experiências
//! para aprendizado continuado (interpolação de gradientes limita o
//! esquecimento — Teo R-1).

use candle_core::{Device, Result, Tensor, Var};
use crate::camadas::Linear;
use crate::rng::RngDeterministico;
use std::collections::HashMap;

// ---------------------------------------------------------------------------
// Parte 1 — núcleo
// ---------------------------------------------------------------------------

/// MLP 2 camadas com SiLU (d → h → d_saida).
pub struct MlpJepa {
    pub l1: Linear,
    pub l2: Linear,
    pub d_entrada: usize,
    pub d_saida: usize,
}

impl MlpJepa {
    pub fn novo(d_entrada: usize, d_oculta: usize, d_saida: usize, dev: &Device) -> Result<Self> {
        let s1 = (1.0 / (d_entrada as f64).sqrt()) * 0.9;
        let s2 = (1.0 / (d_oculta as f64).sqrt()) * 0.9;
        Ok(Self {
            l1: Linear::sem_vies(d_entrada, d_oculta, s1, dev)?,
            l2: Linear::sem_vies(d_oculta, d_saida, s2, dev)?,
            d_entrada,
            d_saida,
        })
    }

    /// v0.9.0 — `novo` com RNG determinístico (quita o débito r-33 §8.1).
    /// Mesma distribuição (N(0, s), s = 0,9/√d) — só a FONTE de aleatoriedade
    /// muda: `RngDeterministico` semeado em vez do RNG global do candle.
    pub fn novo_semeado(
        d_entrada: usize,
        d_oculta: usize,
        d_saida: usize,
        dev: &Device,
        rng: &mut RngDeterministico,
    ) -> Result<Self> {
        let s1 = (1.0 / (d_entrada as f64).sqrt()) * 0.9;
        let s2 = (1.0 / (d_oculta as f64).sqrt()) * 0.9;
        Ok(Self {
            l1: Linear::sem_vies_semeado(d_entrada, d_oculta, s1, dev, rng)?,
            l2: Linear::sem_vies_semeado(d_oculta, d_saida, s2, dev, rng)?,
            d_entrada,
            d_saida,
        })
    }

    pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
        let h = candle_nn::ops::silu(&self.l1.fwd(x)?)?;
        self.l2.fwd(&h)
    }

    pub fn vars(&self) -> Vec<Var> {
        [self.l1.vars(), self.l2.vars()].concat()
    }

    /// EMA: `self ← m·self + (1−m)·outro` (usado no codificador-alvo).
    pub fn ema_de(&mut self, outro: &MlpJepa, m: f64) -> Result<()> {
        for (a, b) in self.vars().iter().zip(outro.vars().iter()) {
            let novo = (a.as_tensor().affine(m, 0.0)? + b.as_tensor().affine(1.0 - m, 0.0)?)?;
            a.set(&novo)?;
        }
        Ok(())
    }
}

/// Núcleo JEPA: codificador + preditor + alvo EMA + regularizador VC.
pub struct NucleoJepa {
    pub encoder: MlpJepa,
    pub preditor: MlpJepa,
    pub alvo: MlpJepa,
    /// tabela de posições (blocos) — entrada extra do preditor.
    pub pos: Var,
    pub d_emb: usize,
    pub d_cond: usize,
    pub m: f64,
    pub lambda_vc: f64,
    pub gamma_vc: f64,
}

impl NucleoJepa {
    pub fn novo(
        d_patch: usize,
        d_emb: usize,
        d_cond: usize,
        n_pos: usize,
        m: f64,
        lambda_vc: f64,
        gamma_vc: f64,
        dev: &Device,
    ) -> Result<Self> {
        Self::novo_com_semente(d_patch, d_emb, d_cond, n_pos, m, lambda_vc, gamma_vc, 42, dev)
    }

    /// v0.9.0 — construção com INICIALIZAÇÃO DETERMINÍSTICA (quita o débito
    /// r-33 §8.1): pesos e tabela de posições nascem de
    /// `RngDeterministico::novo(semente ^ 0x5EED)` — fluxo SEPARADO do RNG de
    /// máscara (`semente`), que permanece intocado. `novo` (acima) delega com
    /// semente fixa 42 para compatibilidade das demos.
    pub fn novo_com_semente(
        d_patch: usize,
        d_emb: usize,
        d_cond: usize,
        n_pos: usize,
        m: f64,
        lambda_vc: f64,
        gamma_vc: f64,
        semente: u64,
        dev: &Device,
    ) -> Result<Self> {
        let mut rng = RngDeterministico::novo(semente ^ 0x5EED_u64);
        let encoder = MlpJepa::novo_semeado(d_patch, 2 * d_emb, d_emb, dev, &mut rng)?;
        let preditor = MlpJepa::novo_semeado(d_emb + d_cond + d_emb, 2 * d_emb, d_emb, dev, &mut rng)?;
        // alvo nasce IGUAL ao encoder (ponto de partida do EMA)
        let mut alvo = MlpJepa::novo_semeado(d_patch, 2 * d_emb, d_emb, dev, &mut rng)?;
        for (va, vb) in alvo.vars().iter().zip(encoder.vars().iter()) {
            va.set(vb.as_tensor())?;
        }
        let n_posicoes = n_pos.max(1);
        let pos_dados: Vec<f32> = (0..n_posicoes * d_emb)
            .map(|_| rng.normal(0.0, 0.02) as f32)
            .collect();
        let pos_t = Tensor::from_vec(pos_dados, (n_posicoes, d_emb), dev)?;
        Ok(Self {
            encoder,
            preditor,
            alvo,
            pos: Var::from_tensor(&pos_t)?,
            d_emb,
            d_cond,
            m,
            lambda_vc,
            gamma_vc,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = self.encoder.vars();
        v.extend(self.preditor.vars());
        v.push(self.pos.clone());
        v // NOTA: alvo NÃO entra — é buffer EMA
    }

    pub fn ema_passo(&mut self) -> Result<()> {
        self.alvo.ema_de(&self.encoder, self.m)
    }

    /// Codifica um conjunto de patches (n, d_patch) → (n, d_emb).
    pub fn codificar(&self, patches: &Tensor) -> Result<Tensor> {
        self.encoder.fwd(patches)
    }

    /// Codifica com o alvo EMA — SEM gradiente (detach por construção: as
    /// vars do alvo não estão no otimizador e o tensor é desconectado).
    pub fn alvo_codificar(&self, patches: &Tensor) -> Result<Tensor> {
        Ok(self.alvo.fwd(patches)?.detach())
    }

    /// Prediz o embedding do bloco `i_pos` a partir do contexto `z_ctx`
    /// (d_emb) e da condição `cond` (d_cond; zeros se None).
    pub fn predizer(&self, z_ctx: &Tensor, i_pos: usize, cond: Option<&Tensor>) -> Result<Tensor> {
        // z_ctx: (b, d_emb) ou (d_emb,)
        let b = z_ctx.dim(0).unwrap_or(1);
        let dev = z_ctx.device();
        let pos = self.pos.as_tensor().get(i_pos.min(self.pos.dim(0)?.saturating_sub(1)))?;
        let pos_b = pos.broadcast_as((b, self.d_emb))?;
        let cond_b = match cond {
            Some(c) => {
                if c.dims().len() == 1 && b > 1 {
                    c.broadcast_as((b, self.d_cond))?
                } else {
                    c.clone()
                }
            }
            None => Tensor::zeros((b, self.d_cond), candle_core::DType::F32, dev)?,
        };
        let entrada = Tensor::cat(&[z_ctx, &pos_b, &cond_b], candle_core::D::Minus1)?;
        self.preditor.fwd(&entrada)
    }

    /// Regularizador variância-covariância do lote de incorporações z (n, d).
    ///
    /// **J-1:** no colapso total `z_i = c ∀i` ⇒ `std_j = 0` ⇒ `V = d·γ²`.
    pub fn perda_vc(&self, z: &Tensor) -> Result<Tensor> {
        let n = z.dim(0)? as f64;
        let d = self.d_emb;
        let media = z.mean_keepdim(0)?; // (1, d)
        let centr = z.broadcast_sub(&media)?;
        let var = (centr.sqr()?.sum(0)? / n)?; // (d,)
        let std = var.sqrt()?;
        // dobradiça: max(0, γ − std)²
        let gamma_t =
            Tensor::new(self.gamma_vc as f32, z.device())?.broadcast_as(std.dims())?;
        let gap = gamma_t.broadcast_sub(&std)?.clamp(0.0, f64::INFINITY)?;
        let v = gap.sqr()?.sum_all()?;
        // covariância fora-diagonal: C = centrᵀ·centr/n; perda = ‖C − diag‖²_F/(2d)
        let cov = (centr.transpose(0, 1)?.matmul(&centr)? / n)?; // (d, d)
        let mask_diag = {
            let mut m = vec![0f32; d * d];
            for i in 0..d {
                m[i * d + i] = 1.0;
            }
            Tensor::from_vec(m, (d, d), z.device())?
        };
        let fora = ((&cov * &mask_diag)?.sqr()?.sum_all()? / (2.0 * d.max(1) as f64))?;
        let v_norm = (v / d.max(1) as f64)?;
        Ok((v_norm + fora)?)
    }

    /// Perda JEPA padrão para UM lote: contexto (número variável de patches),
    /// lista de blocos-alvo (posições), condição opcional.
    ///
    /// `z_ctx` = incorporação agregada do contexto (b, d_emb);
    /// `alvos`: para cada elemento do lote, Vec<(i_pos, patches_alvo)>.
    pub fn perda_lote(
        &self,
        z_ctx: &Tensor,
        alvos: &[Vec<(usize, Tensor)>],
        cond: Option<&Tensor>,
        z_ctx_todos: &Tensor,
    ) -> Result<(Tensor, f64, f64)> {
        let b = z_ctx.dim(0)?;
        let dev = z_ctx.device();
        let mut perdas: Vec<Tensor> = Vec::new();
        let mut n_alvos = 0usize;
        for (i, blocos) in alvos.iter().enumerate() {
            for (i_pos, patches_alvo) in blocos {
                // ẑ_b = média das incorporações-alvo dos patches do bloco
                let zt = self.alvo_codificar(patches_alvo)?; // (m, d_emb)
                let zt_m = zt.mean(0)?; // (d_emb,)
                let zt_b = zt_m.broadcast_as((1, self.d_emb))?;
                let zc = z_ctx.get(i)?.broadcast_as((1, self.d_emb))?;
                let c_i = cond.map(|c| c.get(i)?.broadcast_as((1, self.d_cond))).transpose()?;
                let pred = self.predizer(&zc, *i_pos, c_i.as_ref())?;
                let dif = (&pred - &zt_b)?;
                let l = (dif.sqr()?.sum_all()? / self.d_emb as f64)?;
                perdas.push(l);
                n_alvos += 1;
            }
        }
        let l_pred = if perdas.is_empty() {
            Tensor::zeros((), candle_core::DType::F32, &dev)?
        } else {
            (Tensor::stack(&perdas, 0)?.sum_all()? / n_alvos.max(1) as f64)?
        };
        let l_vc = self.perda_vc(z_ctx_todos)?;
        let total = (l_pred.affine(1.0, 0.0)? + l_vc.affine(self.lambda_vc, 0.0)?)?;
        let lp = l_pred.to_scalar::<f32>()? as f64;
        let lv = l_vc.to_scalar::<f32>()? as f64;
        Ok((total, lp, lv))
    }
}

// ---------------------------------------------------------------------------
// Parte 2 — máscaras
// ---------------------------------------------------------------------------

/// Máscara de blocos 2D (Image-JEPA): grade (h, w) de patches; sorteia
/// `n_blocos` retângulos aleatórios com fração alvo `frac` de patches.
pub fn mascara_blocos_2d(h: usize, w: usize, n_blocos: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> {
    let total = h * w;
    let alvo = (total as f64 * frac) as usize;
    let mut masc = vec![false; total];
    let mut coberto = 0usize;
    for _ in 0..n_blocos.max(1) {
        if coberto >= alvo {
            break;
        }
        let bh = 1 + rng.inteiro(h.min(4));
        let bw = 1 + rng.inteiro(w.min(4));
        let y0 = rng.inteiro(h + 1 - bh);
        let x0 = rng.inteiro(w + 1 - bw);
        for y in y0..y0 + bh {
            for x in x0..x0 + bw {
                let i = y * w + x;
                if !masc[i] {
                    masc[i] = true;
                    coberto += 1;
                }
            }
        }
    }
    masc
}

/// Máscara de tubos 3D (V-JEPA 2): grade (t, h, w); tubos = retângulos que
/// atravessam o tempo.
pub fn mascara_tubos_3d(t: usize, h: usize, w: usize, n_tubos: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> {
    let total = t * h * w;
    let alvo = (total as f64 * frac) as usize;
    let mut masc = vec![false; total];
    let mut coberto = 0usize;
    for _ in 0..n_tubos.max(1) {
        if coberto >= alvo {
            break;
        }
        let bh = 1 + rng.inteiro(h.min(3));
        let bw = 1 + rng.inteiro(w.min(3));
        let y0 = rng.inteiro(h + 1 - bh);
        let x0 = rng.inteiro(w + 1 - bw);
        for tt in 0..t {
            for y in y0..y0 + bh {
                for x in x0..x0 + bw {
                    let i = (tt * h + y) * w + x;
                    if !masc[i] {
                        masc[i] = true;
                        coberto += 1;
                    }
                }
            }
        }
    }
    masc
}

/// Máscara de bandas (Audio-JEPA): grade (f, tau) — bandas de frequência
/// horizontais com extensão temporal.
pub fn mascara_bandas(f: usize, tau: usize, n_bandas: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> {
    mascara_blocos_2d(f, tau, n_bandas, frac, rng) // blocos na grade (f, tau)
}

/// Máscara de blocos 3D (3D-JEPA): grade (x, y, z) de voxels.
pub fn mascara_blocos_3d(x: usize, y: usize, z: usize, n_blocos: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> {
    mascara_tubos_3d(x, y, z, n_blocos, frac, rng)
}

/// Máscara de segmentos canal×tempo (Brain-JEPA): canais inteiros OU
/// janelas temporais (aleatório por chamada).
pub fn mascara_segmentos(c: usize, t: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> {
    let total = c * t;
    let alvo = (total as f64 * frac) as usize;
    let mut masc = vec![false; total];
    if rng.inteiro(2) == 0 {
        // canais inteiros
        let mut coberto = 0;
        let mut canais: Vec<usize> = (0..c).collect();
        // embaralhar
        for i in (1..canais.len()).rev() {
            let j = rng.inteiro(i + 1);
            canais.swap(i, j);
        }
        for &ch in canais.iter() {
            if coberto >= alvo {
                break;
            }
            for tt in 0..t {
                masc[ch * t + tt] = true;
                coberto += 1;
            }
        }
    } else {
        // janelas temporais
        let mut coberto = 0;
        let mut t0 = 0;
        while coberto < alvo && t0 < t {
            let len = 1 + rng.inteiro((t / 4).max(1));
            for ch in 0..c {
                for tt in t0..(t0 + len).min(t) {
                    masc[ch * t + tt] = true;
                    coberto += 1;
                }
            }
            t0 += len + 1;
        }
    }
    masc
}

/// Máscara 1D (T-JEPA): intervalos contíguos em n posições.
pub fn mascara_intervalos_1d(n: usize, frac: f64, rng: &mut RngDeterministico) -> Vec<bool> {
    let alvo = (n as f64 * frac) as usize;
    let mut masc = vec![false; n];
    let mut coberto = 0;
    while coberto < alvo {
        let len = 1 + rng.inteiro((n / 8).max(1));
        let i0 = rng.inteiro(n);
        for i in i0..(i0 + len).min(n) {
            if !masc[i] {
                masc[i] = true;
                coberto += 1;
            }
        }
    }
    masc
}

/// Amostragem dos pontos mais distantes (FPS) para Point-JEPA.
pub fn pontos_mais_distantes(pontos: &[Vec<f64>], k: usize, rng: &mut RngDeterministico) -> Vec<usize> {
    let n = pontos.len();
    if n == 0 || k == 0 {
        return vec![];
    }
    let k = k.min(n);
    let mut centros = vec![rng.inteiro(n)];
    let mut dist_min: Vec<f64> = pontos
        .iter()
        .map(|p| distancia_eucl(p, &pontos[centros[0]]))
        .collect();
    while centros.len() < k {
        let prox = dist_min
            .iter()
            .enumerate()
            .filter(|(i, _)| !centros.contains(i))
            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
            .map(|(i, _)| i)
            .unwrap_or(0);
        centros.push(prox);
        for (i, p) in pontos.iter().enumerate() {
            let d = distancia_eucl(p, &pontos[prox]);
            if d < dist_min[i] {
                dist_min[i] = d;
            }
        }
    }
    centros
}

fn distancia_eucl(a: &[f64], b: &[f64]) -> f64 {
    a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f64>().sqrt()
}

// ---------------------------------------------------------------------------
// Parte 3 — adaptadores de modalidade
// ---------------------------------------------------------------------------

/// (C, H, W) → patches (H/ph · W/pw, C·ph·pw) — Image-JEPA.
pub fn em_patches_2d(img: &Tensor, ph: usize, pw: usize) -> Result<Tensor> {
    let (c, h, w) = (img.dim(0)?, img.dim(1)?, img.dim(2)?);
    let dev = img.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for y in (0..h).step_by(ph) {
        for x in (0..w).step_by(pw) {
            let hh = ph.min(h - y);
            let ww = pw.min(w - x);
            let patch = img.narrow(0, 0, c)?.narrow(1, y, hh)?.narrow(2, x, ww)?;
            // padding para (C, ph, pw) se borda
            let patch = if hh < ph || ww < pw {
                let mut buf = patch.flatten_all()?.to_vec1::<f32>()?;
                let falta = c * ph * pw - buf.len();
                buf.extend(std::iter::repeat(0f32).take(falta));
                Tensor::from_vec(buf, (c * ph * pw,), &dev)?
            } else {
                patch.flatten_all()?
            };
            linhas.push(patch);
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (T, H, W) → tubos (T/pt · H/ph · W/pw, pt·ph·pw) — V-JEPA 2.
pub fn em_tubos_3d(vid: &Tensor, pt: usize, ph: usize, pw: usize) -> Result<Tensor> {
    let (t, h, w) = (vid.dim(0)?, vid.dim(1)?, vid.dim(2)?);
    let dev = vid.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for tt in (0..t).step_by(pt) {
        let tl = pt.min(t - tt);
        for y in (0..h).step_by(ph) {
            for x in (0..w).step_by(pw) {
                let hh = ph.min(h - y);
                let ww = pw.min(w - x);
                let tubo = vid
                    .narrow(0, tt, tl)?
                    .narrow(1, y, hh)?
                    .narrow(2, x, ww)?
                    .flatten_all()?;
                let mut buf = tubo.to_vec1::<f32>()?;
                let falta = pt * ph * pw - buf.len();
                buf.extend(std::iter::repeat(0f32).take(falta));
                linhas.push(Tensor::from_vec(buf, (pt * ph * pw,), &dev)?);
            }
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (F, τ) → patches (F/pf · τ/pt, pf·pt) — Audio-JEPA.
pub fn em_patches_espectrograma(spec: &Tensor, pf: usize, pt: usize) -> Result<Tensor> {
    let (f, tau) = (spec.dim(0)?, spec.dim(1)?);
    let dev = spec.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for i in (0..f).step_by(pf) {
        for j in (0..tau).step_by(pt) {
            let ff = pf.min(f - i);
            let tt = pt.min(tau - j);
            let p = spec.narrow(0, i, ff)?.narrow(1, j, tt)?.flatten_all()?;
            let mut buf = p.to_vec1::<f32>()?;
            let falta = pf * pt - buf.len();
            buf.extend(std::iter::repeat(0f32).take(falta));
            linhas.push(Tensor::from_vec(buf, (pf * pt,), &dev)?);
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (C, T) → patches canal-independentes (C · T/pt, pt) — T-JEPA/Brain-JEPA.
pub fn em_patches_serie(ts: &Tensor, pt: usize) -> Result<Tensor> {
    let (c, t) = (ts.dim(0)?, ts.dim(1)?);
    let dev = ts.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for ch in 0..c {
        for j in (0..t).step_by(pt) {
            let tt = pt.min(t - j);
            let p = ts.narrow(0, ch, 1)?.narrow(1, j, tt)?.flatten_all()?;
            let mut buf = p.to_vec1::<f32>()?;
            let falta = pt - buf.len();
            buf.extend(std::iter::repeat(0f32).take(falta));
            linhas.push(Tensor::from_vec(buf, (pt,), &dev)?);
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (X, Y, Z) → voxels (X/px · Y/py · Z/pz, px·py·pz) — 3D-JEPA.
pub fn em_voxels(vol: &Tensor, px: usize, py: usize, pz: usize) -> Result<Tensor> {
    let (x, y, z) = (vol.dim(0)?, vol.dim(1)?, vol.dim(2)?);
    let dev = vol.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for i in (0..x).step_by(px) {
        for j in (0..y).step_by(py) {
            for k in (0..z).step_by(pz) {
                let lx = px.min(x - i);
                let ly = py.min(y - j);
                let lz = pz.min(z - k);
                let p = vol.narrow(0, i, lx)?.narrow(1, j, ly)?.narrow(2, k, lz)?.flatten_all()?;
                let mut buf = p.to_vec1::<f32>()?;
                let falta = px * py * pz - buf.len();
                buf.extend(std::iter::repeat(0f32).take(falta));
                linhas.push(Tensor::from_vec(buf, (px * py * pz,), &dev)?);
            }
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// Nuvem de pontos (n, d) → grupos FPS (k, m·d): k centros por FPS, cada
/// grupo = os m vizinhos mais próximos do centro (concatenados) — Point-JEPA.
pub fn em_grupos_pontos(pc: &Tensor, k: usize, m: usize, rng: &mut RngDeterministico) -> Result<Tensor> {
    let n = pc.dim(0)?;
    let d = pc.dim(1)?;
    let coords = pc.to_vec2::<f32>()?;
    let coords_f: Vec<Vec<f64>> = coords.iter().map(|r| r.iter().map(|&v| v as f64).collect()).collect();
    let centros = pontos_mais_distantes(&coords_f, k, rng);
    let dev = pc.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for &ci in &centros {
        // m vizinhos mais próximos do centro ci (incluindo ele próprio)
        let mut ord: Vec<usize> = (0..n).collect();
        ord.sort_by(|&a, &b| {
            let da = distancia_eucl(&coords_f[a], &coords_f[ci]);
            let db = distancia_eucl(&coords_f[b], &coords_f[ci]);
            da.partial_cmp(&db).unwrap()
        });
        let mut buf: Vec<f32> = Vec::with_capacity(m * d);
        for &idx in ord.iter().take(m) {
            for j in 0..d {
                buf.push(coords[idx][j]);
            }
        }
        let falta = m * d - buf.len();
        buf.extend(std::iter::repeat(0f32).take(falta));
        linhas.push(Tensor::from_vec(buf, (m * d,), &dev)?);
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

// ---------------------------------------------------------------------------
// 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_com_semente(
        d_patch,
        cfg.d_emb,
        if cond.is_some() { cond.unwrap().dim(1)? } else { 0 },
        n,
        cfg.m,
        cfg.lambda_vc,
        cfg.gamma_vc,
        cfg.semente,
        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()
}

// ---------------------------------------------------------------------------
// v0.9.0 — fração de máscara ADAPTATIVA por Thompson Sampling
// (crítica Qwen3.8 C; provas T-J1/T-J2 em r-34 §7)
// ---------------------------------------------------------------------------

/// Relatório do treino com fração adaptativa.
#[derive(Debug, Clone)]
pub struct RelatorioJepaAdaptativa {
    pub base: RelatorioJepa,
    pub fracs: Vec<f64>,
    /// (época, braço, frac, perda_pred) por época — trilha auditável.
    pub escolhas: Vec<(usize, usize, f64, f64)>,
    /// nº de épocas por braço.
    pub jogadas_por_braco: Vec<usize>,
    /// médias posteriores finais (estimativa de qualidade relativa).
    pub posterior: Vec<f64>,
}

/// **T-J1 (conexão Information Bottleneck — o porquê de a fração importar):**
/// o JEPA minimiza `E‖s(y) − f(s(x_ctx))‖²` com alvo mascarado. Sob ruído
/// gaussiano homocedástico, minimizar o MSE equivale a maximizar um limite
/// inferior de `I(Z; Y | X_ctx)` (o preditor ótimo é `E[s(Y)|X_ctx]` e o
/// resíduo é o ruído epistêmico — princípio InfoMax). A fração `f` controla o
/// TRADE-OFF: f→0 torna o alvo visível no contexto (atalho de identidade,
/// informação trivial); f→1 destrói o contexto (erro de estimação domina).
/// A perda preditiva em função de f é, portanto, em U — com ótimo dependente
/// dos dados. NENHUM f fixo é ótimo universalmente (por isso a adaptação).
///
/// **T-J2 (esquema e recompensa honesta):** um bandit Thompson (T-TS1) sobre
/// `f ∈ fracs` escolhe a fração de cada época. A recompensa usa a PERDA
/// PREDITIVA ISOLADA `l_p` (critério honesto herdado da v0.5.0 — a VC pode
/// crescer legitimamente) e mede MELHORIA RELATIVA do próprio braço:
/// `r = ½ + (l_prev − l)/(2·l_prev)` ∈ [0,1] — ½ = estagnação, >½ =
/// melhorando. Escalar por l_prev torna r comparável entre regimes de perda
/// diferentes; oscilações simétricas não enganam o posterior em expectativa
/// (r médio = ½). O RNG do bandit é um FLUXO SEPARADO (semente derivada) —
/// o fluxo de máscara (cfg.semente) permanece intocado: com um único braço,
/// o treino é BIT-A-BIT idêntico ao `treinar_variante` com fração fixa
/// (teste `jepa_adaptativa_um_braco_e_fixa`).
pub fn treinar_variante_adaptativa(
    patches: &Tensor,
    grid: &[usize],
    fracs: &[f64],
    cfg: &CfgJepa,
    variante: VarianteJepa,
    cond: Option<&Tensor>,
    dev: &Device,
) -> Result<RelatorioJepaAdaptativa> {
    if fracs.is_empty() {
        candle_core::bail!("treinar_variante_adaptativa: fracs não pode ser vazio");
    }
    let fracs: Vec<f64> = fracs.iter().cloned().map(|f| f.clamp(0.05, 0.95)).collect();
    // fluxo de RNG SEPARADO do bandit — preserva o fluxo de máscara
    let mut bandit = crate::controlador::ThompsonSampling::novo(
        fracs.iter().map(|f| format!("frac={f:.2}")).collect(),
        cfg.semente ^ 0xBADC0DEu64,
    );
    let mut ultima_perda_braco = vec![f64::NAN; fracs.len()];
    let mut escolhas: Vec<(usize, usize, f64, f64)> = Vec::with_capacity(cfg.epocas);

    // ---- corpo: réplica fiel do loop de `treinar_variante` (mesma ordem de
    // ops; única diferença: a fração vem do bandit por época) ----
    let n = patches.dim(0)?;
    let d_patch = patches.dim(1)?;
    let mut rng = RngDeterministico::novo(cfg.semente);
    let mut nucleo = NucleoJepa::novo_com_semente(
        d_patch,
        cfg.d_emb,
        if cond.is_some() { cond.unwrap().dim(1)? } else { 0 },
        n,
        cfg.m,
        cfg.lambda_vc,
        cfg.gamma_vc,
        cfg.semente,
        dev,
    )?;
    let otimo_vars = nucleo.vars();
    let mut otimizador = crate::treino::AdamWManual::novo(otimo_vars, cfg.lr, 0.01);
    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 {
        // === escolha do braço (fração) desta época ===
        let k = bandit.escolher();
        let frac_braco = fracs[k];
        let (ctx_idx, blocos) = loop {
            let mut masc = mascara_por_variante(variante, grid, frac_braco, &mut rng);
            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);
            }
        };
        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)?;
        let z_ctx = z_ctx_todos.mean_keepdim(0)?;
        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)?)]);
        }
        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;
        }
        // === recompensa do braço (T-J2): melhoria relativa da perda PREDITIVA ===
        if lp.is_finite() {
            let anterior = ultima_perda_braco[k];
            let r = if anterior.is_finite() && anterior > 0.0 {
                (0.5 + (anterior - lp) / (2.0 * anterior)).clamp(0.0, 1.0)
            } else {
                0.5 // primeira medida do braço: neutra
            };
            bandit.registrar_continuo(k, r);
            ultima_perda_braco[k] = lp;
            escolhas.push((epoca, k, frac_braco, lp));
        }
        // decaimento cosseno do lr + clipping (idêntico ao fixo)
        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);
    let jogadas_por_braco = bandit.jogadas.iter().map(|&c| c as usize).collect();
    let posterior = bandit.medias_posteriori();
    Ok(RelatorioJepaAdaptativa {
        base: 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,
        },
        fracs,
        escolhas,
        jogadas_por_braco,
        posterior,
    })
}

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)
}

// ---------------------------------------------------------------------------
// Parte 5 — SonicMoE (IO/tile-aware)
// ---------------------------------------------------------------------------

/// MoE top-k com capacidade por especialista e perda de balanceamento.
///
/// **Prova SM-1 (limite de capacidade):** o roteamento tile-aware mantém no
/// máximo `C` tokens por especialista por ladrilho — o drop é determinístico
/// (ordem de chegada) e o nº de tokens descartados é `Σ_e max(0, load_e − C)`.
///
/// **Prova SM-2 (balanceamento):** `L_aux = E·Σ_e f_e·p_e`; com roteamento
/// auto-consistente (`f = p`), `L_aux = E·Σ p_e² ≥ 1` com igualdade sse
/// `p_e = 1/E` (Cauchy–Schchwarz) — o mínimo força o roteador uniforme.
pub struct SonicMoe {
    pub roteador: Linear,            // d → E
    pub especialistas: Vec<MlpJepa>, // E experts d → d
    pub n_top: usize,
    pub capacidade: usize,
    pub lambda_aux: f64,
}

pub struct SaidaMoe {
    pub saida: Tensor,
    pub perda_aux: Tensor,
    pub carga: Vec<usize>,  // tokens efetivamente processados por especialista
    pub descartados: usize, // tokens com especialistas cortados pela capacidade
}

impl SonicMoe {
    pub fn novo(d: usize, n_expertos: usize, n_top: usize, capacidade: usize, dev: &Device) -> Result<Self> {
        let roteador = Linear::sem_vies(d, n_expertos, 0.1, dev)?;
        let especialistas = (0..n_expertos)
            .map(|_| MlpJepa::novo(d, 2 * d, d, dev))
            .collect::<Result<Vec<_>>>()?;
        Ok(Self { roteador, especialistas, n_top, capacidade, lambda_aux: 0.01 })
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = self.roteador.vars();
        for e in &self.especialistas {
            v.extend(e.vars());
        }
        v
    }

    /// Forward tile-aware: (n, d) → (n, d).
    pub fn fwd(&self, x: &Tensor) -> Result<SaidaMoe> {
        let n = x.dim(0)?;
        let d = x.dim(1)?;
        let e = self.especialistas.len();
        let dev = x.device();
        let logits = self.roteador.fwd(x)?; // (n, E)
        let probs = candle_nn::ops::softmax(&logits, candle_core::D::Minus1)?;

        // top-k por token (não diferenciável — pesos do top-k são dif. via probs)
        let probs_v = probs.to_vec2::<f32>()?;
        let mut carga = vec![0usize; e];
        let mut descartados = 0usize;
        // acumulador de saída (n, d) — construído por soma de contribuições
        let mut saida = Tensor::zeros((n, d), candle_core::DType::F32, &dev)?;
        let mut peso_total = Tensor::zeros((n, 1), candle_core::DType::F32, &dev)?;
        // por especialista (tile): tokens cujo top-k contém e e com carga < C
        for ei in 0..e {
            // candidatos: tokens ordenados por prob do especialista ei
            let mut cand: Vec<(usize, f32)> = (0..n)
                .map(|i| (i, probs_v[i][ei]))
                .collect();
            cand.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
            let mut escolhidos: Vec<usize> = Vec::new();
            for (i, _) in cand {
                if carga[ei] >= self.capacidade {
                    break;
                }
                // verifica se ei está no top-k do token i
                let mut ordem: Vec<(f32, usize)> =
                    (0..e).map(|k| (probs_v[i][k], k)).collect();
                ordem.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
                let top: Vec<usize> = ordem.iter().rev().take(self.n_top).map(|(_, k)| *k).collect();
                if top.contains(&ei) {
                    escolhidos.push(i);
                    carga[ei] += 1;
                }
            }
            if escolhidos.is_empty() {
                continue;
            }
            // tile do especialista: uma única passagem em lote
            let idx: Vec<u32> = escolhidos.iter().map(|&i| i as u32).collect();
            let n_e = idx.len();
            let idx_t = Tensor::from_vec(idx, (n_e,), &dev)?;
            let tile = x.index_select(&idx_t, 0)?;
            let y = self.especialistas[ei].fwd(&tile)?; // (m, d)
            // pesos: prob do especialista ei para cada token escolhido
            let w: Vec<f32> = escolhidos.iter().map(|&i| probs_v[i][ei]).collect();
            let n_w = escolhidos.len();
            let w_t = Tensor::from_vec(w.clone(), (n_w, 1), &dev)?;
            let contrib = y.broadcast_mul(&w_t)?;
            // scatter-add: acumula nas posições
            let mut buf = saida.to_vec2::<f32>()?;
            let contrib_v = contrib.to_vec2::<f32>()?;
            for (k, &i) in escolhidos.iter().enumerate() {
                for j in 0..d {
                    buf[i][j] += contrib_v[k][j];
                }
            }
            saida = Tensor::from_vec(
                buf.into_iter().flatten().collect::<Vec<f32>>(),
                (n, d),
                &dev,
            )?;
            let mut pt = peso_total.to_vec2::<f32>()?;
            for (k, &i) in escolhidos.iter().enumerate() {
                pt[i][0] += w[k];
            }
            peso_total = Tensor::from_vec(
                pt.into_iter().flatten().collect::<Vec<f32>>(),
                (n, 1),
                &dev,
            )?;
        }
        // tokens sem especialista (todos cortados) → descartados; saída 0
        let pt_v = peso_total.to_vec2::<f32>()?;
        for i in 0..n {
            if pt_v[i][0] == 0.0 {
                descartados += 1;
            }
        }
        // renormaliza pelo peso efetivo (evita viés da capacidade)
        let denom = peso_total.clamp(1e-6, f64::INFINITY)?;
        let saida = saida.broadcast_div(&denom)?;

        // perda auxiliar: E·Σ f_e·p_e (f = fração de tokens cujo top-1 é e)
        let mut f = vec![0f32; e];
        for i in 0..n {
            let (mut melhor, mut mp) = (0usize, f32::NEG_INFINITY);
            for k in 0..e {
                if probs_v[i][k] > mp {
                    mp = probs_v[i][k];
                    melhor = k;
                }
            }
            f[melhor] += 1.0 / n as f32;
        }
        let f_t = Tensor::from_vec(f, (e,), &dev)?;
        let p_media = probs.mean(0)?; // (E,)
        let perda_aux = ((&f_t * &p_media)?.sum_all()? * (e as f64))?;
        Ok(SaidaMoe { saida, perda_aux, carga, descartados })
    }
}

// ---------------------------------------------------------------------------
// Parte 6 — ReplaySSM
// ---------------------------------------------------------------------------

/// SSM de porta diagonal com replay para aprendizado continuado.
///
/// Recorrência: `h_t = σ(a) ⊙ h_{t−1} + (1 − σ(a)) ⊙ (B x_t)`.
///
/// **Prova R-1 (estabilidade):** `|σ(a)| < 1` ⇒ perturbações de `h_0`
/// decaem geometricamente (mesma estrutura do Teo S-2' do SSD-Δ).
///
/// **Prova R-2 (limite de esquecimento):** o passo de treino mistura
/// gradientes `(1−α)·∇L_novo + α·∇L_replay`; para `L_antiga` μ-fortemente
/// convexa com gradientes novos limitados por `G`, o afastamento do ótimo
/// antigo por passo é `≤ (1−α)·η·G` (sem replay: `≤ η·G`) — o replay
/// contrai o drift por um fator `1/(1−α)` no regime linear.
pub struct ReplaySsm {
    pub a: Var,   // (d,) — porta (sigmoid no uso)
    pub b: Var,   // (d_in, d) — projeção de entrada
    pub d: usize,
    pub alpha_replay: f64,
    pub buffer: Vec<Tensor>,
    pub capacidade_buffer: usize,
}

impl ReplaySsm {
    pub fn novo(d_in: usize, d: usize, alpha_replay: f64, capacidade_buffer: usize, dev: &Device) -> Result<Self> {
        let a_t = Tensor::randn(0f32, 0.5, (d,), dev)?;
        let b_t = Tensor::randn(0f32, 0.1, (d_in, d), dev)?;
        Ok(Self {
            a: Var::from_tensor(&a_t)?,
            b: Var::from_tensor(&b_t)?,
            d,
            alpha_replay,
            buffer: Vec::new(),
            capacidade_buffer,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        vec![self.a.clone(), self.b.clone()]
    }

    /// Scan: x (T, d_in) → h (T, d). Diferenciável (candle autograd).
    pub fn scan(&self, x: &Tensor) -> Result<Tensor> {
        let t = x.dim(0)?;
        let dev = x.device();
        let porta = candle_nn::ops::sigmoid(&self.a.as_tensor().broadcast_as((1, self.d))?)?; // (1, d)
        let um_menos = porta.affine(-1.0, 1.0)?; // 1 − σ(a)
        let entrada = x.matmul(&self.b.as_tensor())?; // (T, d)
        let mut h = Tensor::zeros((1, self.d), candle_core::DType::F32, &dev)?;
        let mut saidas: Vec<Tensor> = Vec::with_capacity(t);
        for i in 0..t {
            let xt = entrada.narrow(0, i, 1)?;
            h = ((&porta * &h)? + (&um_menos * &xt)?)?;
            saidas.push(h.clone());
        }
        Ok(Tensor::stack(&saidas, 0)?)
    }

    /// Empurra uma janela de treino no buffer (FIFO).
    pub fn guardar(&mut self, x: &Tensor) {
        if self.buffer.len() >= self.capacidade_buffer {
            self.buffer.remove(0);
        }
        self.buffer.push(x.detach());
    }

    /// Um passo de treino no lote novo (com mistura de replay).
    /// Devolve (perda_total, perda_novo, perda_replay).
    pub fn passo_treino(
        &mut self,
        x: &Tensor,
        alvo: &Tensor,
        otimizador: &mut crate::treino::AdamWManual,
    ) -> Result<(f64, f64, f64)> {
        // perda nova
        let h = self.scan(x)?;
        let dif = (&h - alvo)?;
        let perda_novo = dif.sqr()?.mean_all()?;
        // replay
        let mut perda_total = perda_novo.clone();
        let mut p_r = 0.0;
        if !self.buffer.is_empty() {
            let idx = self.buffer.len() - 1;
            let antigo = self.buffer[idx].clone();
            let hr = self.scan(&antigo)?;
            // alvo do replay: manter a dinâmica aprendida estável (magnitudes
            // finitas) — auto-supervisão de estabilidade
            let dr = hr.sqr()?.mean_all()?;
            perda_total = (perda_total.affine(1.0 - self.alpha_replay, 0.0)?
                + dr.affine(self.alpha_replay, 0.0)?)?;
            p_r = dr.to_scalar::<f32>()? as f64;
        }
        let pn = perda_novo.to_scalar::<f32>()? as f64;
        otimizador.passo(crate::treino::grads_achatados(&perda_total, &self.vars())?);
        self.guardar(x);
        let pt = perda_total.to_scalar::<f32>()? as f64;
        Ok((pt, pn, p_r))
    }
}

// ---------------------------------------------------------------------------
// 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_com_semente(
        d_patch, cfg.d_emb, d_cond, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, 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_com_semente(d_a, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, dev)?;
    let mut na = NucleoJepa::novo_com_semente(d_b, cfg.d_emb, 0, 1, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, 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_com_semente(
        d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, dev,
    )?;
    // preditor anti-causal: núcleo separado
    let mut ninv = NucleoJepa::novo_com_semente(
        d_patch, cfg.d_emb, 0, 2, cfg.m, cfg.lambda_vc, cfg.gamma_vc, cfg.semente, 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"
}

#[cfg(test)]
mod testes {
    use super::*;

    fn tensor_2d(dados: &[Vec<f64>], dev: &Device) -> Tensor {
        let flat: Vec<f32> = dados.iter().flat_map(|v| v.iter().map(|&x| x as f32)).collect();
        Tensor::from_vec(flat, (dados.len(), dados[0].len()), dev).unwrap()
    }


    /// T-J3 (falsificação da neutralidade): com UM único braço (frac fixa),
    /// o treino adaptativo deve ser BIT-A-BIT idêntico ao `treinar_variante`
    /// com a mesma fração — o RNG do bandit é fluxo separado e não toca o
    /// fluxo de máscara. Qualquer vazamento de aleatoriedade quebra a igualdade.
    #[test]
    fn jepa_adaptativa_um_braco_e_fixa() {
        let dev = Device::Cpu;
        let mut rng = RngDeterministico::novo(3);
        let dados: Vec<Vec<f64>> = (0..96)
            .map(|_| (0..12).map(|_| rng.normal(0.0, 1.0)).collect())
            .collect();
        let t = tensor_2d(&dados, &dev);
        let mut cfg = CfgJepa::default();
        cfg.epocas = 12;
        cfg.d_emb = 16;
        cfg.frac_mascara = 0.35;
        let grid = [12usize, 8];
        let var = VarianteJepa::ImageJepa;
        // fixo
        let rel_fixo = treinar_variante(
            &t,
            &|_n, r| mascara_por_variante(var, &grid, 0.35, r),
            &cfg,
            var,
            None,
            &dev,
        )
        .unwrap();
        // adaptativo com um braço = mesma fração
        let rel_adapt = treinar_variante_adaptativa(
            &t,
            &grid,
            &[0.35],
            &cfg,
            var,
            None,
            &dev,
        )
        .unwrap();
        assert_eq!(
            rel_fixo.perda_pred_final.to_bits(),
            rel_adapt.base.perda_pred_final.to_bits(),
            "perda preditiva divergiu: fixa {} vs adapt {}",
            rel_fixo.perda_pred_final,
            rel_adapt.base.perda_pred_final
        );
        assert_eq!(rel_fixo.perda_final.to_bits(), rel_adapt.base.perda_final.to_bits());
        assert_eq!(rel_adapt.jogadas_por_braco, vec![cfg.epocas]);
        // determinismo do adaptativo multi-braço
        let roda = || {
            treinar_variante_adaptativa(&t, &grid, &[0.2, 0.35, 0.5], &cfg, var, None, &dev)
                .unwrap()
        };
        let a1 = roda();
        let a2 = roda();
        assert_eq!(a1.base.perda_pred_final.to_bits(), a2.base.perda_pred_final.to_bits());
        assert_eq!(a1.escolhas.len(), a2.escolhas.len());
        // multi-braço joga todos os braços ao menos uma vez (exploração TS)
        assert!(a1.jogadas_por_braco.iter().all(|&j| j > 0), "braços não explorados: {:?}", a1.jogadas_por_braco);
    }

    #[test]
    fn teste_mascaras_validas() {
        // fração respeitada ±, particionamento contexto/alvo não-vazio
        let mut rng = RngDeterministico::novo(1);
        let m2 = mascara_blocos_2d(8, 8, 4, 0.3, &mut rng);
        assert_eq!(m2.len(), 64);
        let frac = m2.iter().filter(|&&b| b).count() as f64 / 64.0;
        assert!(frac >= 0.15 && frac <= 0.55, "frac 2d = {frac}");
        assert!(m2.iter().any(|&b| !b), "contexto vazio");

        let m3 = mascara_tubos_3d(4, 6, 6, 3, 0.3, &mut rng);
        assert_eq!(m3.len(), 144);
        let f3 = m3.iter().filter(|&&b| b).count() as f64 / 144.0;
        assert!(f3 >= 0.15 && f3 <= 0.6, "frac 3d = {f3}");

        let ms = mascara_segmentos(4, 32, 0.25, &mut rng);
        assert_eq!(ms.len(), 128);
        assert!(ms.iter().any(|&b| b) && ms.iter().any(|&b| !b));

        let m1 = mascara_intervalos_1d(64, 0.3, &mut rng);
        let f1 = m1.iter().filter(|&&b| b).count() as f64 / 64.0;
        assert!(f1 >= 0.25 && f1 <= 0.45, "frac 1d = {f1}");
    }

    #[test]
    fn teste_fps_cobre() {
        let mut rng = RngDeterministico::novo(2);
        let pontos: Vec<Vec<f64>> = (0..40)
            .map(|i| {
                let a = i as f64 * 0.1;
                vec![a.cos() * 3.0, a.sin() * 3.0, (i % 7) as f64 * 0.3]
            })
            .collect();
        let centros = pontos_mais_distantes(&pontos, 6, &mut rng);
        assert_eq!(centros.len(), 6);
        // centros distintos
        let mut s = centros.clone();
        s.sort();
        s.dedup();
        assert_eq!(s.len(), 6);
    }

    #[test]
    fn teste_j1_colapso_impossivel() {
        // No colapso total (z constante), a perda VC = γ² > 0 — verificando a
        // fórmula implementada com um lote constante.
        let dev = Device::Cpu;
        let nucleo = NucleoJepa::novo(4, 8, 0, 4, 0.99, 1.0, 1.0, &dev).unwrap();
        let z = Tensor::from_vec(vec![0.5f32; 8 * 8], (8, 8), &dev).unwrap();
        let lvc = nucleo.perda_vc(&z).unwrap();
        let v = lvc.to_scalar::<f32>().unwrap() as f64;
        // std=0 ⇒ dobradiça = γ = 1 ⇒ V/d = 1; cov = 0
        assert!(v > 0.9, "perda VC no colapso = {v} (deveria ser ~1)");
    }

    #[test]
    fn teste_j2_ema_ponto_fixo() {
        // encoder congelado ⇒ após muitas passadas EMA, alvo = encoder
        let dev = Device::Cpu;
        let mut nucleo = NucleoJepa::novo(4, 8, 0, 4, 0.5, 0.1, 1.0, &dev).unwrap();
        for _ in 0..40 {
            nucleo.ema_passo().unwrap();
        }
        let a = nucleo.alvo.vars()[0].as_tensor().flatten_all().unwrap().to_vec1::<f32>().unwrap();
        let b = nucleo.encoder.vars()[0].as_tensor().flatten_all().unwrap().to_vec1::<f32>().unwrap();
        let dif: f32 = a.iter().zip(&b).map(|(x, y)| (x - y).abs()).fold(0.0, f32::max);
        assert!(dif < 1e-5, "EMA não convergiu: dif máx = {dif}");
    }

    #[test]
    fn teste_jepa_aprende_latente_previsivel() {
        // J-3: com alvo previsível (função determinística do contexto), a
        // perda pred must cair substancialmente e o std das incorporações
        // NÃO colapsa (J-1).
        let dev = Device::Cpu;
        let mut rng = RngDeterministico::novo(3);
        // patches = 64 vetores 8-D em 4 grupos (estrutura previsível)
        let mut dados: Vec<Vec<f64>> = Vec::new();
        for g in 0..4 {
            let c = g as f64 * 2.0;
            for _ in 0..16 {
                dados.push((0..8).map(|k| c + rng.normal(0.0, 0.3) + (k as f64) * 0.05).collect());
            }
        }
        let patches = tensor_2d(&dados, &dev);
        let cfg = CfgJepa { d_emb: 16, epocas: 40, lr: 5e-3, ..Default::default() };
        let rel = treinar_variante(
            &patches,
            &|n, r| mascara_intervalos_1d(n, 0.3, r),
            &cfg,
            VarianteJepa::TJepa,
            None,
            &dev,
        )
        .unwrap();
        assert!(
            rel.perda_final < rel.perda_inicial * 0.8,
            "perda deveria cair: {} → {}",
            rel.perda_inicial,
            rel.perda_final
        );
        assert!(
            rel.std_emb_final > 0.05,
            "sem colapso: std_final = {}",
            rel.std_emb_final
        );
    }

    #[test]
    fn teste_adaptadores_formas() {
        let dev = Device::Cpu;
        let img = Tensor::randn(0f32, 1.0, (3, 16, 16), &dev).unwrap();
        let p = em_patches_2d(&img, 4, 4).unwrap();
        assert_eq!(p.dims().to_vec(), vec![16, 48]);

        let vid = Tensor::randn(0f32, 1.0, (4, 8, 8), &dev).unwrap();
        let t3 = em_tubos_3d(&vid, 2, 4, 4).unwrap();
        assert_eq!(t3.dims().to_vec(), vec![8, 32]);

        let spec = Tensor::randn(0f32, 1.0, (16, 32), &dev).unwrap();
        let ps = em_patches_espectrograma(&spec, 4, 8).unwrap();
        assert_eq!(ps.dims().to_vec(), vec![16, 32]);

        let ts = Tensor::randn(0f32, 1.0, (4, 32), &dev).unwrap();
        let pt_ = em_patches_serie(&ts, 8).unwrap();
        assert_eq!(pt_.dims().to_vec(), vec![16, 8]);

        let vol = Tensor::randn(0f32, 1.0, (8, 8, 8), &dev).unwrap();
        let vx = em_voxels(&vol, 4, 4, 4).unwrap();
        assert_eq!(vx.dims().to_vec(), vec![8, 64]);

        let pc = Tensor::randn(0f32, 1.0, (50, 3), &dev).unwrap();
        let mut rng = RngDeterministico::novo(5);
        let gr = em_grupos_pontos(&pc, 6, 4, &mut rng).unwrap();
        assert_eq!(gr.dims().to_vec(), vec![6, 12]);
    }

    #[test]
    fn teste_sonic_moe_capacidade_e_balanceamento() {
        // SM-1: carga por especialista ≤ capacidade; SM-2: aux ≥ ~1 quando
        // roteador é ~uniforme (início).
        let dev = Device::Cpu;
        let moe = SonicMoe::novo(8, 4, 2, 6, &dev).unwrap();
        let x = Tensor::randn(0f32, 1.0, (20, 8), &dev).unwrap();
        let s = moe.fwd(&x).unwrap();
        for (e, c) in s.carga.iter().enumerate() {
            assert!(*c <= 6, "especialista {e} com carga {c} > capacidade 6");
        }
        assert_eq!(s.saida.dims().to_vec(), vec![20, 8]);
        let aux = s.perda_aux.to_scalar::<f32>().unwrap() as f64;
        // com roteador quase uniforme, aux ≈ E·Σ (1/E)·(1/E) = 1
        assert!(aux > 0.5 && aux < 3.0, "aux = {aux}");
        // cada token processado por ≤ n_top especialistas
        let total_carga: usize = s.carga.iter().sum();
        assert!(total_carga <= 20 * 2, "carga total {total_carga} > n·topk");
    }

    #[test]
    fn teste_replay_ssm_estabilidade_e_replay() {
        // R-1: estados finitos e limitados; R-2: com replay α>0, após treinar
        // a tarefa nova, a dinâmica antiga (magnitude do estado) permanece
        // estável (não diverge).
        let dev = Device::Cpu;
        let mut ssm = ReplaySsm::novo(4, 6, 0.5, 8, &dev).unwrap();
        let vars = ssm.vars();
        let mut otim = crate::treino::AdamWManual::novo(vars, 1e-2, 0.0);

        // tarefa "antiga": sinal constante 1.0
        let x_antigo = Tensor::from_vec(vec![1.0f32; 16 * 4], (16, 4), &dev).unwrap();
        for _ in 0..30 {
            let h = ssm.scan(&x_antigo).unwrap();
            let alvo = h.detach();
            ssm.passo_treino(&x_antigo, &alvo, &mut otim).unwrap();
        }
        // estado após treino antigo: finito
        let h1 = ssm.scan(&x_antigo).unwrap();
        let m1 = h1.abs().unwrap().max_all().unwrap().to_scalar::<f32>().unwrap();
        assert!(m1.is_finite() && m1 < 50.0, "estado explodiu: {m1}");

        // tarefa "nova": sinal −1.0 — treino com replay
        let x_novo = Tensor::from_vec(vec![-1.0f32; 16 * 4], (16, 4), &dev).unwrap();
        for _ in 0..30 {
            let h = ssm.scan(&x_novo).unwrap();
            let alvo = h.detach();
            ssm.passo_treino(&x_novo, &alvo, &mut otim).unwrap();
        }
        // a dinâmica antiga continua estável (sem explosão)
        let h2 = ssm.scan(&x_antigo).unwrap();
        let m2 = h2.abs().unwrap().max_all().unwrap().to_scalar::<f32>().unwrap();
        assert!(m2.is_finite() && m2 < 50.0, "estado antigo explodiu: {m2}");
        // portas dentro de (0,1)
        let a = ssm.a.as_tensor().to_vec1::<f32>().unwrap();
        for v in a {
            assert!(v.is_finite());
        }
    }
}