//! fnh.rs — Camada FNH (Flexible Network of Hypotheses) — port Rust
//! **melhorado** de `PowerMachine/gru-ring-v13-9-2/xavante_work/xavante/model/
//! {fnh_layer.py, rotation_module.py, smoothquant_compressor.py}` (v0.7.0).
//!
//! Melhorias sobre o original (exigidas pelas provas da casa — r-29):
//! (a) **Ortogonalidade EXATA**: o original usa a série de Taylor com 3 termos
//!     (erro O(θ⁴) — com θ=0,5 pode alcançar ~10⁻²); aqui a série truncada é
//!     PROJETADA de volta ao SO(n) pela retração polar (Newton–Schulz de
//!     `stiefel.rs`), dando δ_ort ≤ 10⁻⁶ comprovado por teste F-1.
//! (b) **Nascimento neutro**: θ inicializado em 0 (o original usa N(0, 0,01))
//!     ⇒ R = I bit-a-bit ⇒ a camada nasce EXATAMENTE igual à linear base para
//!     QUALQUER gate (Teo F-2; padrão da casa Teo R.2).
//! (c) **Gate de punição** (requisito do usuário): a camada só ativa quando a
//!     punição normalizada π_t excede θ_π — `g = σ((π_t−θ_π)/τ_g)` — e o
//!     efeito é convexo; fora da região de punição o gradiente às rotações é
//!     estrangulado por g (homeostase — Teo F-4).
//!
//! Fórmula (r-28 §4.1):
//! ```text
//! s_j   = max|X_j|^α / max|W_j|^(1−α)          (EMA, α=0,5 — SmoothQuant)
//! ŝ_j   = s_j / (∏ s_k)^(1/d)                   (balanço det-preservante)
//! R_wp, R_Atv = polar(expm_trunc(A(θ)))          (SO(n) exato)
//! W_eff = (1−g)·W + g·((W·diag(ŝ)) @ R_wp)      (convexo no gate)
//! y     = Atv((x @ R_Atvᵀ) @ W_effᵀ) + b
//! ```

use crate::camadas::Linear;
use crate::rng::RngDeterministico;
use crate::stiefel::{erro_ortogonalidade, retrair_polar};
use candle_core::{DType, Device, Result, Tensor, Var};

/// Ativações flexíveis do FNH original (learnable fica como extensão futura;
/// gelu não existe em candle-nn 0.11 — tanh-approx idêntica à do PyTorch).
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum Ativacao {
    Gelu,
    Swish,
    Mish,
    Selu,
    Relu,
}

impl Ativacao {
    pub fn aplicar(&self, z: &Tensor) -> Result<Tensor> {
        match self {
            Ativacao::Swish => candle_nn::ops::silu(z),
            Ativacao::Mish => candle_nn::ops::mish(z),
            Ativacao::Selu => candle_nn::ops::selu(z, 1.0507, 1.6733),
            Ativacao::Relu => z.relu(),
            // gelu-tanh: 0,5·z·(1 + tanh(√(2/π)·(z + 0,044715·z³)))
            Ativacao::Gelu => {
                let raiz = (2.0 / std::f64::consts::PI).sqrt();
                let interno = z.affine(1.0, 0.0)?.mul(&(z.powf(3.0)?.affine(0.044715, 0.0)?))?;
                let t = interno.affine(raiz, 0.0)?.tanh()?;
                t.affine(0.5, 0.5)?.mul(z)
            }
        }
    }
}

/// Modo de construção da rotação (v0.8.0 — proposta Qwen-3.8 nº 3).
/// O padrão `Polar` preserva a semântica v0.7.0 bit-a-bit (port fiel do
/// gru-ring: série truncada projetada ao SO(n)); `Cayley` é a alternativa
/// algébrica EXATA sem série infinita (ver `rotacao_cayley`).
#[derive(Clone, Copy, Debug, PartialEq, Default)]
pub enum ModoRotacao {
    /// R = polar(I + A + A²/2 + A³/6) — Newton–Schulz (~82 matmuls n³).
    #[default]
    Polar,
    /// R = (I − A/2)(I + A/2)⁻¹ — Cayley (Teo F-9); inversão pela iteração
    /// de Newton X ← X(2I − MX) com saída antecipada (típico 3–8 matmuls
    /// duplos, teto 12) — diferenciável e EXATA em aritmética exata.
    Cayley,
}

/// Rotação aprendível em SO(n): R = polar(I + A + A²/2 + A³/6) com A skew
/// parametrizada por θ ∈ R^{n(n−1)/2} (port de `SkewRotation`; a retração
/// polar é o acréscimo v0.7.0 — Teo F-1).
pub struct RotacaoSkew {
    /// θ (n(n−1)/2) — INICIALIZADO EM ZERO: R = I bit-a-bit no nascimento.
    pub theta: Var,
    pub n: usize,
    /// clamp de θ por elemento (original: 0,5) — mantém a série truncada
    /// perto da exponencial verdadeira antes da projeção polar.
    pub clamp_theta: f64,
    /// iterações Newton–Schulz da retração polar.
    pub iters_polar: usize,
    /// modo de construção (v0.8.0 — padrão Polar: compatibilidade v0.7.0).
    pub modo: ModoRotacao,
    /// teto de iterações da inversão de Newton do modo Cayley (F-10).
    pub iters_cayley: usize,
}

