//! Perdas multi-tarefa: PCGrad (Teo D.7 / Teo 35.1 do projeto) e λ com
//! hipergradiente-âncora exato (Teo 29.2). Port de `treino/perdas.py`.
//!
//! PCGrad: para cada par (g_i, g_j) com conflito (⟨g_i,g_j⟩ < 0), projeta
//! g_i no complemento ortogonal de g_j: g_i' = g_i − (⟨g_i,g_j⟩/‖g_j‖²)·g_j.
//! Invariantes I1/I2/I3 de Yu et al. 2020: projeção exata, contração,
//! não-piora por par. Com λ=0 o gradiente extra é nulo (bit-a-bit — T10(ii)).

use candle_core::{Result, Tensor, Var};

/// Gradiente achatado de cada perda w.r.t. às variáveis (ordem fixa de `vars`).
pub fn gradientes_por_perda(perdas: &[&Tensor], vars: &[Var]) -> Result<Vec<Vec<Tensor>>> {
    let mut saida = Vec::with_capacity(perdas.len());
    for p in perdas {
        let gs = p.backward()?;
        let mut um = Vec::with_capacity(vars.len());
        for v in vars {
            let g = gs.get(v).ok_or_else(|| {
                candle_core::Error::Msg("variável sem gradiente (λ=0 ou grafo desconectado)".into())
            })?;
            um.push(g.clone());
        }
        saida.push(um);
    }
    Ok(saida)
}

/// PCGrad sobre gradientes por perda (lista de listas, mesmas formas por índice).
/// Devolve o gradiente combinado (média das projeções por perda).
pub fn pcgrad(grads: &[Vec<Tensor>]) -> Result<Vec<Tensor>> {
    let n_perdas = grads.len();
    let n_vars = grads[0].len();
    let mut combinado: Vec<Option<Tensor>> = (0..n_vars).map(|_| None).collect();
    for i in 0..n_perdas {
        // g_i pode ser projetado contra cada g_j (j ≠ i)
        let mut proj = grads[i].clone();
        for j in 0..n_perdas {
            if i == j {
                continue;
            }
            // produto interno global ⟨g_i, g_j⟩ e ‖g_j‖²
            let mut ip = 0f64;
            let mut nj = 0f64;
            for v in 0..n_vars {
                let a = grads[i][v].flatten_all()?.to_vec1::<f32>()?;
                let b = grads[j][v].flatten_all()?.to_vec1::<f32>()?;
                for k in 0..a.len() {
                    ip += a[k] as f64 * b[k] as f64;
                    nj += b[k] as f64 * b[k] as f64;
                }
            }
            if ip < 0.0 && nj > 1e-12 {
                let alfa = ip / nj;
                for v in 0..n_vars {
                    let b = grads[j][v].flatten_all()?;
                    let escala = Tensor::full(alfa as f32, b.shape(), b.device())?;
                    let sub = grads[i][v].broadcast_sub(&b.mul(&escala)?)?;
                    proj[v] = sub.reshape(grads[i][v].shape())?;
                }
            }
        }
        for v in 0..n_vars {
            let media = proj[v].affine(1.0 / n_perdas as f64, 0.0)?;
            combinado[v] = Some(match &combinado[v] {
                Some(a) => (a + &media)?,
                None => media,
            });
        }
    }
    Ok(combinado.into_iter().map(|t| t.unwrap()).collect())
}

/// Verificação V8-análogo: a projeção PCGrad é exata (⟨g_i', g_j⟩ = 0).
pub fn verificar_projecao(gi: &[Tensor], gj: &[Tensor]) -> Result<bool> {
    let mut ip = 0f64;
    for (a, b) in gi.iter().zip(gj.iter()) {
        let av = a.flatten_all()?.to_vec1::<f32>()?;
        let bv = b.flatten_all()?.to_vec1::<f32>()?;
        for k in 0..av.len() {
            ip += av[k] as f64 * bv[k] as f64;
        }
    }
    Ok(ip.abs() < 1e-3)
}

/// λ com hipergradiente-âncora (Teo 29.2): L(λ) = λ·L_diff + ρ(λ−λ₀)² é
/// afim+quadrática em λ ⇒ hipergradiente EXATO ∂L/∂λ = L_diff + 2ρ(λ−λ₀);
/// ótimo interior λ* = λ₀ − L_diff/(2ρ). Começa em 0 e cresce até λ*.
pub struct LambdaAncora {
    pub lambda: f64,
    pub lambda0: f64,
    pub rho: f64,
    pub minimo: f64,
    pub maximo: f64,
}

impl LambdaAncora {
    pub fn novo(rho: f64) -> Self {
        LambdaAncora { lambda: 0.0, lambda0: 1.0, rho, minimo: 0.0, maximo: 2.0 }
    }

    /// Um passo de gradiente descendente em λ com hipergradiente exato.
    pub fn passo(&mut self, l_diff: f64, taxa: f64) -> f64 {
        let hg = l_diff + 2.0 * self.rho * (self.lambda - self.lambda0);
        self.lambda = (self.lambda - taxa * hg).clamp(self.minimo, self.maximo);
        self.lambda
    }

    pub fn otimo_analitico(&self, l_diff: f64) -> f64 {
        (self.lambda0 - l_diff / (2.0 * self.rho)).clamp(self.minimo, self.maximo)
    }
}

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

    #[test]
    fn lambda_ancora_otimo_interior() {
        let mut la = LambdaAncora::novo(0.05);
        let l_diff = -0.2; // perda nova REDUZ a total (queremos λ > 0)
        let alvo = la.otimo_analitico(l_diff);
        for _ in 0..3000 {
            la.passo(l_diff, 0.05);
        }
        assert!((la.lambda - alvo).abs() < 2e-2, "{} vs {}", la.lambda, alvo);
    }
}