Blame for src/jepa/nucleo.rs

//! `jepa::nucleo` — 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 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,
        })
    }

    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> {
        let encoder = MlpJepa::novo(d_patch, 2 * d_emb, d_emb, dev)?;
        let preditor = MlpJepa::novo(d_emb + d_cond + d_emb, 2 * d_emb, d_emb, dev)?;
        // alvo nasce IGUAL ao encoder (ponto de partida do EMA)
        let mut alvo = MlpJepa::novo(d_patch, 2 * d_emb, d_emb, dev)?;
        for (va, vb) in alvo.vars().iter().zip(encoder.vars().iter()) {
            va.set(vb.as_tensor())?;
        }
        let pos_t = Tensor::randn(0f32, 0.02, (n_pos.max(1), 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))
    }
}