impl RotacaoSkew {
    pub fn nova(n: usize, dev: &Device) -> Result<Self> {
        let theta = Tensor::zeros((n * (n - 1) / 2,), DType::F32, dev)?;
        Ok(RotacaoSkew {
            theta: Var::from_tensor(&theta)?,
            n,
            clamp_theta: 0.5,
            iters_polar: 40,
            modo: ModoRotacao::Polar,
            iters_cayley: 12,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        vec![self.theta.clone()]
    }

    /// A(θ) skew-simétrica com clamp elemento-a-elemento — construção
    /// DIFERENCIÁVEL por matriz de seleção (matmul): A = U − Uᵀ com U
    /// triangular superior esparsa preenchida por θ (sem slice_assign).
    fn skew(&self) -> Result<Tensor> {
        let dev = self.theta.as_tensor().device();
        let n = self.n;
        let m = n * (n - 1) / 2;
        let th = self
            .theta
            .as_tensor()
            .clamp(-self.clamp_theta as f32, self.clamp_theta as f32)?;
        // S (m, n·n) constante: S[k, i·n+j] = 1 para o k-ésimo par (i<j)
        let mut sel = vec![0f32; m * n * n];
        let mut k = 0usize;
        for i in 0..n {
            for j in (i + 1)..n {
                sel[k * n * n + i * n + j] = 1.0;
                k += 1;
            }
        }
        let s = Tensor::from_vec(sel, (m, n * n), dev)?;
        let tri = th.reshape((1, m))?.matmul(&s)?.reshape((n, n))?; // U(θ)
        Ok(tri.sub(&tri.transpose(0, 1)?)?) // A = U − Uᵀ (skew exato)
    }

    /// R ∈ SO(n) no modo configurado (v0.8.0: despacha para polar ou
    /// Cayley — mesmo contrato: ortogonal exata, nascimento neutro,
    /// diferenciável em θ).
    pub fn rotacao(&self) -> Result<Tensor> {
        match self.modo {
            ModoRotacao::Polar => self.rotacao_polar(),
            ModoRotacao::Cayley => self.rotacao_cayley(),
        }
    }

    /// R = polar(expm_trunc(A)) ∈ SO(n) EXATO (δ_ort ≤ 1e-6 — teste F-1).
    /// Diferenciável em θ (todas as ops são matmul/afim — autograd flui).
    /// Atalho de nascimento: se o valor do truncado é EXATAMENTE I (θ=0 ⇒
    /// I+0+0+0 = I bit-a-bit), devolve o PRÓPRIO trunc — o valor é I exato
    /// (evita a flakiness da σ aleatória da iteração de potência, 3% dos
    /// casos deixa resíduo ~1e-8) e o GRAFO permanece conectado a θ
    /// (gradiente flui — sem deadlock).
    pub fn rotacao_polar(&self) -> Result<Tensor> {
        let a = self.skew()?;
        let a2 = a.matmul(&a)?;
        let a3 = a2.matmul(&a)?;
        let dev = a.device();
        let ident = Tensor::eye(self.n, DType::F32, &dev)?;
        let trunc = ident
            .add(&a)?
            .add(&a2.affine(0.5, 0.0)?)?
            .add(&a3.affine(1.0 / 6.0, 0.0)?)?;
        let dif = (&trunc - &ident)?.abs()?.max_all()?.to_scalar::<f32>()?;
        if dif == 0.0 {
            return Ok(trunc); // nascimento: valor I exato, grafo preservado
        }
        retrair_polar(&trunc, self.iters_polar)
    }

    /// **F-9/F-10 (v0.8.0):** Transformada de Cayley
    /// `R = (I + A/2)(I − A/2)⁻¹` com A = skew(θ) — proposta Qwen-3.8 nº 3
    /// CORRIGIDA em DOIS pontos: (a) candle 0.11 **não possui `Tensor::inv()``
    /// (a proposta assumia a API) — a inversão é feita pela iteração de
    /// Newton para inversos `X ← X(2I − MX)`, 100% diferenciável, com
    /// escalonamento inicial `X₀ = αMᵀ, α = 1/(1+‖B‖²_F)` de convergência
    /// PROVADA (F-10); (b) **convenção de sinal** — a forma original
    /// `(I − A/2)(I + A/2)⁻¹` é ortogonal, mas equivale a `expm(−A + O(A³))`
    /// (identidade escalar (1−t)/(1+t) = e^{−2·artanh t} elevada a matrizes
    /// comutantes); a forma aqui adotada `(I + A/2)(I − A/2)⁻¹ = expm(+A +
    /// O(A³))` preserva a paridade com a série de Taylor do gru-ring
    /// (verificado contra oráculo NumPy/SVD — r-33 §4.6).
    ///
    /// Propriedades (provas r-33 §4):
    /// - **SO(n) exato em aritmética exata**: RᵀR = I e det R = +1
    ///   (comutação de polinômios em A; F-9.ii/iii);
    /// - **Denominador sempre invertível E bem condicionado**: eig(I−A/2)
    ///   = 1 ± iλ/2 ⇒ |eig| = √(1+λ²/4) ≥ 1 ⇒ κ ≤ √(1+‖A‖₂²/4), limitado
    ///   pelo clamp — sem necessidade de LU/pivoteamento (F-9.iv);
    /// - **Nascimento neutro**: θ=0 ⇒ B=0 ⇒ α=1 ⇒ X₀=Mᵀ=I e o resíduo
    ///   |1−α| = 0 já encerra na 1ª iteração ⇒ R = I bit-a-bit COM o grafo
    ///   conectado a θ (gradiente flui no nascimento — melhor que o atalho
    ///   do polar, que precisa de caso especial);
    /// - **Limitação honesta**: −1 nunca é autovalor (ângulo exatamente π é
    ///   inatingível — conjunto de medida zero, irrelevante ao treino).
    pub fn rotacao_cayley(&self) -> Result<Tensor> {
        let a = self.skew()?;
        let dev = a.device();
        let ident = Tensor::eye(self.n, DType::F32, &dev)?;
        let half = a.affine(0.5, 0.0)?; // B = A/2 (skew)
        let m = ident.sub(&half)?; // M = I − A/2 — invertível (F-9.i)
        // X₀ = αMᵀ com α = 1/(1+‖B‖²_F) ⇒ ‖I − MX₀‖₂ < 1 (F-10)
        let bf2 = half.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
        let alpha = 1.0 / (1.0 + bf2);
        let mut x = m.transpose(0, 1)?.affine(alpha, 0.0)?;
        // Newton para inversos: E_{k+1} = E_k² (quadrático); saída
        // antecipada quando o resíduo ‖I − MX‖∞ < 1e-7 (típico: 1–3 iters
        // perto do nascimento; ≤ 9 no pior caso do clamp, teto 12)
        for _ in 0..self.iters_cayley {
            let mx = m.matmul(&x)?;
            let residuo = (&ident - &mx)?.abs()?.max_all()?.to_scalar::<f32>()? as f64;
            if residuo < 1e-7 {
                break;
            }
            // 2I − MX: soma 2 SÓ na diagonal (affine somaria em TODAS as
            // entradas — bug corrigido: 2·J − MX diverge)
            let dois_i_menos = ident.affine(2.0, 0.0)?.sub(&mx)?;
            x = x.matmul(&dois_i_menos)?;
        }
        let num = ident.add(&half)?; // I + A/2
        num.matmul(&x) // R = (I + A/2)(I − A/2)⁻¹ ≈ expm(+A)
    }

    /// δ_ort(R) = ‖RᵀR − I‖_F (monitor).
    pub fn erro_ortogonalidade(&self) -> Result<f64> {
        erro_ortogonalidade(&self.rotacao()?)
    }
}

/// **F-11 (v0.8.0):** rotação por planos canônicos (bloco-diagonal 2×2)
/// via Cayley em fórmula FECHADA — SEM inversão de matriz e SEM iterações.
/// Com s_k = θ_k/2:
/// `R = ⊕_k 1/(1+s_k²) · [[1−s_k², −2s_k], [2s_k, 1−s_k²]]`
/// (inversão analítica exata de [[1, s], [−s, 1]]; convenção expm(+A) —
/// Teo F-11 — r-33 §4.5).
/// `theta_planos`: (n/2,) — a rotação age em n/2 planos canônicos sem
/// mistura entre eles (família RESTRITA, n/2 graus de liberdade).
/// Diferenciável (elementwise + 1 matmul de seleção) e nascimento neutro
/// (θ=0 ⇒ R = I bit-a-bit com grafo conectado a θ). Indicada quando o
/// custo O(n³) das rotações completas domina (n grande).
pub fn rotacao_cayley_bloco(theta_planos: &Tensor) -> Result<Tensor> {
    let dev = theta_planos.device();
    let m = theta_planos.elem_count();
    let n = 2 * m;
    // s = θ/2; c = (1−s²)/(1+s²); t = 2s/(1+s²) — elementwise, grafo VIVO
    let s = theta_planos.flatten_all()?.affine(0.5, 0.0)?;
    let s2 = s.sqr()?;
    let um_menos = s2.affine(-1.0, 1.0)?; // 1−s²
    let um_mais = s2.affine(1.0, 1.0)?; // 1+s²
    let c = um_menos.broadcast_div(&um_mais)?;
    let t = s.affine(2.0, 0.0)?.broadcast_div(&um_mais)?;
    // intercala [c_0, t_0, c_1, t_1, ...] → (1, 2m)
    let par = Tensor::stack(&[&c, &t], 0)? // (2, m)
        .transpose(0, 1)? // (m, 2) = [[c_0, t_0], ...]
        .reshape((1, 2 * m))?;
    // matriz de seleção S (2m, n²) CONSTANTE: linha 2k (c_k) marca as
    // diagonais do bloco k; linha 2k+1 (t_k) marca (−acima, +abaixo)
    // (convenção expm(+A): R_k = [[c, −t], [t, c]])
    let mut sel = vec![0f32; 2 * m * n * n];
    for k in 0..m {
        let (a, b) = (2 * k, 2 * k + 1);
        sel[2 * k * n * n + a * n + a] = 1.0;
        sel[2 * k * n * n + b * n + b] = 1.0;
        sel[(2 * k + 1) * n * n + a * n + b] = -1.0;
        sel[(2 * k + 1) * n * n + b * n + a] = 1.0;
    }
    let s_mat = Tensor::from_vec(sel, (2 * m, n * n), &dev)?;
    Ok(par.matmul(&s_mat)?.reshape((n, n))?)
}

/// SmoothQuant (Xiao et al. 2022) — port do `SmoothQuantCompressor`.
/// s_j = max|X_j|^α / max|W_j|^(1−α) via EMA; a migração compensada
/// (W·diag(s), X/diag(s)) é EXATAMENTE neutra na saída (Teo F-6 — prova
/// algébrica); o balanço ŝ = s/geomean(s) redistribui energia de canais com
/// det preservado (Teo F-7 — AM-GM).
pub struct SuavizadorQuant {
    pub alpha: f64,
    pub beta: f64,
    pub x_max: Option<Vec<f32>>,
    pub w_max: Option<Vec<f32>>,
    pub n_bits: u32,
}

impl SuavizadorQuant {
    pub fn novo(alpha: f64, beta: f64) -> Self {
        SuavizadorQuant { alpha, beta, x_max: None, w_max: None, n_bits: 8 }
    }

    /// Calibração EMA dos máximos por canal de entrada.
    /// w: (d_out, d_in); x: (N, d_in).
    pub fn calibrar(&mut self, w: &Tensor, x: &Tensor) -> Result<()> {
        let d_in = w.dims()[1];
        let mut wmax = vec![1e-5f32; d_in];
        let wv = w.flatten_all()?.to_vec1::<f32>()?;
        for r in 0..w.dims()[0] {
            for c in 0..d_in {
                let a = wv[r * d_in + c].abs();
                if a > wmax[c] {
                    wmax[c] = a;
                }
            }
        }
        let mut xmax = vec![1e-5f32; d_in];
        if x.dims().len() >= 2 && x.dims().last() == Some(&d_in) {
            let xv = x.flatten_all()?.to_vec1::<f32>()?;
            let linhas = xv.len() / d_in;
            for r in 0..linhas {
                for c in 0..d_in {
                    let a = xv[r * d_in + c].abs();
                    if a > xmax[c] {
                        xmax[c] = a;
                    }
                }
            }
        }
        let ema = |atual: &Option<Vec<f32>>, novo: Vec<f32>| -> Vec<f32> {
            match atual {
                None => novo,
                Some(a) => a
                    .iter()
                    .zip(novo.iter())
                    .map(|(o, n)| (*o as f64 * self.beta + *n as f64 * (1.0 - self.beta)) as f32)
                    .collect(),
            }
        };
        self.w_max = Some(ema(&self.w_max, wmax));
        self.x_max = Some(ema(&self.x_max, xmax));
        Ok(())
    }

    /// s_j = max|X_j|^α / max|W_j|^(1−α).
    pub fn escala(&self) -> Vec<f32> {
        let (Some(xm), Some(wm)) = (&self.x_max, &self.w_max) else {
            return vec![1.0];
        };
        xm.iter()
            .zip(wm.iter())
            .map(|(&x, &w)| {
                ((x as f64).powf(self.alpha) / (w as f64).powf(1.0 - self.alpha)) as f32
            })
            .collect()
    }

    /// ŝ = s / geomean(s) — det-preservante (∏ŝ = 1).
    pub fn escala_equilibrada(&self) -> Vec<f32> {
        let s = self.escala();
        if s.is_empty() {
            return s;
        }
        let g = s.iter().map(|v| (*v as f64).ln()).sum::<f64>() / s.len() as f64;
        let gm = g.exp() as f32;
        s.iter().map(|v| v / gm).collect()
    }

    /// Quantização fake simétrica per-tensor com STE explícito:
    /// out = q + (qc − q).detach() — gradiente passa por q (atalho reto),
    /// forward usa o valor clipado qc (padrão QAT).
    pub fn quantizar_fake(&self, t: &Tensor) -> Result<Tensor> {
        let qmax = 2f32.powi(self.n_bits as i32 - 1) - 1.0;
        let amax = t.abs()?.max_all()?.to_scalar::<f32>()?.max(1e-8);
        let escala = amax / qmax;
        let q = t.affine(1.0 / escala as f64, 0.0)?.round()?;
        let qc = q.clamp(-qmax, qmax)?;
        let ste = q.add(&(qc.sub(&q)?.detach()))?;
        Ok(ste.affine(escala as f64, 0.0)?)
    }
}

/// Configuração do gate de punição.
#[derive(Clone, Copy, Debug)]
pub struct GatePunicao {
    pub theta_pi: f64, // limiar de ativação (default 0,6)
    pub tau_g: f64,    // temperatura da sigmoide (default 0,1)
}

impl Default for GatePunicao {
    fn default() -> Self {
        GatePunicao { theta_pi: 0.6, tau_g: 0.1 }
    }
}

impl GatePunicao {
    /// g(π) = max(0, 2σ((π − θ_π)/τ_g) − 1) — **sigmoide retificada**:
    /// g(θ_π) = 0 EXATO (contínua, monótona não-decrescente — Teo F-3) e
    /// g = 0 bit-a-bit para π ≤ θ_π ⇒ gradiente às rotações estrangulado
    /// por construção (homeostase F-4 sem epsilon de ponto flutuante).
    pub fn valor(&self, punicao: f64) -> f64 {
        let z = (punicao - self.theta_pi) / self.tau_g;
        let sig = 1.0 / (1.0 + (-z).exp());
        (2.0 * sig - 1.0).max(0.0)
    }
}

/// Camada FNH — uma hipótese flexível com rotações aprendíveis, suavização de
/// canais e ativação flexível, GATADA pela punição do treino (r-28 §4.1).
pub struct CamadaFNH {
    pub base: Linear, // W (d_out, d_in) + b
    pub r_wp: RotacaoSkew,
    pub r_atv: RotacaoSkew,
    pub suavizador: SuavizadorQuant,
    pub gate_cfg: GatePunicao,
    pub ativacao: Ativacao,
    /// aplica diag(ŝ) no peso quando gate aberto — OPT-IN (default FALSE:
    /// o caminho não-quantizado do FNH original usa SÓ rotações; o balanço
    /// AM-GM é uma decisão explícita de treino, como ativar a ADBA).
    pub equilibrar: bool,
    pub gate_ultimo: f64,
}

impl CamadaFNH {
    pub fn nova(d_in: usize, d_out: usize, dev: &Device) -> Result<Self> {
        Ok(CamadaFNH {
            base: Linear::novo(d_in, d_out, 0.02, dev)?,
            r_wp: RotacaoSkew::nova(d_in, dev)?,
            r_atv: RotacaoSkew::nova(d_in, dev)?,
            suavizador: SuavizadorQuant::novo(0.5, 0.98),
            gate_cfg: GatePunicao::default(),
            ativacao: Ativacao::Gelu,
            equilibrar: false,
            gate_ultimo: 0.0,
        })
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = vec![self.base.w.clone()];
        if let Some(b) = &self.base.b {
            v.push(b.clone());
        }
        v.extend(self.r_wp.vars());
        v.extend(self.r_atv.vars());
        v
    }

    /// Peso efetivo: W_eff = W + g·((W·diag(ŝ))@R_wp − W) — forma RESIDUAL
    /// da mistura convexa (matemática idêntica; numericamente superior):
    /// θ=0 (nascimento) ⇒ R_wp = I, ŝ=1 ⇒ termo = 0 ⇒ W_eff = W bit-a-bit
    /// para QUALQUER gate (Teo F-2 — o mesmo padrão de M-R2; a forma
    /// (1−g)·W + g·W' acumula erro de arredondamento ~1e-9 em f32).
    pub fn peso_efetivo(&self, gate: f64) -> Result<Tensor> {
        let w = self.base.w.as_tensor();
        if gate <= 1e-9 {
            return Ok(w.clone());
        }
        let mut ws = w.clone();
        if self.equilibrar {
            let s = self.suavizador.escala_equilibrada();
            if s.len() == w.dims()[1] && s.iter().all(|v| v.is_finite()) {
                let st = Tensor::from_vec(s, (1, w.dims()[1]), w.device())?;
                ws = w.broadcast_mul(&st)?;
            }
        }
        let rot = self.r_wp.rotacao()?;
        let w_rot = ws.matmul(&rot)?;
        w.add(&w_rot.sub(w)?.affine(gate, 0.0)?)
    }

    /// Forward com punição π_t ∈ [0,1] (medida ANTES do passo — causal, F-4).
    /// x: (..., d_in) → (..., d_out).
    pub fn forward(&mut self, x: &Tensor, punicao: f64) -> Result<Tensor> {
        let gate = self.gate_cfg.valor(punicao);
        self.gate_ultimo = gate;
        let w_eff = self.peso_efetivo(gate)?;
        let x_rot = if gate > 1e-9 {
            let r_atv = self.r_atv.rotacao()?;
            let forma = x.dims().to_vec();
            let n = forma.iter().take(forma.len() - 1).product::<usize>();
            let x2 = x.reshape((n, forma[forma.len() - 1]))?;
            x2.matmul(&r_atv.transpose(0, 1)?)?.reshape(forma)?
        } else {
            x.clone()
        };
        // z = x_rot @ W_effᵀ (usa o achatamento de lotes do Linear)
        let z = self.linear_sem_bias(&x_rot, &w_eff)?;
        let y = self.ativacao.aplicar(&z)?;
        match &self.base.b {
            Some(b) => y.broadcast_add(b.as_tensor()),
            None => Ok(y),
        }
    }

    /// x @ wᵀ com achatamento de lotes — devolve (..., d_out).
    pub fn linear_sem_bias(&self, x: &Tensor, w: &Tensor) -> Result<Tensor> {
        let forma = x.dims().to_vec();
        let d_out = w.dims()[0];
        let mut saida = forma.clone();
        *saida.last_mut().unwrap() = d_out;
        let n = forma.iter().take(forma.len() - 1).product::<usize>();
        let x2 = x.reshape((n, forma[forma.len() - 1]))?;
        let z = x2.matmul(&w.transpose(0, 1)?)?;
        z.reshape(saida)
    }

    /// Caminho quantizado W8A8 (SmoothQuant migrado + fake quant STE) —
    /// equivalente ao `quantize_mode=True` do original. Ativa apenas quando
    /// calibrado; a migração é neutra (F-6): (X/diag(s))@(W·diag(s))ᵀ = X@Wᵀ.
    pub fn forward_quantizado(&mut self, x: &Tensor, punicao: f64) -> Result<Tensor> {
        let gate = self.gate_cfg.valor(punicao);
        self.gate_ultimo = gate;
        let (d_in, _d_out) = (self.base.w.as_tensor().dims()[1], self.base.w.as_tensor().dims()[0]);
        let s = self.suavizador.escala();
        let s_t = if s.len() == d_in && s.iter().all(|v| v.is_finite() && *v > 0.0) {
            Tensor::from_vec(s, (1, d_in), x.device())?
        } else {
            Tensor::ones((1, d_in), DType::F32, x.device())?
        };
        let w_s = self.base.w.as_tensor().broadcast_mul(&s_t)?;
        // x_suave = x / diag(s) — compensa a migração (identidade F-6)
        let x_s = x.broadcast_div(&s_t)?;
        let w_q = self.suavizador.quantizar_fake(&w_s)?;
        let x_q = self.suavizador.quantizar_fake(&x_s)?;
        let z = self.linear_sem_bias(&x_q, &w_q)?;
        let y = self.ativacao.aplicar(&z)?;
        match &self.base.b {
            Some(b) => y.broadcast_add(b.as_tensor()),
            None => Ok(y),
        }
    }
}

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

    fn rng() -> RngDeterministico {
        RngDeterministico::novo(4242)
    }

    fn randn(rng: &mut RngDeterministico, forma: &[usize], dev: &Device) -> Result<Tensor> {
        let n: usize = forma.iter().product();
        let v: Vec<f32> = (0..n).map(|_| rng.normal_padrao() as f32).collect();
        Tensor::from_vec(v, forma, dev)
    }

    /// F-1: R ∈ SO(n) EXATO — δ_ort ≤ 1e-6; det = +1 por continuidade
    /// (polar de matriz próxima de I — o fator polar de M nonsingular tem
    /// det = sign(det M) e det(trunc) > 0 perto de I; ver r-29 F-1).
    #[test]
    fn f1_rotacao_so_n_exata() -> Result<()> {
        let dev = Device::Cpu;
        let r = RotacaoSkew::nova(12, &dev)?;
        // inicializa θ aleatório dentro do clamp
        let th: Vec<f32> = (0..12 * 11 / 2).map(|i| ((i as f32) * 0.07) % 0.5).collect();
        r.theta.set(&Tensor::from_vec(th, (12 * 11 / 2,), &dev)?)?;
        let rot = r.rotacao()?;
        let delta = erro_ortogonalidade(&rot)?;
        assert!(delta < 1e-6, "F-1: δ_ort = {delta} (esperado ≤ 1e-6)");
        // nascimento: θ=0 ⇒ R = I bit-a-bit
        let r0 = RotacaoSkew::nova(8, &dev)?;
        let ident = r0.rotacao()?;
        let dif = (&ident - &Tensor::eye(8, DType::F32, &dev)?)?
            .abs()?
            .max_all()?
            .to_scalar::<f32>()?;
        assert_eq!(dif, 0.0, "F-1: θ=0 ⇒ R = I bit-a-bit (nascimento neutro)");
        Ok(())
    }

    /// F-2: nascimento neutro — θ=0 ⇒ R=I ⇒ camada ≡ linear base
    /// bit-a-bit para QUALQUER punição (com equilibrar=false, o padrão —
    /// caminho fiel ao FNH original não-quantizado; e com equilibrar=true
    /// SEM calibração: ŝ indefinida ⇒ sem balanço).
    #[test]
    fn f2_nascimento_neutro() -> Result<()> {
        let dev = Device::Cpu;
        let mut rng = rng();
        for equilibrar in [false, true] {
            let mut camada = CamadaFNH::nova(8, 6, &dev)?;
            camada.equilibrar = equilibrar;
            let x = randn(&mut rng, &[3, 8], &dev)?;
            for pi in [0.0f64, 0.3, 0.85, 1.0] {
                let y = camada.forward(&x, pi)?;
                let z = camada.linear_sem_bias(&x, camada.base.w.as_tensor())?;
                let y_ref = camada.ativacao.aplicar(&z)?;
                let y_ref = match &camada.base.b {
                    Some(b) => y_ref.broadcast_add(b.as_tensor())?,
                    None => y_ref,
                };
                let dif = (&y - &y_ref)?.abs()?.max_all()?.to_scalar::<f32>()?;
                assert_eq!(
                    dif, 0.0,
                    "F-2: equilibrar={equilibrar}, π={pi}: dif = {dif} (deve ser 0 bit-a-bit)"
                );
            }
        }
        Ok(())
    }

    /// F-3: gate monótono e com zona morta EXATA — π ≤ θ_π ⇒ g = 0;
    /// π ≥ θ_π+6τ ⇒ g > 0,99; monotonicidade.
    #[test]
    fn f3_gate_monotono() {
        let g = GatePunicao::default();
        let pis = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0];
        let mut anterior = -1.0;
        for &pi in &pis {
            let v = g.valor(pi);
            assert!(v >= anterior - 1e-12, "F-3: g deve ser não-decrescente");
            anterior = v;
        }
        for pi in [0.0f64, 0.3, 0.6] {
            assert_eq!(g.valor(pi), 0.0, "F-3: g(π≤θ_π) deve ser EXATAMENTE 0");
        }
        assert!(g.valor(1.0) > 0.95, "F-3: g(1) = {} deve ser ~1", g.valor(1.0));
        assert!(g.valor(1.4) > 0.999, "F-3: saturação além do domínio");
    }

