use burn::module::Module;
use burn::module::ParamId;
use burn::nn::{Embedding, EmbeddingConfig, RmsNorm, RmsNormConfig};
use burn::prelude::*;
use burn::tensor::activation::{silu, softmax};
use burn::tensor::Bool;

use crate::config::{HfConfig, LoraConfig};
use crate::lora::LoraLinear;

/// RoPE with the Hugging Face Llama "half-split" convention: head_dim is
/// split in half and rotated as (x1, x2) -> (-x2, x1).
/// Note: `burn::nn::RotaryEncoding` uses the interleaved convention (it
/// rotates adjacent pairs), which is incompatible with these weights.
#[derive(Module, Debug)]
pub struct Rope<B: Backend> {
    cos: Tensor<B, 2>,
    sin: Tensor<B, 2>,
}

impl<B: Backend> Rope<B> {
    pub fn new(head_dim: usize, max_seq: usize, theta: f32, device: &B::Device) -> Self {
        let half = head_dim / 2;
        let inv_freq: Vec<f32> = (0..half)
            .map(|i| 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32))
            .collect();
        let inv_freq = Tensor::<B, 1>::from_floats(inv_freq.as_slice(), device);
        let pos = Tensor::<B, 1, Int>::arange(0..max_seq as i64, device).float();

        // [max_seq, half] = outer(pos, inv_freq)
        let freqs = pos
            .unsqueeze_dim::<2>(1)
            .matmul(inv_freq.unsqueeze_dim::<2>(0));
        // duplicate to cover the whole head_dim: [f0..fh, f0..fh]
        let emb = Tensor::cat(vec![freqs.clone(), freqs], 1);

        Self {
            cos: emb.clone().cos(),
            sin: emb.sin(),
        }
    }

    /// x: [batch, heads, seq, head_dim]
    pub fn apply(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
        let [b, h, t, hd] = x.dims();
        let half = hd / 2;

        let cos = self.cos.clone().slice([0..t, 0..hd]).unsqueeze::<4>();
        let sin = self.sin.clone().slice([0..t, 0..hd]).unsqueeze::<4>();

        let x1 = x.clone().slice([0..b, 0..h, 0..t, 0..half]);
        let x2 = x.clone().slice([0..b, 0..h, 0..t, half..hd]);
        let rotated = Tensor::cat(vec![-x2, x1], 3);

        x * cos + rotated * sin
    }
}

#[derive(Module, Debug)]
pub struct Attention<B: Backend> {
    pub q: LoraLinear<B>,
    pub k: LoraLinear<B>,
    pub v: LoraLinear<B>,
    pub o: LoraLinear<B>,
    pub n_heads: usize,
    pub n_kv_heads: usize,
    pub head_dim: usize,
}

impl<B: Backend> Attention<B> {
    fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self {
        let hd = cfg.head_dim();
        let d = cfg.hidden_size;
        let kv = cfg.num_key_value_heads * hd;
        let r = |name: &str| {
            if lora.targets(name) {
                Some((lora.rank, lora.scaling()))
            } else {
                None
            }
        };
        Self {
            q: LoraLinear::new(d, d, r("q"), device),
            k: LoraLinear::new(d, kv, r("k"), device),
            v: LoraLinear::new(d, kv, r("v"), device),
            o: LoraLinear::new(d, d, r("o"), device),
            n_heads: cfg.num_attention_heads,
            n_kv_heads: cfg.num_key_value_heads,
            head_dim: hd,
        }
    }

    /// Repeats each KV head `groups` times, interleaved, like Hugging Face
    /// `repeat_kv`: [k0,k0,k1,k1,...].
    fn repeat_kv(&self, x: Tensor<B, 4>, groups: usize) -> Tensor<B, 4> {
        if groups == 1 {
            return x;
        }
        let [b, n_kv, t, hd] = x.dims();
        x.reshape([b, n_kv, 1, t, hd])
            .repeat_dim(2, groups)
            .reshape([b, n_kv * groups, t, hd])
    }

    pub fn forward(&self, x: Tensor<B, 3>, rope: &Rope<B>) -> Tensor<B, 3> {
        let [b, t, d] = x.dims();
        let hd = self.head_dim;
        let groups = self.n_heads / self.n_kv_heads;
        let device = x.device();

        let q = self
            .q
            .forward(x.clone())
            .reshape([b, t, self.n_heads, hd])
            .swap_dims(1, 2);
        let k = self
            .k
            .forward(x.clone())
            .reshape([b, t, self.n_kv_heads, hd])
            .swap_dims(1, 2);
        let v = self
            .v
            .forward(x)
            .reshape([b, t, self.n_kv_heads, hd])
            .swap_dims(1, 2);

        let q = rope.apply(q);
        let k = rope.apply(k);

        let k = self.repeat_kv(k, groups);
        let v = self.repeat_kv(v, groups);

        let scale = 1.0 / (hd as f64).sqrt();
        let scores = q.matmul(k.swap_dims(2, 3)) * scale;

        // Causal mask. Mind burn's semantics: `tril_mask` is `false` on the
        // triangle that is KEPT (lower + diagonal) and `true` elsewhere,
        // which is what gets filled with -inf.
        let mask = Tensor::<B, 2, Bool>::tril_mask([t, t], 0, &device).unsqueeze::<4>();
        let scores = scores.mask_fill(mask, f32::NEG_INFINITY);
        let attn = softmax(scores, 3);

        let out = attn.matmul(v).swap_dims(1, 2).reshape([b, t, d]);
        self.o.forward(out)
    }

    fn freeze_base(self) -> Self {
        Self {
            q: self.q.freeze_base(),
            k: self.k.freeze_base(),
            v: self.v.freeze_base(),
            o: self.o.freeze_base(),
            ..self
        }
    }

