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