"""LoRA adapter for a frozen nn.Linear."""
from __future__ import annotations

import math

import torch
import torch.nn as nn


class LoRALinear(nn.Module):
    def __init__(
        self,
        base_linear: nn.Linear,
        rank: int = 8,
        alpha: float | None = None,
        dropout: float = 0.0,
    ):
        super().__init__()
        if rank <= 0:
            raise ValueError(f"LoRA rank must be positive; got {rank}")
        if alpha is None:
            alpha = float(rank)
        self.rank = rank
        self.alpha = float(alpha)
        self.scale = self.alpha / rank
        self.base = base_linear

        in_f = base_linear.in_features
        out_f = base_linear.out_features
        # Match base's device/dtype so post-construction LoRA wrapping
        # (e.g. apply_lora_and_freeze after the model is on GPU) doesn't leave
        # the adapter params on CPU.
        device = base_linear.weight.device
        dtype  = base_linear.weight.dtype
        self.lora_A = nn.Parameter(torch.empty(rank, in_f, device=device, dtype=dtype))
        self.lora_B = nn.Parameter(torch.zeros(out_f, rank, device=device, dtype=dtype))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

        self.lora_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        # When False the adapter is skipped and forward == base(x). Lets callers
        # evaluate the pristine frozen base (e.g. a self-distillation teacher)
        # without unwrapping. Defaults True — backward-compatible.
        self.enabled = True

    @property
    def in_features(self) -> int:
        return self.base.in_features

    @property
    def out_features(self) -> int:
        return self.base.out_features

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.base(x)
        if not self.enabled:
            return out
        update = self.lora_dropout(x) @ self.lora_A.t() @ self.lora_B.t()
        return out + update * self.scale