"""Small transparent LoRA implementation for the reference model's linear layers.

For external production models use PEFT; this is an inspectable training experiment.
"""
import math
import torch
from torch import nn


class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, rank=4, alpha=8):
        super().__init__()
        if rank < 1 or alpha <= 0:
            raise ValueError("Invalid LoRA hyperparameters")
        self.base, self.scale = base, alpha/rank
        self.base.requires_grad_(False)
        self.a = nn.Parameter(torch.empty(rank, base.in_features, device=base.weight.device, dtype=base.weight.dtype))
        self.b = nn.Parameter(torch.zeros(base.out_features, rank, device=base.weight.device, dtype=base.weight.dtype))
        nn.init.kaiming_uniform_(self.a, a=math.sqrt(5))

    def forward(self, x):
        return self.base(x) + (x @ self.a.T @ self.b.T)*self.scale

    def merged(self):
        base = nn.Linear(self.base.in_features, self.base.out_features, bias=self.base.bias is not None,
                         device=self.base.weight.device, dtype=self.base.weight.dtype)
        with torch.no_grad():
            base.weight.copy_(self.base.weight + (self.b @ self.a)*self.scale)
            if base.bias is not None:
                base.bias.copy_(self.base.bias)
        return base


def inject_lora(model, targets=("q", "v"), rank=4, alpha=8):
    model.requires_grad_(False)
    replaced = []
    for name, module in list(model.named_modules()):
        if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in targets:
            parent_name, _, leaf = name.rpartition(".")
            parent = model.get_submodule(parent_name) if parent_name else model
            setattr(parent, leaf, LoRALinear(module, rank, alpha))
            replaced.append(name)
    if not replaced:
        raise ValueError("No target linear modules found")
    return replaced


def merge_lora(model):
    for name, module in list(model.named_modules()):
        if isinstance(module, LoRALinear):
            parent_name, _, leaf = name.rpartition(".")
            parent = model.get_submodule(parent_name) if parent_name else model
            setattr(parent, leaf, module.merged())
    return model