    fn trainable_ids(&self, out: &mut Vec<ParamId>) {
        for l in [&self.q, &self.k, &self.v, &self.o] {
            l.trainable_ids(out);
        }
    }
}

#[derive(Module, Debug)]
pub struct Mlp<B: Backend> {
    pub gate: LoraLinear<B>,
    pub up: LoraLinear<B>,
    pub down: LoraLinear<B>,
}

impl<B: Backend> Mlp<B> {
    fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self {
        let (d, i) = (cfg.hidden_size, cfg.intermediate_size);
        let r = |name: &str| {
            if lora.targets(name) {
                Some((lora.rank, lora.scaling()))
            } else {
                None
            }
        };
        Self {
            gate: LoraLinear::new(d, i, r("gate"), device),
            up: LoraLinear::new(d, i, r("up"), device),
            down: LoraLinear::new(i, d, r("down"), device),
        }
    }

    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        let gate = silu(self.gate.forward(x.clone()));
        self.down.forward(gate * self.up.forward(x))
    }

    fn freeze_base(self) -> Self {
        Self {
            gate: self.gate.freeze_base(),
            up: self.up.freeze_base(),
            down: self.down.freeze_base(),
        }
    }

    fn trainable_ids(&self, out: &mut Vec<ParamId>) {
        for l in [&self.gate, &self.up, &self.down] {
            l.trainable_ids(out);
        }
    }
}

#[derive(Module, Debug)]
pub struct Block<B: Backend> {
    pub attn_norm: RmsNorm<B>,
    pub attn: Attention<B>,
    pub ffn_norm: RmsNorm<B>,
    pub ffn: Mlp<B>,
}

impl<B: Backend> Block<B> {
    fn new(cfg: &HfConfig, lora: &LoraConfig, device: &B::Device) -> Self {
        let norm = || {
            RmsNormConfig::new(cfg.hidden_size)
                .with_epsilon(cfg.rms_norm_eps)
                .init(device)
        };
        Self {
            attn_norm: norm(),
            attn: Attention::new(cfg, lora, device),
            ffn_norm: norm(),
            ffn: Mlp::new(cfg, lora, device),
        }
    }

    pub fn forward(&self, x: Tensor<B, 3>, rope: &Rope<B>) -> Tensor<B, 3> {
        let h = x.clone() + self.attn.forward(self.attn_norm.forward(x), rope);
        h.clone() + self.ffn.forward(self.ffn_norm.forward(h))
    }

    fn freeze_base(self) -> Self {
        Self {
            attn: self.attn.freeze_base(),
            ffn: self.ffn.freeze_base(),
            ..self
        }
    }

    fn trainable_ids(&self, out: &mut Vec<ParamId>) {
        self.attn.trainable_ids(out);
        self.ffn.trainable_ids(out);
    }
}

#[derive(Module, Debug)]
pub struct Llama<B: Backend> {
    pub embed: Embedding<B>,
    pub layers: Vec<Block<B>>,
    pub norm: RmsNorm<B>,
    pub rope: Rope<B>,
}

impl<B: Backend> Llama<B> {
    pub fn new(cfg: &HfConfig, lora: &LoraConfig, max_seq: usize, device: &B::Device) -> Self {
        Self {
            embed: EmbeddingConfig::new(cfg.vocab_size, cfg.hidden_size).init(device),
            layers: (0..cfg.num_hidden_layers)
                .map(|_| Block::new(cfg, lora, device))
                .collect(),
            norm: RmsNormConfig::new(cfg.hidden_size)
                .with_epsilon(cfg.rms_norm_eps)
                .init(device),
            rope: Rope::new(cfg.head_dim(), max_seq, cfg.rope_theta, device),
        }
    }

    /// tokens: [batch, seq] -> logits [batch, seq, vocab]
    pub fn forward(&self, tokens: Tensor<B, 2, Int>) -> Tensor<B, 3> {
        let mut h = self.embed.forward(tokens);
        for layer in &self.layers {
            h = layer.forward(h, &self.rope);
        }
        let h = self.norm.forward(h);
        // tied embeddings: lm_head = embed.weight^T
        h.matmul(self.embed.weight.val().transpose().unsqueeze())
    }

    /// Freezes all base weights; only the LoRA A/B stay trainable.
    pub fn freeze_base(self) -> Self {
        let Self {
            embed,
            layers,
            norm,
            rope,
        } = self;
        Self {
            embed: Embedding {
                weight: embed.weight.set_require_grad(false),
            },
            layers: layers.into_iter().map(|l| l.freeze_base()).collect(),
            norm: RmsNorm {
                gamma: norm.gamma.set_require_grad(false),
                epsilon: norm.epsilon,
            },
            rope,
        }
    }

    /// L2 norm of all B matrices together. They start at zero, so this is a
    /// fingerprint to check that an adapter was trained and that it loads
    /// back identical.
    pub fn lora_norm(&self) -> f32 {
        let mut acc = 0.0f32;
        for l in &self.layers {
            let mut push = |lin: &LoraLinear<B>| {
                if let Some(ad) = &lin.lora {
                    let v: f32 = ad.b.val().powf_scalar(2.0).sum().into_scalar().elem();
                    acc += v;
                }
            };
            for lin in [&l.attn.q, &l.attn.k, &l.attn.v, &l.attn.o] {
                push(lin);
            }
            for lin in [&l.ffn.gate, &l.ffn.up, &l.ffn.down] {
                push(lin);
            }
        }
        acc.sqrt()
    }

    pub fn trainable_ids(&self) -> Vec<ParamId> {
        let mut out = Vec::new();
        for l in &self.layers {
            l.trainable_ids(&mut out);
        }
        out
    }
}