    /// F-4: homeostase — com g=0 o gradiente às rotações é ZERO (congelado);
    /// com g=1 é não-nulo (aprendizado ativo). Causalidade: π medida antes.
    #[test]
    fn f4_homeostase_gradiente() -> Result<()> {
        let dev = Device::Cpu;
        let mut rng = rng();
        let mut camada = CamadaFNH::nova(6, 4, &dev)?;
        let x = randn(&mut rng, &[2, 6], &dev)?;
        // π baixa: gradiente a θ é zero
        let y = camada.forward(&x, 0.0)?;
        let perda = y.sqr()?.sum_all()?;
        let grads = perda.backward()?;
        let g_theta = grads
            .get(&camada.r_wp.theta)
            .map(|t| t.flatten_all().unwrap().to_vec1::<f32>().unwrap())
            .unwrap_or_default();
        assert!(
            g_theta.iter().all(|v| v.abs() < 1e-9),
            "F-4: com g=0 o gradiente a θ deve ser 0"
        );
        // θ ≠ 0 para haver sinal: gradiente a θ flui com g>0
        let th: Vec<f32> = (0..6 * 5 / 2).map(|i| 0.1 + 0.01 * i as f32).collect();
        camada
            .r_wp
            .theta
            .set(&Tensor::from_vec(th, (6 * 5 / 2,), &dev)?)?;
        let y2 = camada.forward(&x, 1.0)?;
        let perda2 = y2.sqr()?.sum_all()?;
        let grads2 = perda2.backward()?;
        let g2 = grads2
            .get(&camada.r_wp.theta)
            .map(|t| t.flatten_all().unwrap().to_vec1::<f32>().unwrap())
            .unwrap_or_default();
        assert!(
            g2.iter().any(|v| v.abs() > 1e-8),
            "F-4: com g=1 o gradiente a θ deve fluir (homeostase)"
        );
        Ok(())
    }

