use burn::module::{Module, Param, ParamId};
use burn::nn::{Linear, LinearConfig};
use burn::prelude::*;
use burn::tensor::Distribution;

/// Low-rank adapter: delta_W = (A @ B) * scaling.
/// A starts as Gaussian noise and B as zeros, so at the start of training
/// delta_W = 0 and the model is identical to the base one.
#[derive(Module, Debug)]
pub struct LoraAdapter<B: Backend> {
    /// [d_in, rank]
    pub a: Param<Tensor<B, 2>>,
    /// [rank, d_out]
    pub b: Param<Tensor<B, 2>>,
    pub scaling: f64,
}

impl<B: Backend> LoraAdapter<B> {
    pub fn new(d_in: usize, d_out: usize, rank: usize, scaling: f64, device: &B::Device) -> Self {
        let std = 1.0 / (rank as f64).sqrt();
        let a = Tensor::random([d_in, rank], Distribution::Normal(0.0, std), device);
        let b = Tensor::zeros([rank, d_out], device);
        Self {
            a: Param::from_tensor(a),
            b: Param::from_tensor(b),
            scaling,
        }
    }
}

/// Frozen Linear + optional LoRA adapter.
#[derive(Module, Debug)]
pub struct LoraLinear<B: Backend> {
    pub base: Linear<B>,
    pub lora: Option<LoraAdapter<B>>,
}

impl<B: Backend> LoraLinear<B> {
    pub fn new(d_in: usize, d_out: usize, rank: Option<(usize, f64)>, device: &B::Device) -> Self {
        let base = LinearConfig::new(d_in, d_out).with_bias(false).init(device);
        let lora = rank.map(|(r, scaling)| LoraAdapter::new(d_in, d_out, r, scaling, device));
        Self { base, lora }
    }

    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
        let out = self.base.forward(x.clone());
        match &self.lora {
            None => out,
            Some(lora) => {
                let delta = x
                    .matmul(lora.a.val().unsqueeze())
                    .matmul(lora.b.val().unsqueeze());
                out + delta * lora.scaling
            }
        }
    }

    /// Freezes the base weight: it stops accumulating gradients in the autodiff graph.
    pub fn freeze_base(self) -> Self {
        let Self { base, lora } = self;
        let base = Linear {
            weight: base.weight.set_require_grad(false),
            bias: base.bias.map(|b| b.set_require_grad(false)),
        };
        Self { base, lora }
    }

    /// IDs of the trainable parameters (only A and B).
    pub fn trainable_ids(&self, out: &mut Vec<ParamId>) {
        if let Some(lora) = &self.lora {
            out.push(lora.a.id);
            out.push(lora.b.id);
        }
    }
}