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

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