    /// F-6: migração compensada é identidade — (X/s)@(W·s)ᵀ = X@Wᵀ.
    #[test]
    fn f6_migracao_neutra() -> Result<()> {
        let dev = Device::Cpu;
        let mut rng = rng();
        let w = randn(&mut rng, &[5, 7], &dev)?;
        let x = randn(&mut rng, &[4, 7], &dev)?;
        let s: Vec<f32> = (1..=7).map(|i| 0.3 * i as f32).collect();
        let st = Tensor::from_vec(s.clone(), (1, 7), &dev)?;
        let ws = w.broadcast_mul(&st)?;
        let xs = x.broadcast_div(&st)?;
        let a = xs.matmul(&ws.transpose(0, 1)?)?;
        let b = x.matmul(&w.transpose(0, 1)?)?;
        let dif = (&a - &b)?.abs()?.max_all()?.to_scalar::<f32>()?;
        assert!(dif < 1e-4, "F-6: dif da migração compensada = {dif}");
        Ok(())
    }

    /// F-7: balanço det-preservante equaliza a energia das colunas
    /// (AM-GM: Σ_j ŝ_j²‖W_:,j‖² mínimo quando ŝ_j‖W_:,j‖ iguais).
    #[test]
    fn f7_balanceamento_am_gm() -> Result<()> {
        let dev = Device::Cpu;
        let mut rng = rng();
        let mut suav = SuavizadorQuant::novo(0.5, 0.98);
        // W com energias de coluna MUITO desiguais
        let mut wv = Vec::new();
        for r in 0..4 {
            for c in 0..6 {
                let escala = (c + 1) as f32; // colunas 5× maiores que a 1ª
                wv.push(rng.normal_padrao() as f32 * escala);
            }
        }
        let w = Tensor::from_vec(wv, (4, 6), &dev)?;
        let xv: Vec<f32> = (0..3 * 6).map(|_| rng.normal_padrao() as f32).collect();
        let x = Tensor::from_vec(xv, (3, 6), &dev)?;
        suav.calibrar(&w, &x)?;
        let s = suav.escala_equilibrada();
        // produto de ŝ = 1 (det preservado)
        let prod: f64 = s.iter().map(|v| (*v as f64).ln()).sum::<f64>().exp();
        assert!((prod - 1.0).abs() < 1e-3, "F-7: ∏ŝ = {prod} (deve ser 1)");
        // energias das colunas após balanço mais iguais que antes
        let wv = w.flatten_all()?.to_vec1::<f32>()?;
        let energia = |escala: &dyn Fn(usize) -> f32| -> (f32, f32) {
            let mut mins = f32::MAX;
            let mut maxs = 0f32;
            for c in 0..6 {
                let e: f32 = (0..4)
                    .map(|r| (wv[r * 6 + c] * escala(c)).powi(2))
                    .sum::<f32>()
                    .sqrt();
                mins = mins.min(e);
                maxs = maxs.max(e);
            }
            (mins, maxs)
        };
        let antes = energia(&|_| 1.0);
        let depois = energia(&|c| s[c]);
        let razao_antes = antes.1 / antes.0.max(1e-9);
        let razao_depois = depois.1 / depois.0.max(1e-9);
        assert!(
            razao_depois < razao_antes,
            "F-7: razão max/min de energia {razao_depois} deve < {razao_antes}"
        );
        Ok(())
    }

