"""
Normalized Low-Rank Adaptation (NoRA & NoRA-Init).
Reference: SphereLab, Microsoft Research, CUHK (arXiv:2608.31036, Aug 31, 2026).

Normalizes the down-projection matrix A along the rank dimension:
    delta_W = alpha * B * Norm(A)
Provides improved convergence, zero parameter overhead, and exact weight mergeability.
"""

from __future__ import annotations

import math
from typing import List, Optional

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


class NoRALinear(nn.Module):
    """
    Linear layer with Normalized Low-Rank Adaptation (NoRA).
    Applies rank-dimensional normalization to down-projection A:
        y = W_0 x + (alpha / r) * B * Norm_r(A) * x
    """

    def __init__(
        self,
        base_linear: nn.Linear,
        rank: int = 16,
        alpha: float = 32.0,
        dropout: float = 0.0,
        nora_mode: str = "dynamic",  # "dynamic" (re-normalizes per step) or "init" (NoRA-init)
    ):
        super().__init__()
        self.in_features = base_linear.in_features
        self.out_features = base_linear.out_features
        self.rank = rank
        self.alpha = alpha
        self.scaling = alpha / rank
        self.nora_mode = nora_mode
        self.merged = False

        # Base frozen weights
        self.weight = base_linear.weight
        self.bias = base_linear.bias
        self.weight.requires_grad = False
        if self.bias is not None:
            self.bias.requires_grad = False

        # NoRA adapter weights: A (r, in_features), B (out_features, r)
        self.nora_a = nn.Parameter(torch.empty(rank, self.in_features))
        self.nora_b = nn.Parameter(torch.zeros(self.out_features, rank))
        self.dropout = nn.Dropout(p=dropout) if dropout > 0.0 else nn.Identity()

        self.reset_parameters()

    def reset_parameters(self):
        # Kaiming uniform for down-projection A
        nn.init.kaiming_uniform_(self.nora_a, a=math.sqrt(5))
        # Zero-initialize B so adapter starts as exact identity
        nn.init.zeros_(self.nora_b)

        # Apply NoRA-init: normalize rows along the rank dimension
        with torch.no_grad():
            self.nora_a.copy_(self._normalize_a(self.nora_a))

    def _normalize_a(self, a_matrix: torch.Tensor) -> torch.Tensor:
        """Normalizes projection vectors along the rank dimension (dim 0)."""
        norm = torch.norm(a_matrix, p=2, dim=-1, keepdim=True) + 1e-8
        # Scale by sqrt(in_features) to preserve expected activation variance
        return (a_matrix / norm) * math.sqrt(self.in_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.merged:
            return F.linear(x, self.weight, self.bias)

        base_out = F.linear(x, self.weight, self.bias)

        # Normalize A dynamically if in dynamic mode
        if self.nora_mode == "dynamic" and self.training:
            a_effective = self._normalize_a(self.nora_a)
        else:
            a_effective = self.nora_a

        # Low-rank forward: (B, S, D_in) -> (B, S, r) -> (B, S, D_out)
        dropped_x = self.dropout(x)
        latent = F.linear(dropped_x, a_effective)
        adapter_out = F.linear(latent, self.nora_b) * self.scaling

        return base_out + adapter_out

    def merge_weights(self):
        """Exact merge of adapter weights into base weights."""
        if self.merged:
            return
        with torch.no_grad():
            a_effective = self._normalize_a(self.nora_a) if self.nora_mode == "dynamic" else self.nora_a
            delta_w = (self.nora_b @ a_effective) * self.scaling
            self.weight.add_(delta_w)
            self.merged = True

    def unmerge_weights(self):
        """Reverses the weight merge."""
        if not self.merged:
            return
        with torch.no_grad():
            a_effective = self._normalize_a(self.nora_a) if self.nora_mode == "dynamic" else self.nora_a
            delta_w = (self.nora_b @ a_effective) * self.scaling
            self.weight.sub_(delta_w)
            self.merged = False


def apply_nora(
    model: nn.Module,
    rank: int = 16,
    alpha: float = 32.0,
    target_names: List[str] = ["q_proj", "v_proj", "k_proj"],
    nora_mode: str = "dynamic",
) -> List[str]:
    """Applies NoRA adapters to designated linear layers in the model."""
    adapted_layers = []
    for name, module in model.named_modules():
        for child_name, child in module.named_children():
            if isinstance(child, nn.Linear) and any(t in child_name for t in target_names):
                nora_layer = NoRALinear(child, rank=rank, alpha=alpha, nora_mode=nora_mode)
                setattr(module, child_name, nora_layer)
                adapted_layers.append(f"{name}.{child_name}")
    return adapted_layers