//! Máscaras e interpolação multiscale (port de `conceitos/mascaras.py` do n-04 §2.2;
//! contratos dos Teos D.1(ii)/B.2 — verificação V4: vazamento zero).

use candle_core::{D, DType, Device, Result, Tensor};

/// Máscara aditiva block-causal B (Teo D.1(ii)): o slot do bloco ℓ atende ao
/// prefixo (sempre visível) e a TODOS os blocos ≤ ℓ (bidirecional no bloco
/// próprio). `tamanhos` = [n_1, …, n_L]; `n_prefixo` = |cond|.
///
/// Retorna matriz (T+P, T+P) f32 com 0 onde permitido e f32::MIN onde bloqueado,
/// pronta para `scores + mask` (broadcast).
pub fn mascara_block_causal(tamanhos: &[usize], n_prefixo: usize, dev: &Device) -> Result<Tensor> {
    let mut cortes = vec![n_prefixo];
    for &t in tamanhos {
        let ult = *cortes.last().unwrap();
        cortes.push(ult + t);
    }
    let total = *cortes.last().unwrap();
    let mut m = vec![f32::MIN; total * total];
    for (bloco_i, (&ini_i, &fim_i)) in cortes.iter().zip(cortes.iter().skip(1)).enumerate() {
        for i in ini_i..fim_i {
            // prefixo sempre visível
            for j in 0..n_prefixo {
                m[i * total + j] = 0.0;
            }
            // blocos ≤ bloco_i inteiros visíveis
            let fim_visivel = cortes[bloco_i + 1];
            for j in 0..fim_visivel {
                m[i * total + j] = 0.0;
            }
        }
    }
    Tensor::from_vec(m, (total, total), dev)?.to_dtype(DType::F32)
}

/// Verificação V4 — nenhuma posição do bloco ℓ acessa bloco > ℓ.
/// Devolve o número de vazamentos (deve ser 0).
pub fn contar_vazamentos(tamanhos: &[usize], n_prefixo: usize, mascara: &Tensor) -> Result<usize> {
    let total = mascara.dim(0)?;
    let dados = mascara.flatten_all()?.to_vec1::<f32>()?;
    let mut cortes = vec![n_prefixo];
    for &t in tamanhos {
        let ult = *cortes.last().unwrap();
        cortes.push(ult + t);
    }
    let mut vazamentos = 0;
    for (bloco_i, (&ini_i, &fim_i)) in cortes.iter().zip(cortes.iter().skip(1)).enumerate() {
        let inicio_proibido = cortes.get(bloco_i + 2).copied().unwrap_or(*cortes.last().unwrap());
        for i in ini_i..fim_i {
            for j in inicio_proibido..total {
                if dados[i * total + j] >= 0.0 {
                    vazamentos += 1;
                }
            }
        }
    }
    Ok(vazamentos)
}

/// Matriz de interpolação linear `Interp_n` (n-02 §0): R^{n_de × d} → R^{n_para × d}.
/// - downsample (n_para < n_de): média por segmento com pesos de sobreposição (área),
///   linhas somam 1 (contrato: "down = avg-pool por segmento");
/// - upsample: interpolação linear entre os 2 vizinhos.
/// Linear, determinística, sem parâmetros (gradiente flui por matmul).
pub fn matriz_interp(n_de: usize, n_para: usize, dev: &Device) -> Result<Tensor> {
    let mut w = vec![0f32; n_para * n_de];
    if n_de == n_para {
        for i in 0..n_para {
            w[i * n_de + i] = 1.0;
        }
    } else if n_para < n_de {
        // média por segmento com bordas fracionárias
        let razao = n_de as f64 / n_para as f64;
        for i in 0..n_para {
            let ini = i as f64 * razao;
            let fim = ini + razao;
            for j in 0..n_de {
                let jj = j as f64;
                let inter = (fim.min(jj + 1.0)) - (ini.max(jj));
                if inter > 0.0 {
                    w[i * n_de + j] = (inter / razao) as f32;
                }
            }
        }
    } else {
        // upsample linear
        let razao = n_de as f64 / n_para as f64;
        for i in 0..n_para {
            let centro = (i as f64 + 0.5) * razao - 0.5;
            let j0 = centro.floor().clamp(0.0, (n_de - 1) as f64) as usize;
            let j1 = (j0 + 1).min(n_de - 1);
            let fr = (centro - j0 as f64).clamp(0.0, 1.0);
            w[i * n_de + j0] = (1.0 - fr) as f32;
            w[i * n_de + j1] += fr as f32;
        }
    }
    Tensor::from_vec(w, (n_para, n_de), dev)?.to_dtype(DType::F32)
}

/// Interp_m(x): redimensiona a última dimensão espacial (n_de → n_para) de
/// x com forma (..., n_de, d) via produto com a matriz (broadcast).
pub fn interp_para(x: &Tensor, n_para: usize, dev: &Device) -> Result<Tensor> {
    let n_de = x.dim(D::Minus2)?;
    let w = matriz_interp(n_de, n_para, dev)?;
    // (n_para, n_de) x (..., n_de, d) -> (..., n_para, d)
    w.broadcast_matmul(x)
}

#[cfg(test)]
mod testes {
    use super::*;

    #[test]
    fn mascara_sem_vazamento() -> candle_core::Result<()> {
        let dev = Device::Cpu;
        let tamanhos = [4usize, 8, 16];
        let m = mascara_block_causal(&tamanhos, 3, &dev)?;
        assert_eq!(contar_vazamentos(&tamanhos, 3, &m)?, 0);
        // posição do bloco 3 (último) vê tudo
        let total = 3 + 4 + 8 + 16;
        let dados = m.flatten_all()?.to_vec1::<f32>()?;
        let ultima = (total - 1) * total;
        assert!(dados[ultima..ultima + total].iter().all(|&v| v == 0.0));
        Ok(())
    }

    #[test]
    fn interp_preserva_constante() -> candle_core::Result<()> {
        let dev = Device::Cpu;
        // downsample de vetor constante → constante (linhas somam 1)
        let x = Tensor::from_vec(vec![2.0f32; 8], (8, 1), &dev)?;
        let y = interp_para(&x, 4, &dev)?;
        let v = y.flatten_all()?.to_vec1::<f32>()?;
        assert!(v.iter().all(|&a| (a - 2.0).abs() < 1e-5), "{v:?}");
        Ok(())
    }
}