    /// F-8: forward quantizado calibra e roda finito (W8A8 fake com STE).
    #[test]
    fn f8_forward_quantizado_finito() -> Result<()> {
        let dev = Device::Cpu;
        let mut rng = rng();
        let mut camada = CamadaFNH::nova(8, 5, &dev)?;
        let x = randn(&mut rng, &[3, 8], &dev)?;
        camada.suavizador.calibrar(camada.base.w.as_tensor(), &x)?;
        let y = camada.forward_quantizado(&x, 0.9)?;
        let mx = y.abs()?.max_all()?.to_scalar::<f32>()?;
        assert!(mx.is_finite() && mx > 0.0, "F-8: saída quantizada finita");
        Ok(())
    }

    /// determinante por eliminação gaussiana em f64 (auxiliar dos testes).
    fn det_f64(t: &Tensor) -> f64 {
        let n = t.dims()[0];
        let mut m: Vec<Vec<f64>> = t
            .to_vec2::<f32>()
            .unwrap()
            .iter()
            .map(|r| r.iter().map(|x| *x as f64).collect())
            .collect();
        let mut det = 1.0;
        for col in 0..n {
            let mut pivo = col;
            for r in col + 1..n {
                if m[r][col].abs() > m[pivo][col].abs() {
                    pivo = r;
                }
            }
            if m[pivo][col].abs() < 1e-12 {
                return 0.0;
            }
            if pivo != col {
                m.swap(pivo, col);
                det = -det;
            }
            det *= m[col][col];
            for r in col + 1..n {
                let f = m[r][col] / m[col][col];
                for c in col..n {
                    m[r][c] -= f * m[col][c];
                }
            }
        }
        det
    }

