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


#[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()
    }

    #[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());
        }
    }
}