import torch
import torch.nn as nn
import torch.nn.functional as F


class LoRALinear(nn.Module):
    """Wraps a frozen nn.Linear with a trainable low-rank update."""

    def __init__(self, base_linear: nn.Linear, r=8, alpha=16):
        super().__init__()
        self.base = base_linear
        for p in self.base.parameters():
            p.requires_grad = False

        d_out, d_in = base_linear.weight.shape
        self.r = r
        self.scaling = alpha / r
        self.A = nn.Parameter(torch.randn(r, d_in) * 0.01)
        self.B = nn.Parameter(torch.zeros(d_out, r))

    def forward(self, x):
        base_out = self.base(x)
        lora_out = (x @ self.A.T) @ self.B.T
        return base_out + self.scaling * lora_out

    def merged_weight(self):
        """Fold the LoRA update into a single weight matrix, as
        would be done before deploying the adapted model."""
        return self.base.weight + self.scaling * (self.B @ self.A)


def apply_lora_to_model(model, r=8, alpha=16, target_attr="qkv"):
    """Replace every nn.Linear named `target_attr` in the model
    with a LoRA-wrapped version, freezing everything else."""
    for name, module in model.named_modules():
        for child_name, child in module.named_children():
            if child_name == target_attr and isinstance(child, nn.Linear):
                setattr(module, child_name, LoRALinear(child, r, alpha))
    return model


def count_trainable_params(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def count_total_params(model):
    return sum(p.numel() for p in model.parameters())
