//! 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]
}

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,
        }
    }

    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]`. This
    /// is the batch-1 special case of [`forward_batched`](Self::forward_batched)
    /// (used by the autoregressive sampler / inference path).
    pub fn forward(&self, ids: &Tensor) -> Var {
        self.forward_batched(ids, 1)
    }

    /// Batched forward over `batch` sequences of equal length `seq`, flattened to
    /// `[batch*seq]` I32 ids in sequence-major order (sequence 0's `seq` tokens,
    /// then sequence 1's, …). Returns logits `[batch*seq, vocab]` in the SAME flat
    /// layout. The whole graph runs on the flattened tokens so every linear
    /// projection is ONE big `[batch*seq, dim] × [dim, out]` GEMM (the
    /// GPU-filling win); only attention is sequence-aware (per-sequence causal
    /// mask + RoPE position, NO cross-sequence attention).
    pub fn forward_batched(&self, ids: &Tensor, batch: usize) -> Var {
        let total = ids.shape()[0];
        assert_eq!(
            total % batch,
            0,
            "ids len {total} not divisible by batch {batch}"
        );
        let seq = total / batch;

        let mut h = ops::embedding(&self.embed, ids); // [batch*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, batch, 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) // [batch*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)
    }

    /// Batched cross-entropy mean loss: `forward_batched(ids, batch)` against
    /// flat `targets` (`[batch*seq]` I32, same sequence-major layout). The CE mean
    /// is over all `batch*seq` rows — identical to averaging the per-sequence
    /// losses, so the loss value matches the looped single-sequence path.
    pub fn loss_batched(&self, ids: &Tensor, targets: &Tensor, batch: usize) -> Var {
        let logits = self.forward_batched(ids, batch);
        ops::cross_entropy(&logits, targets)
    }

    /// Multi-head causal self-attention over a flattened batch. `x`:[batch*seq,dim]
    /// (already normed), laid out sequence-major. The Q/K/V/O projections are big
    /// `[batch*seq, dim]` GEMMs; the scaled-dot-product attention itself runs as a
    /// fused BATCHED op over the `batch·n_heads` (sequence,head) blocks — each
    /// attends within its own `[seq,seq]` causal window (NO cross-sequence
    /// attention), with RoPE positions reset per sequence (`period = seq`). Causal
    /// masking is applied inside the fused op's softmax kernel (no additive
    /// `[seq,seq]` mask tensor).
    fn attention(&self, b: &Block, x: &Var, batch: usize, seq: usize) -> Var {
        let (nh, hd) = (self.cfg.n_heads, self.cfg.head_dim);
        let total = batch * seq;
        let bh = batch * nh;
        let scale = 1.0 / (hd as f32).sqrt();

        // Project, qk-norm + RoPE, then lay out as a batched [B*nh, seq, hd] tensor.
        //   [B*S,dim] @ [dim,dim] = [B*S,dim]
        //   reshape   [B*S, nh, hd]
        //   qk-norm   per-head RMSNorm over hd (Qwen3-style; Q/K only, before RoPE)
        //   rope      [B*S, nh, hd] with per-sequence position (period = seq)
        //   reshape   [B, S, nh, hd] → transpose(1,2) → [B, nh, S, hd] → [B*nh, S, hd]
        let to_bh = |proj: Var, norm: Option<&Var>| -> Var {
            let r = ops::reshape(&proj, &[total, nh, hd]);
            let r = match norm {
                // Per-head RMSNorm: flatten the (B*S,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, &[total * nh, hd]);
                    let normed = ops::rms_norm(&flat, gamma, self.cfg.eps);
                    let r = ops::reshape(&normed, &[total, nh, hd]);
                    ops::rope(&r, self.cfg.rope_theta, seq)
                }
                None => r,
            };
            let r = ops::reshape(&r, &[batch, seq, nh, hd]);
            let t = ops::transpose_4d12(&r); // [B, nh, S, hd]
            ops::reshape(&t, &[bh, seq, hd]) // [B*nh, S, hd]
        };

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

        // Fused batched causal SDPA over all B*nh (sequence,head) blocks at once
        // (2 batched GEMMs + 1 causal-softmax kernel; no per-head/per-seq loop).
        let out = ops::attention(&q, &k, &v, scale); // [B*nh, S, hd]

        // Back to [B*S, dim]: [B*nh,S,hd] → [B,nh,S,hd] → transpose(1,2) →
        // [B,S,nh,hd] → [B*S, dim].
        let out = ops::reshape(&out, &[batch, nh, seq, hd]);
        let out = ops::transpose_4d12(&out); // [B, S, nh, hd]
        let concat = ops::reshape(&out, &[total, nh * hd]); // [B*S, dim]
        ops::matmul(&concat, &b.wo) // out projection
    }

    /// SwiGLU MLP: `down( silu(gate(x)) ∘ up(x) )`. `x`:[batch*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]
    }
}

/// 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)
}

/// Flatten `batch` equal-length sequences into one `[batch*seq]` I32 tensor in
/// sequence-major order (the layout `forward_batched` expects). Each row of
/// `seqs` is one sequence; all must have the same length.
pub fn batched_ids_tensor(seqs: &[Vec<i32>], device: Device) -> Tensor {
    assert!(!seqs.is_empty(), "empty batch");
    let seq = seqs[0].len();
    let mut flat = Vec::with_capacity(seqs.len() * seq);
    for s in seqs {
        assert_eq!(s.len(), seq, "ragged batch: sequences must be equal length");
        flat.extend_from_slice(s);
    }
    Tensor::from_slice(&flat, &[flat.len()]).to_device(device)
}