//! Memória: SOM de conceitos + JEPA latente + RM-lite por nível — port
//! minimalista dos contratos do KHTST (docs 01/13/17/28; Teos 17.4, 28.5–28.8,
//! D.3, D.5). Tudo com stop-grad no caminho pós-geração (V14-análogo).

use candle_core::{Device, Result, Tensor, Var};
use crate::camadas::Linear;

/// SOM 2D online sobre o espaço dos códigos de conceito (contrato doc 01 +
/// Teo 17.4: protótipos atualizados com stop-grad; EQ médio monotônico
/// esperado). Grid lado×lado em d_lat.
pub struct SomConceitos {
    pub prototipos: Tensor, // (lado*lado, d_lat) — buffer, SEM gradiente
    pub lado: usize,
    pub lr: f64,
    pub eq_medio: f64,
}

impl SomConceitos {
    pub fn novo(lado: usize, d_lat: usize, dev: &Device) -> Result<Self> {
        Ok(SomConceitos {
            prototipos: Tensor::randn(0f32, 0.1, (lado * lado, d_lat), dev)?,
            lado,
            lr: 0.05,
            eq_medio: 0.0,
        })
    }

    /// Insere vetores (n, d_lat) com atualização online (sem gradiente — V14).
    pub fn inserir(&mut self, vetores: &Tensor) -> Result<f64> {
        let p = self.prototipos.flatten_all()?.to_vec1::<f32>()?;
        let v = vetores.flatten_all()?.to_vec1::<f32>()?;
        let (n, d) = (vetores.dim(0)?, vetores.dim(1)?);
        let total = self.lado * self.lado;
        let mut eq_soma = 0f64;
        let mut protos = p;
        for i in 0..n {
            // BMU
            let (mut melhor, mut dist_melhor) = (0usize, f32::INFINITY);
            for c in 0..total {
                let mut dist = 0f32;
                for t in 0..d {
                    let dif = v[i * d + t] - protos[c * d + t];
                    dist += dif * dif;
                }
                if dist < dist_melhor {
                    dist_melhor = dist;
                    melhor = c;
                }
            }
            eq_soma += dist_melhor as f64;
            // vizinhança gaussiana no grid
            let (bx, by) = (melhor % self.lado, melhor / self.lado);
            for c in 0..total {
                let (cx, cy) = (c % self.lado, c / self.lado);
                let dist2 = ((cx as i32 - bx as i32).pow(2) + (cy as i32 - by as i32).pow(2)) as f32;
                let w = (-dist2 / 2.0).exp();
                for t in 0..d {
                    protos[c * d + t] += (self.lr as f32 * w) * (v[i * d + t] - protos[c * d + t]);
                }
            }
        }
        self.eq_medio = eq_soma / n.max(1) as f64;
        let dev = self.prototipos.device();
        self.prototipos = Tensor::from_vec(protos, (total, d), dev)?;
        Ok(self.eq_medio)
    }
}

/// JEPA latente: preditor do resíduo do PRÓXIMO nível no espaço latente
/// (contrato doc 30 — jepa_familia; alvo com stop-grad; perda de cosseno).
pub struct JepaLatente {
    pub preditor: Linear, // d_lat → d_lat
}

impl JepaLatente {
    pub fn novo(d_lat: usize, dev: &Device) -> Result<Self> {
        Ok(JepaLatente { preditor: Linear::sem_vies(d_lat, d_lat, 0.02, dev)? })
    }

    /// Perda JEPA entre reconstruções cumulativas de níveis adjacentes:
    /// L = 1 − cos(pred(ẑ_{ℓ−1} interp), sg[ẑ_ℓ − r_ℓ]) — prevê o resíduo.
    pub fn perda(&self, zhat_ant: &Tensor, alvo: &Tensor) -> Result<Tensor> {
        let pred = self.preditor.fwd(zhat_ant)?;
        let alvo_sg = alvo.detach();
        // cosseno por linha
        let dot = (&pred * &alvo_sg)?.sum(candle_core::D::Minus1)?;
        let np = pred.sqr()?.sum(candle_core::D::Minus1)?.sqrt()?;
        let na = alvo_sg.sqr()?.sum(candle_core::D::Minus1)?.sqrt()?;
        let cos = dot.broadcast_div(&(&np * &na)?)?;
        (&cos * -1.0)?.mean_all()?.affine(1.0, 1.0)
    }

    pub fn vars(&self) -> Vec<Var> {
        self.preditor.vars()
    }
}

/// RM-lite por nível (Teo D.3 — condições de Borkar): calibração de confiança
/// por nível com atualização em duas escalas (τ rápidos, sinal lento).
pub struct RmNivel {
    pub medias_confianca: Vec<f64>, // por nível (escala lenta)
    pub tau_por_nivel: Vec<f64>,    // por nível (escala rápida — bissersão D.3)
    pub contagem: usize,
}

impl RmNivel {
    pub fn novo(l: usize) -> Self {
        RmNivel { medias_confianca: vec![0.5; l], tau_por_nivel: vec![1.0; l], contagem: 0 }
    }

    /// Atualização duas-escalas (Borkar Thm 6.2): sinal lento = média de
    /// confiança; rápidos = τ por bisseção contra entropia alvo por nível.
    pub fn passo(&mut self, confiancas_por_nivel: &[f64], entropias_por_nivel: &[f64], alvo: f64) {
        self.contagem += 1;
        let a_lento = 1.0 / (self.contagem as f64).sqrt().max(1.0); // Σa_t=∞, Σa_t²<∞
        for (i, c) in confiancas_por_nivel.iter().enumerate() {
            if i < self.medias_confianca.len() {
                self.medias_confianca[i] += a_lento * (c - self.medias_confianca[i]);
            }
        }
        for (i, h) in entropias_por_nivel.iter().enumerate() {
            if i < self.tau_por_nivel.len() {
                // ajuste proporcional rápido (proxy da bisseção — D.3(ii))
                let delta = h - alvo;
                self.tau_por_nivel[i] = (self.tau_por_nivel[i] * (1.0 - 0.1 * delta)).clamp(0.5, 1.5);
            }
        }
    }
}