"""
lora.py — LoRA adapter for KaizenLM Phase 5 online learning.

Architecture: rank-4 LoRA on Q and V projections of all 8 attention layers.
  Q slice: qkv_out[..., :D]       ← corrected by lora_q
  V slice: qkv_out[..., 2*D:]     ← corrected by lora_v

Adapter params per layer: 2 × (rank×D + D×rank) = 2 × (4×512 + 512×4) = 8,192
Total (8 layers): 65,536 params
Storage per adapter: 65,536 × 4 bytes ≈ 256 KB

Usage:
  adapter = LoRAAdapter(n_layers=8, d=512, rank=4, alpha=8)
  with torch.no_grad():
      out = model_with_lora_forward(base_model, x, adapter)
  loss.backward()  # gradients only on adapter params
  optimizer.step()
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class LoRALinear(nn.Module):
    """Additive low-rank correction: output += B(A(x)) * scale."""

    def __init__(self, in_features: int, out_features: int,
                 rank: int = 4, alpha: float = 8.0):
        super().__init__()
        self.rank  = rank
        self.scale = alpha / rank
        self.A = nn.Linear(in_features, rank, bias=False)
        self.B = nn.Linear(rank, out_features, bias=False)
        nn.init.kaiming_uniform_(self.A.weight, a=math.sqrt(5))
        nn.init.zeros_(self.B.weight)   # zero init → adapter starts as identity

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.B(self.A(x)) * self.scale


class LoRAAdapter(nn.Module):
    """
    Per-layer LoRA corrections for KaizenLM.
    Applies to Q and V slices of the fused QKV projection in each attention block.
    """

    def __init__(self, n_layers: int = 8, d: int = 512,
                 rank: int = 4, alpha: float = 8.0):
        super().__init__()
        self.d = d
        self.lora_q = nn.ModuleList(
            [LoRALinear(d, d, rank, alpha) for _ in range(n_layers)])
        self.lora_v = nn.ModuleList(
            [LoRALinear(d, d, rank, alpha) for _ in range(n_layers)])

    def correct_qkv(self, layer_idx: int, x: torch.Tensor,
                    qkv: torch.Tensor) -> torch.Tensor:
        """
        Add LoRA corrections to Q and V slices.
        qkv shape: [B, T, 3*D]
        Uses torch.cat (not in-place) to preserve autograd graph for backward.
        """
        D = self.d
        q_corrected = qkv[..., :D]   + self.lora_q[layer_idx](x)
        k_unchanged = qkv[..., D:2*D]
        v_corrected = qkv[..., 2*D:] + self.lora_v[layer_idx](x)
        return torch.cat([q_corrected, k_unchanged, v_corrected], dim=-1)

    def param_count(self) -> int:
        return sum(p.numel() for p in self.parameters())

    def size_bytes(self) -> int:
        return sum(p.numel() * p.element_size() for p in self.parameters())

    @classmethod
    def merged(cls, adapters: list, weights: list,
               n_layers: int = 8, d: int = 512,
               rank: int = 4, alpha: float = 8.0) -> 'LoRAAdapter':
        """
        Weighted average merge of multiple adapters.
        weights need not sum to 1 (normalised internally).
        """
        assert len(adapters) == len(weights) and len(adapters) > 0
        total = sum(weights)
        merged = cls(n_layers=n_layers, d=d, rank=rank, alpha=alpha)
        # zero-init merged state
        with torch.no_grad():
            for p in merged.parameters():
                p.zero_()
            norm_w = [w / total for w in weights]
            for adapter, w in zip(adapters, norm_w):
                for p_m, p_a in zip(merged.parameters(), adapter.parameters()):
                    p_m.add_(p_a * w)
        return merged


# ── KaizenLM with LoRA forward ────────────────────────────────────────────────
# Copied model definition with LoRA injection — base weights unchanged.

def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    h = x.shape[-1] // 2
    return torch.cat([-x[..., h:], x[..., :h]], dim=-1)


def _apply_rope(q, k, cos, sin):
    return q * cos + _rotate_half(q) * sin, k * cos + _rotate_half(k) * sin


def _precompute_rope(head_dim, max_seq, theta=10000.0):
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    t   = torch.arange(max_seq).float()
    f   = torch.outer(t, inv)
    return torch.cat([f, f], dim=-1).cos(), torch.cat([f, f], dim=-1).sin()


class RMSNorm(nn.Module):
    def __init__(self, d: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps    = eps

    def forward(self, x):
        n = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * n * self.weight).type_as(x)


class CausalSelfAttentionWithLoRA(nn.Module):
    def __init__(self, d: int, n_heads: int, drop: float = 0.0):
        super().__init__()
        self.nh = n_heads
        self.hd = d // n_heads
        self.drop = drop
        self.qkv      = nn.Linear(d, 3 * d, bias=False)
        self.out_proj = nn.Linear(d, d, bias=False)

    def forward(self, x, cos, sin, lora_adapter=None, layer_idx=None):
        B, T, C = x.shape
        qkv = self.qkv(x)
        if lora_adapter is not None:
            qkv = lora_adapter.correct_qkv(layer_idx, x, qkv)
        q, k, v = qkv.split(C, dim=-1)
        q = q.view(B, T, self.nh, self.hd)
        k = k.view(B, T, self.nh, self.hd)
        v = v.view(B, T, self.nh, self.hd)
        q, k = _apply_rope(q, k, cos[:T].unsqueeze(1), sin[:T].unsqueeze(1))
        q = q.transpose(1, 2); k = k.transpose(1, 2); v = v.transpose(1, 2)
        o = F.scaled_dot_product_attention(q, k, v, is_causal=True,
                dropout_p=self.drop if self.training else 0.0)
        return self.out_proj(o.transpose(1, 2).contiguous().view(B, T, C))


class FFN(nn.Module):
    def __init__(self, d: int, d_ff: int):
        super().__init__()
        self.w1 = nn.Linear(d, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d, bias=False)

    def forward(self, x):
        return self.w2(F.gelu(self.w1(x), approximate='tanh'))


class BlockWithLoRA(nn.Module):
    def __init__(self, d: int, n_heads: int, d_ff: int):
        super().__init__()
        self.ln1  = RMSNorm(d)
        self.attn = CausalSelfAttentionWithLoRA(d, n_heads)
        self.ln2  = RMSNorm(d)
        self.ff   = FFN(d, d_ff)

    def forward(self, x, cos, sin, lora_adapter=None, layer_idx=None):
        x = x + self.attn(self.ln1(x), cos, sin, lora_adapter, layer_idx)
        return x + self.ff(self.ln2(x))


class KaizenWithLoRA(nn.Module):
    """
    KaizenLM base (frozen) + optional LoRAAdapter (trainable).
    Base weights are never updated. LoRA adapter is swappable at inference time.
    """

    VOCAB_SIZE = 32768
    D_MODEL    = 512
    N_HEADS    = 8
    N_LAYERS   = 8
    D_FF       = 2048
    BLOCK_SIZE = 1024

    def __init__(self):
        super().__init__()
        d, h, l, ff, b = (self.D_MODEL, self.N_HEADS, self.N_LAYERS,
                          self.D_FF, self.BLOCK_SIZE)
        self.embed  = nn.Embedding(self.VOCAB_SIZE, d)
        self.blocks = nn.ModuleList(
            [BlockWithLoRA(d, h, ff) for _ in range(l)])
        self.ln_f   = RMSNorm(d)
        self.head   = nn.Linear(d, self.VOCAB_SIZE, bias=False)
        self.head.weight = self.embed.weight

        cos, sin = _precompute_rope(d // h, b * 2)
        self.register_buffer('rope_cos', cos)
        self.register_buffer('rope_sin', sin)

    def load_base(self, ckpt_path: str):
        """Load base weights, freeze all base params."""
        ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=True)
        sd   = ckpt.get('model_state', ckpt.get('model', ckpt))
        self.load_state_dict(sd, strict=True)
        for p in self.parameters():
            p.requires_grad_(False)

    def forward(self, idx: torch.Tensor, targets=None,
                adapter: 'LoRAAdapter' = None) -> torch.Tensor:
        x   = self.embed(idx)
        cos = self.rope_cos[:idx.shape[1]]
        sin = self.rope_sin[:idx.shape[1]]
        for i, blk in enumerate(self.blocks):
            x = blk(x, cos, sin, lora_adapter=adapter, layer_idx=i)
        logits = self.head(self.ln_f(x))
        if targets is None:
            return logits
        loss = F.cross_entropy(
            logits.view(-1, logits.size(-1)),
            targets.view(-1),
            ignore_index=-100,
        )
        return logits, loss

    def embed_task(self, idx: torch.Tensor,
                   adapter: 'LoRAAdapter' = None) -> torch.Tensor:
        """
        Returns mean-pooled last hidden state (before lm_head) as task embedding.
        Shape: [D_MODEL]
        """
        cos = self.rope_cos[:idx.shape[1]]
        sin = self.rope_sin[:idx.shape[1]]
        with torch.no_grad():
            x = self.embed(idx)
            for i, blk in enumerate(self.blocks):
                x = blk(x, cos, sin, lora_adapter=adapter, layer_idx=i)
            x = self.ln_f(x)
        return x[0].mean(0)   # [T, D] → [D]