"""Plain LoRA for the custom ProtoGPT (NOT HF/PEFT — this is a bare nn.Module).
Wraps the four Linear layers per block: attn.qkv, attn.proj, mlp.0, mlp.2.
Base weights frozen; only A,B train. Adapter saves/loads as just the A,B tensors.
"""
import math
import torch
import torch.nn as nn

TARGETS = ("attn.qkv", "attn.proj", "mlp.0", "mlp.2")


class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, rank=16, alpha=32, dropout=0.05):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad_(False)              # base frozen
        self.scaling = alpha / rank
        self.drop = nn.Dropout(dropout)
        self.A = nn.Parameter(torch.empty(rank, base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, rank))   # B=0 -> delta=0 at init
        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))

    def forward(self, x):
        return self.base(x) + (self.drop(x) @ self.A.t() @ self.B.t()) * self.scaling


def inject_lora(model, rank=16, alpha=32, dropout=0.05, targets=TARGETS):
    """Freeze everything, replace target Linears with LoRALinear. Returns #replaced."""
    for p in model.parameters():
        p.requires_grad_(False)
    names = [n for n, m in model.named_modules()
             if isinstance(m, nn.Linear) and any(n.endswith(t) for t in targets)]
    for n in names:
        parent = model.get_submodule(n.rsplit(".", 1)[0])
        child = n.rsplit(".", 1)[1]
        setattr(parent, child, LoRALinear(getattr(parent, child), rank, alpha, dropout))
    return len(names)


def lora_trainable(model):
    return [p for p in model.parameters() if p.requires_grad]


def lora_state_dict(model):
    return {n: p.detach().cpu() for n, p in model.named_parameters() if p.requires_grad}


def load_lora_state(model, sd):
    params = dict(model.named_parameters())
    for n, v in sd.items():
        params[n].data.copy_(v.to(params[n].device, params[n].dtype))