    /// inversão Gauss–Jordan em f64 (auxiliar — oráculo do teste F-10).
    fn inv_f64(t: &Tensor) -> Vec<Vec<f64>> {
        let n = t.dims()[0];
        let mut a: Vec<Vec<f64>> = t
            .to_vec2::<f32>()
            .unwrap()
            .iter()
            .map(|r| r.iter().map(|x| *x as f64).collect())
            .collect();
        let mut inv = vec![vec![0.0; n]; n];
        for i in 0..n {
            inv[i][i] = 1.0;
        }
        for col in 0..n {
            let mut pivo = col;
            for r in col + 1..n {
                if a[r][col].abs() > a[pivo][col].abs() {
                    pivo = r;
                }
            }
            a.swap(pivo, col);
            inv.swap(pivo, col);
            let p = a[col][col];
            for c in 0..n {
                a[col][c] /= p;
                inv[col][c] /= p;
            }
            for r in 0..n {
                if r != col && a[r][col].abs() > 1e-15 {
                    let f = a[r][col];
                    for c in 0..n {
                        a[r][c] -= f * a[col][c];
                        inv[r][c] -= f * inv[col][c];
                    }
                }
            }
        }
        inv
    }

    /// **F-9 (v0.8.0):** Cayley ∈ SO(n) — δ_ort ≤ 1e-6, det = +1, e
    /// nascimento I bit-a-bit COM o grafo conectado a θ (gradiente flui —
    /// propriedade que a forma fechada dá de graça, sem o atalho especial
    /// do caminho polar). Perda assimétrica ⟨W, R⟩ para expor o sinal.
    #[test]
    fn f9_cayley_so_n_exato() -> Result<()> {
        let dev = Device::Cpu;
        let r = RotacaoSkew::nova(12, &dev)?;
        let th: Vec<f32> = (0..12 * 11 / 2).map(|i| ((i as f32) * 0.07) % 0.5).collect();
        r.theta.set(&Tensor::from_vec(th, (12 * 11 / 2,), &dev)?)?;
        let rot = r.rotacao_cayley()?;
        let delta = erro_ortogonalidade(&rot)?;
        assert!(delta < 1e-6, "F-9: δ_ort = {delta} (esperado ≤ 1e-6)");
        let d = det_f64(&rot);
        assert!((d - 1.0).abs() < 1e-4, "F-9: det = {d} (deve ser +1)");
        // nascimento: θ=0 ⇒ R = I bit-a-bit (α=1 exato ⇒ 1ª iteração encerra)
        let r0 = RotacaoSkew::nova(8, &dev)?;
        let ident = r0.rotacao_cayley()?;
        let dif = (&ident - &Tensor::eye(8, DType::F32, &dev)?)?
            .abs()?
            .max_all()?
            .to_scalar::<f32>()?;
        assert_eq!(dif, 0.0, "F-9: θ=0 ⇒ R = I bit-a-bit");
        // grafo VIVO no nascimento: ⟨W, R⟩ com W não-simétrico ⇒ grad ≠ 0
        let mut rng2 = rng();
        let w = randn(&mut rng2, &[8, 8], &dev)?;
        let perda = (&w * &ident)?.sum_all()?;
        let grads = perda.backward()?;
        let g = grads
            .get(&r0.theta)
            .map(|t| t.flatten_all().unwrap().to_vec1::<f32>().unwrap())
            .unwrap_or_default();
        assert!(
            g.iter().any(|v| v.abs() > 1e-9),
            "F-9: gradiente flui a θ no nascimento (grafo conectado sem atalho)"
        );
        Ok(())
    }

