//! Camadas base: Linear, Embedding, RMSNorm, SwiGLU — port minimalista e
//! autocontido (sem depender do API de módulos do candle-nn para reduzir
//! superfície de drift entre versões).

use candle_core::{DType, Device, Result, Tensor, Var};
use candle_nn::ops::rms_norm;

pub struct Linear {
    pub w: Var, // (d_saida, d_entrada)
    pub b: Option<Var>,
}

impl Linear {
    pub fn novo(d_entrada: usize, d_saida: usize, desvio: f64, dev: &Device) -> Result<Self> {
        let w = Tensor::randn(0f32, desvio as f32, (d_saida, d_entrada), dev)?;
        let b = Tensor::zeros((d_saida,), DType::F32, dev)?;
        Ok(Linear { w: Var::from_tensor(&w)?, b: Some(Var::from_tensor(&b)?) })
    }

    pub fn sem_vies(d_entrada: usize, d_saida: usize, desvio: f64, dev: &Device) -> Result<Self> {
        let w = Tensor::randn(0f32, desvio as f32, (d_saida, d_entrada), dev)?;
        Ok(Linear { w: Var::from_tensor(&w)?, b: None })
    }

    /// v0.9.0 — `sem_vies` com RNG DETERMINÍSTICO próprio (quita o débito
    /// r-33 §8.1: a inicialização do JEPA usava o RNG global do candle — não
    /// semeado — e impedia bit-identidade entre execuções).
    pub fn sem_vies_semeado(
        d_entrada: usize,
        d_saida: usize,
        desvio: f64,
        dev: &Device,
        rng: &mut crate::rng::RngDeterministico,
    ) -> Result<Self> {
        let n = d_entrada * d_saida;
        let dados: Vec<f32> = (0..n).map(|_| rng.normal(0.0, desvio) as f32).collect();
        let w = Tensor::from_vec(dados, (d_saida, d_entrada), dev)?;
        Ok(Linear { w: Var::from_tensor(&w)?, b: None })
    }

    /// W = 0 no nascimento — para camadas cujo efeito deve nascer NEUTRO
    /// (ex.: W_f do FiLM, Teo R.2: tanh(0·h)=0 ⇒ identidade bit-a-bit).
    pub fn zeros_sem_vies(d_entrada: usize, d_saida: usize, dev: &Device) -> Result<Self> {
        let w = Tensor::zeros((d_saida, d_entrada), DType::F32, dev)?;
        Ok(Linear { w: Var::from_tensor(&w)?, b: None })
    }

    /// x (..., d_entrada) → (..., d_saida). Achata lotes (candle matmul
    /// não aceita 3-D @ 2-D) e devolve na forma original.
    pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
        let forma: Vec<usize> = x.dims().to_vec();
        let d_saida = self.w.as_tensor().dims()[0];
        let x2 = if forma.len() > 2 {
            let n = forma.iter().take(forma.len() - 1).product::<usize>();
            x.reshape((n, forma[forma.len() - 1]))?
        } else {
            x.clone()
        };
        let y = x2.matmul(&self.w.as_tensor().transpose(0, 1)?)?;
        let y = match &self.b {
            Some(b) => y.broadcast_add(b.as_tensor())?,
            None => y,
        };
        if forma.len() > 2 {
            let mut alvo = forma;
            *alvo.last_mut().unwrap() = d_saida;
            y.reshape(alvo)
        } else {
            Ok(y)
        }
    }

    pub fn vars(&self) -> Vec<Var> {
        match &self.b {
            Some(b) => vec![self.w.clone(), b.clone()],
            None => vec![self.w.clone()],
        }
    }
}

pub struct Embedding {
    pub tabela: Var, // (vocab, d)
}

impl Embedding {
    pub fn novo(vocab: usize, d: usize, dev: &Device) -> Result<Self> {
        let t = Tensor::randn(0f32, 0.02, (vocab, d), dev)?;
        Ok(Embedding { tabela: Var::from_tensor(&t)? })
    }

    /// ids u32 (…) → (..., d) — lookup via index_select no plano.
    pub fn fwd(&self, ids: &Tensor) -> Result<Tensor> {
        let forma: Vec<usize> = ids.dims().to_vec();
        let d = self.tabela.as_tensor().dims()[1];
        let plano = ids.reshape((ids.elem_count(),))?;
        let linhas = self.tabela.as_tensor().index_select(&plano, 0)?;
        let mut alvo = forma;
        alvo.push(d);
        linhas.reshape(alvo)
    }

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

pub struct RmsNorm {
    pub ganho: Var,
    pub eps: f64,
}

impl RmsNorm {
    pub fn novo(d: usize, eps: f64, dev: &Device) -> Result<Self> {
        let g = Tensor::ones((d,), DType::F32, dev)?;
        Ok(RmsNorm { ganho: Var::from_tensor(&g)?, eps })
    }

    pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
        rms_norm(x, self.ganho.as_tensor(), self.eps as f32)
    }

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

/// SwiGLU: W2(silu(W1·x) ⊙ W3·x) — d_ff oculto.
pub struct Swiglu {
    pub w1: Linear,
    pub w3: Linear,
    pub w2: Linear,
}

impl Swiglu {
    pub fn novo(d: usize, d_ff: usize, dev: &Device) -> Result<Self> {
        let desvio = (1.0 / (d as f64).sqrt()) * 0.9;
        Ok(Swiglu {
            w1: Linear::sem_vies(d, d_ff, desvio, dev)?,
            w3: Linear::sem_vies(d, d_ff, desvio, dev)?,
            w2: Linear::sem_vies(d_ff, d, desvio / (d_ff as f64).sqrt() * 10.0, dev)?,
        })
    }

    pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
        let a = self.w1.fwd(x)?;
        let b = self.w3.fwd(x)?;
        let gate = candle_nn::ops::silu(&a)?;
        self.w2.fwd(&(&gate * &b)?)
    }

    pub fn vars(&self) -> Vec<Var> {
        [self.w1.vars(), self.w3.vars(), self.w2.vars()].concat()
    }
}