//! `jepa::patches` — 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 3 — adaptadores de modalidade
// ---------------------------------------------------------------------------

/// (C, H, W) → patches (H/ph · W/pw, C·ph·pw) — Image-JEPA.
pub fn em_patches_2d(img: &Tensor, ph: usize, pw: usize) -> Result<Tensor> {
    let (c, h, w) = (img.dim(0)?, img.dim(1)?, img.dim(2)?);
    let dev = img.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for y in (0..h).step_by(ph) {
        for x in (0..w).step_by(pw) {
            let hh = ph.min(h - y);
            let ww = pw.min(w - x);
            let patch = img.narrow(0, 0, c)?.narrow(1, y, hh)?.narrow(2, x, ww)?;
            // padding para (C, ph, pw) se borda
            let patch = if hh < ph || ww < pw {
                let mut buf = patch.flatten_all()?.to_vec1::<f32>()?;
                let falta = c * ph * pw - buf.len();
                buf.extend(std::iter::repeat(0f32).take(falta));
                Tensor::from_vec(buf, (c * ph * pw,), &dev)?
            } else {
                patch.flatten_all()?
            };
            linhas.push(patch);
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (T, H, W) → tubos (T/pt · H/ph · W/pw, pt·ph·pw) — V-JEPA 2.
pub fn em_tubos_3d(vid: &Tensor, pt: usize, ph: usize, pw: usize) -> Result<Tensor> {
    let (t, h, w) = (vid.dim(0)?, vid.dim(1)?, vid.dim(2)?);
    let dev = vid.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for tt in (0..t).step_by(pt) {
        let tl = pt.min(t - tt);
        for y in (0..h).step_by(ph) {
            for x in (0..w).step_by(pw) {
                let hh = ph.min(h - y);
                let ww = pw.min(w - x);
                let tubo = vid
                    .narrow(0, tt, tl)?
                    .narrow(1, y, hh)?
                    .narrow(2, x, ww)?
                    .flatten_all()?;
                let mut buf = tubo.to_vec1::<f32>()?;
                let falta = pt * ph * pw - buf.len();
                buf.extend(std::iter::repeat(0f32).take(falta));
                linhas.push(Tensor::from_vec(buf, (pt * ph * pw,), &dev)?);
            }
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (F, τ) → patches (F/pf · τ/pt, pf·pt) — Audio-JEPA.
pub fn em_patches_espectrograma(spec: &Tensor, pf: usize, pt: usize) -> Result<Tensor> {
    let (f, tau) = (spec.dim(0)?, spec.dim(1)?);
    let dev = spec.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for i in (0..f).step_by(pf) {
        for j in (0..tau).step_by(pt) {
            let ff = pf.min(f - i);
            let tt = pt.min(tau - j);
            let p = spec.narrow(0, i, ff)?.narrow(1, j, tt)?.flatten_all()?;
            let mut buf = p.to_vec1::<f32>()?;
            let falta = pf * pt - buf.len();
            buf.extend(std::iter::repeat(0f32).take(falta));
            linhas.push(Tensor::from_vec(buf, (pf * pt,), &dev)?);
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (C, T) → patches canal-independentes (C · T/pt, pt) — T-JEPA/Brain-JEPA.
pub fn em_patches_serie(ts: &Tensor, pt: usize) -> Result<Tensor> {
    let (c, t) = (ts.dim(0)?, ts.dim(1)?);
    let dev = ts.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for ch in 0..c {
        for j in (0..t).step_by(pt) {
            let tt = pt.min(t - j);
            let p = ts.narrow(0, ch, 1)?.narrow(1, j, tt)?.flatten_all()?;
            let mut buf = p.to_vec1::<f32>()?;
            let falta = pt - buf.len();
            buf.extend(std::iter::repeat(0f32).take(falta));
            linhas.push(Tensor::from_vec(buf, (pt,), &dev)?);
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// (X, Y, Z) → voxels (X/px · Y/py · Z/pz, px·py·pz) — 3D-JEPA.
pub fn em_voxels(vol: &Tensor, px: usize, py: usize, pz: usize) -> Result<Tensor> {
    let (x, y, z) = (vol.dim(0)?, vol.dim(1)?, vol.dim(2)?);
    let dev = vol.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for i in (0..x).step_by(px) {
        for j in (0..y).step_by(py) {
            for k in (0..z).step_by(pz) {
                let lx = px.min(x - i);
                let ly = py.min(y - j);
                let lz = pz.min(z - k);
                let p = vol.narrow(0, i, lx)?.narrow(1, j, ly)?.narrow(2, k, lz)?.flatten_all()?;
                let mut buf = p.to_vec1::<f32>()?;
                let falta = px * py * pz - buf.len();
                buf.extend(std::iter::repeat(0f32).take(falta));
                linhas.push(Tensor::from_vec(buf, (px * py * pz,), &dev)?);
            }
        }
    }
    Ok(Tensor::stack(&linhas, 0)?)
}

/// Nuvem de pontos (n, d) → grupos FPS (k, m·d): k centros por FPS, cada
/// grupo = os m vizinhos mais próximos do centro (concatenados) — Point-JEPA.
pub fn em_grupos_pontos(pc: &Tensor, k: usize, m: usize, rng: &mut RngDeterministico) -> Result<Tensor> {
    let n = pc.dim(0)?;
    let d = pc.dim(1)?;
    let coords = pc.to_vec2::<f32>()?;
    let coords_f: Vec<Vec<f64>> = coords.iter().map(|r| r.iter().map(|&v| v as f64).collect()).collect();
    let centros = pontos_mais_distantes(&coords_f, k, rng);
    let dev = pc.device();
    let mut linhas: Vec<Tensor> = Vec::new();
    for &ci in &centros {
        // m vizinhos mais próximos do centro ci (incluindo ele próprio)
        let mut ord: Vec<usize> = (0..n).collect();
        ord.sort_by(|&a, &b| {
            let da = distancia_eucl(&coords_f[a], &coords_f[ci]);
            let db = distancia_eucl(&coords_f[b], &coords_f[ci]);
            da.partial_cmp(&db).unwrap()
        });
        let mut buf: Vec<f32> = Vec::with_capacity(m * d);
        for &idx in ord.iter().take(m) {
            for j in 0..d {
                buf.push(coords[idx][j]);
            }
        }
        let falta = m * d - buf.len();
        buf.extend(std::iter::repeat(0f32).take(falta));
        linhas.push(Tensor::from_vec(buf, (m * d,), &dev)?);
    }
    Ok(Tensor::stack(&linhas, 0)?)
}