    /// **F-10 (v0.8.0):** Newton-inverso converge no PIOR CASO do clamp
    /// (todos os θ = 0,5) e o resultado confere com a inversão explícita
    /// Gauss–Jordan em f64.
    #[test]
    fn f10_newton_inverso_pior_caso() -> Result<()> {
        let dev = Device::Cpu;
        let n = 16;
        let r = RotacaoSkew::nova(n, &dev)?;
        let th: Vec<f32> = vec![0.5f32; n * (n - 1) / 2];
        r.theta.set(&Tensor::from_vec(th, (n * (n - 1) / 2,), &dev)?)?;
        let rot = r.rotacao_cayley()?;
        let delta = erro_ortogonalidade(&rot)?;
        assert!(delta < 1e-6, "F-10: δ_ort pior caso = {delta}");
        // oráculo explícito: R = (I + B)·(I − B)⁻¹ em f64 (convenção expm(+A))
        let a = r.skew()?;
        let av = a.to_vec2::<f32>()?;
        let mut b = vec![vec![0.0f64; n]; n];
        for i in 0..n {
            for j in 0..n {
                b[i][j] = 0.5 * av[i][j] as f64;
            }
        }
        let bt = Tensor::from_vec(
            b.iter().flatten().map(|x| *x as f32).collect::<Vec<f32>>(),
            (n, n),
            &dev,
        )?;
        let um_b = Tensor::eye(n, DType::F32, &dev)?.add(&bt)?;
        let um_menos_b = Tensor::eye(n, DType::F32, &dev)?.sub(&bt)?;
        let inv = inv_f64(&um_menos_b);
        let inv_t = Tensor::from_vec(
            inv.iter().flatten().map(|x| *x as f32).collect::<Vec<f32>>(),
            (n, n),
            &dev,
        )?;
        let explicito = um_b.matmul(&inv_t)?;
        let dif = (&rot - &explicito)?.abs()?.max_all()?.to_scalar::<f32>()?;
        assert!(dif < 1e-3, "F-10: max|R_newton − R_gaussjordan| = {dif}");
        Ok(())
    }

