//! The tiny transformer forward graph + parameter container (Phase T5).

#![cfg(not(no_cuda))]

use crate::config::Config;
use xtrain_autodiff::ops;
use xtrain_autodiff::tape::Var;
use xtrain_tensor::{Device, Tensor};

/// One decoder block's learnable tensors.
struct Block {
    attn_norm: Var, // [dim]
    wq: Var,        // [dim, dim]
    wk: Var,        // [dim, dim]
    wv: Var,        // [dim, dim]
    q_norm: Var,    // [head_dim] — per-head QK-norm (Qwen3-style)
    k_norm: Var,    // [head_dim]
    wo: Var,        // [dim, dim]
    ffn_norm: Var,  // [dim]
    w_gate: Var,    // [dim, ffn_hidden]
    w_up: Var,      // [dim, ffn_hidden]
    w_down: Var,    // [ffn_hidden, dim]
}

/// A tiny RoPE+RMSNorm+SwiGLU decoder. Holds every parameter as a leaf [`Var`];
/// `forward` builds an autograd graph over them.
pub struct TinyTransformer {
    cfg: Config,
    embed: Var, // [vocab, dim]
    blocks: Vec<Block>,
    final_norm: Var, // [dim]
    lm_head: Var,    // [dim, vocab]
    device: Device,
}

impl TinyTransformer {
    /// Build a model with parameters initialised from `init(shape) -> host data`.
    /// The caller controls initialisation (deterministic for tests / PyTorch
    /// parity). `init` receives the logical shape and returns row-major data.
    pub fn new(cfg: Config, device: Device, mut init: impl FnMut(&[usize]) -> Vec<f32>) -> Self {
        let leaf = |data: Vec<f32>, shape: &[usize]| -> Var {
            Var::leaf(Tensor::from_slice(&data, shape).to_device(device))
        };
        let mut mk = |shape: &[usize]| -> Var {
            let data = init(shape);
            assert_eq!(data.len(), shape.iter().product::<usize>(), "init size");
            leaf(data, shape)
        };

        let embed = mk(&[cfg.vocab, cfg.dim]);
        let blocks = (0..cfg.n_layers)
            .map(|_| Block {
                attn_norm: mk(&[cfg.dim]),
                wq: mk(&[cfg.dim, cfg.dim]),
                wk: mk(&[cfg.dim, cfg.dim]),
                wv: mk(&[cfg.dim, cfg.dim]),
                q_norm: mk(&[cfg.head_dim]),
                k_norm: mk(&[cfg.head_dim]),
                wo: mk(&[cfg.dim, cfg.dim]),
                ffn_norm: mk(&[cfg.dim]),
                w_gate: mk(&[cfg.dim, cfg.ffn_hidden]),
                w_up: mk(&[cfg.dim, cfg.ffn_hidden]),
                w_down: mk(&[cfg.ffn_hidden, cfg.dim]),
            })
            .collect();
        let final_norm = mk(&[cfg.dim]);
        let lm_head = mk(&[cfg.dim, cfg.vocab]);

        Self {
            cfg,
            embed,
            blocks,
            final_norm,
            lm_head,
            device,
        }
    }

    pub fn config(&self) -> &Config {
        &self.cfg
    }

    /// All learnable parameters, in a stable order. The optimizer (a hand-written
    /// GD step in T5, AdamW in T6) iterates this; each holds its `.grad()` after
    /// `backward()`.
    pub fn params(&self) -> Vec<Var> {
        let mut ps = vec![self.embed.clone()];
        for b in &self.blocks {
            ps.extend([
                b.attn_norm.clone(),
                b.wq.clone(),
                b.wk.clone(),
                b.wv.clone(),
                b.q_norm.clone(),
                b.k_norm.clone(),
                b.wo.clone(),
                b.ffn_norm.clone(),
                b.w_gate.clone(),
                b.w_up.clone(),
                b.w_down.clone(),
            ]);
        }
        ps.push(self.final_norm.clone());
        ps.push(self.lm_head.clone());
        ps
    }

    /// Forward over a single sequence of token `ids` (`[seq]` I32 on this
    /// model's device). Returns the logits [`Var`] of shape `[seq, vocab]`.
    pub fn forward(&self, ids: &Tensor) -> Var {
        let seq = ids.shape()[0];
        let mask = self.causal_mask(seq);

        let mut h = ops::embedding(&self.embed, ids); // [seq, dim]
        for b in &self.blocks {
            // --- Attention sub-block (pre-norm + residual) ---
            let normed = ops::rms_norm(&h, &b.attn_norm, self.cfg.eps);
            let attn = self.attention(b, &normed, &mask, seq);
            h = ops::add(&h, &attn);

            // --- MLP sub-block (pre-norm + residual) ---
            let normed = ops::rms_norm(&h, &b.ffn_norm, self.cfg.eps);
            let mlp = self.swiglu_mlp(b, &normed);
            h = ops::add(&h, &mlp);
        }

        let h = ops::rms_norm(&h, &self.final_norm, self.cfg.eps);
        ops::matmul(&h, &self.lm_head) // [seq, vocab]
    }