    /// **F-11 (v0.8.0):** bloco 2×2 — ortogonal exata, bloque-diagonal
    /// (zero fora dos planos), nascimento I bit-a-bit, gradiente flui.
    #[test]
    fn f11_cayley_bloco() -> Result<()> {
        let dev = Device::Cpu;
        let th = Tensor::from_vec(vec![0.3f32, -0.7, 1.2, 0.05], (4,), &dev)?;
        let r = rotacao_cayley_bloco(&th)?;
        assert_eq!(r.dims(), [8, 8], "F-11: 4 planos ⇒ 8×8");
        let delta = erro_ortogonalidade(&r)?;
        assert!(delta < 1e-6, "F-11: δ_ort = {delta}");
        // estrutura bloque-diagonal: zero EXATO fora dos blocos 2×2
        let rv = r.to_vec2::<f32>()?;
        for i in 0..8 {
            for j in 0..8 {
                if (i / 2) != (j / 2) {
                    assert_eq!(rv[i][j], 0.0, "F-11: fora do bloco ({i},{j})");
                }
            }
        }
        // nascimento: θ=0 ⇒ R = I bit-a-bit
        let th_var = Var::from_tensor(&Tensor::zeros((4,), DType::F32, &dev)?)?;
        let rz = rotacao_cayley_bloco(th_var.as_tensor())?;
        let dif = (&rz - &Tensor::eye(8, DType::F32, &dev)?)?
            .abs()?
            .max_all()?
            .to_scalar::<f32>()?;
        assert_eq!(dif, 0.0, "F-11: nascimento neutro bit-a-bit");
        // grafo vivo: gradiente de ⟨W, R⟩ em θ flui
        let mut rng2 = rng();
        let w = randn(&mut rng2, &[8, 8], &dev)?;
        let perda = (&w * &rz)?.sum_all()?;
        let grads = perda.backward()?;
        let g = grads
            .get(&th_var)
            .map(|t| t.flatten_all().unwrap().to_vec1::<f32>().unwrap())
            .unwrap_or_default();
        assert!(
            g.iter().any(|v| v.abs() > 1e-9),
            "F-11: gradiente flui a θ (grafo conectado)"
        );
        Ok(())
    }

    /// **F-12 (v0.8.0):** Cayley ≈ polar para θ pequeno (ambos ≈ expm(+A) —
    /// com a convenção corrigida F-9(b): polar-série com erro O(θ⁴),
    /// Cayley = expm(2·artanh(A/2)) = expm(+A − A³/12 + O(A⁵)); a distância
    /// mútua é O(‖A‖³/12), verificada contra oráculo NumPy/SVD — r-33 §4.6).
    #[test]
    fn f12_cayley_vs_polar_teta_pequeno() -> Result<()> {
        let dev = Device::Cpu;
        let n = 8;
        let r = RotacaoSkew::nova(n, &dev)?;
        let th: Vec<f32> = (0..n * (n - 1) / 2).map(|i| 0.02 + 0.005 * i as f32).collect();
        r.theta.set(&Tensor::from_vec(th, (n * (n - 1) / 2,), &dev)?)?;
        let rp = r.rotacao_polar()?;
        let rc = r.rotacao_cayley()?;
        let dif = (&rp - &rc)?.sqr()?.sum_all()?.to_scalar::<f32>()?.sqrt();
        let d_id = (&rp - &Tensor::eye(n, DType::F32, &dev)?)?
            .sqr()?
            .sum_all()?
            .to_scalar::<f32>()?
            .sqrt();
        assert!(
            dif < 2e-2,
            "F-12: ‖R_polar − R_cayley‖_F = {dif} (θ pequeno deve coincidir)"
        );
        // honestidade: desvio O(‖A‖³/12) contra distância O(‖A‖) à identidade
        let razao = dif / d_id.max(1e-9);
        assert!(razao < 0.25, "F-12: razão desvio/distância = {razao}");
        Ok(())
    }
}