    /// Cross-entropy mean loss of `forward(ids)` against `targets` (`[seq]` I32).
    pub fn loss(&self, ids: &Tensor, targets: &Tensor) -> Var {
        let logits = self.forward(ids);
        ops::cross_entropy(&logits, targets)
    }

    /// Multi-head causal self-attention. `x`:[seq,dim] (already normed).
    fn attention(&self, b: &Block, x: &Var, mask: &Var, seq: usize) -> Var {
        let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
        let scale = 1.0 / (hd as f32).sqrt();

        // Project, then lay out as per-head [seq, head_dim] tensors.
        //   [seq,dim] @ [dim,dim] = [seq,dim]
        //   reshape   [seq, nh, hd]
        //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
        //   rope      (kernel expects exactly [tokens, heads, head_dim])
        //   transpose [nh, seq, hd]  → split into nh × [seq, hd]
        let to_heads = |proj: Var, norm: Option<&Var>| -> Vec<Var> {
            let r = ops::reshape(&proj, &[seq, nh, hd]);
            let r = match norm {
                // Per-head RMSNorm: flatten the (seq,nh) head rows, norm over hd,
                // restore. RoPE follows on the normed Q/K (mirrors xserv qwen3.rs).
                Some(gamma) => {
                    let flat = ops::reshape(&r, &[seq * nh, hd]);
                    let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
                    let r = ops::reshape(&normed, &[seq, nh, hd]);
                    ops::rope(&r, self.cfg.rope_theta)
                }
                None => r,
            };
            let t = ops::transpose_3d01(&r); // [nh, seq, hd]
            ops::split_heads(&t)
        };

        let q = to_heads(ops::matmul(x, &b.wq), Some(&b.q_norm));
        let k = to_heads(ops::matmul(x, &b.wk), Some(&b.k_norm));
        let v = to_heads(ops::matmul(x, &b.wv), None);

        // Per-head scaled-dot-product attention with causal mask.
        let heads_out: Vec<Var> = (0..nh)
            .map(|i| {
                let kt = ops::transpose_2d(&k[i]); // [hd, seq]
                let scores = ops::scale(&ops::matmul(&q[i], &kt), scale); // [seq,seq]
                let scores = ops::add(&scores, mask); // causal
                let probs = ops::softmax(&scores);
                ops::matmul(&probs, &v[i]) // [seq, hd]
            })
            .collect();

        // Stack heads back: nh × [seq,hd] → [nh,seq,hd] → [seq,nh,hd] → [seq,dim].
        let merged = ops::merge_heads(&heads_out); // [nh, seq, hd]
        let t = ops::transpose_3d01(&merged); // [seq, nh, hd]
        let concat = ops::reshape(&t, &[seq, nh * hd]); // [seq, dim]
        ops::matmul(&concat, &b.wo) // out projection
    }

    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[seq,dim].
    fn swiglu_mlp(&self, b: &Block, x: &Var) -> Var {
        let gate = ops::matmul(x, &b.w_gate); // [seq, ffn_hidden]
        let up = ops::matmul(x, &b.w_up); // [seq, ffn_hidden]
        let act = ops::swiglu(&gate, &up); // silu(gate) ∘ up
        ops::matmul(&act, &b.w_down) // [seq, dim]
    }

    /// Additive causal mask `[seq,seq]`: 0 on/below the diagonal, −1e9 above it
    /// (so softmax zeros out future positions). A constant leaf (no grad needed,
    /// but harmless if it accumulates one — it has no consumers downstream of x).
    fn causal_mask(&self, seq: usize) -> Var {
        let mut m = vec![0.0f32; seq * seq];
        for i in 0..seq {
            for j in (i + 1)..seq {
                m[i * seq + j] = -1.0e9;
            }
        }
        Var::leaf(Tensor::from_slice(&m, &[seq, seq]).to_device(self.device))
    }
}

/// Materialise a parameter's value back to a host `Vec<f32>` (for the GD step
/// and PyTorch parity export).
pub fn param_to_host(v: &Var) -> Vec<f32> {
    v.value().to_device(Device::Cpu).as_slice::<f32>().to_vec()
}

/// Build an I32 id tensor on `device` from token ids.
pub fn ids_tensor(ids: &[i32], device: Device) -> Tensor {
    Tensor::from_slice(ids, &[ids.len()]).to_device(device)
}
