"""Minimal LoRA for the Stage-2 fake-score critic (no external dependency).

``inject_lora`` replaces selected ``nn.Linear`` layers inside ``model.blocks``
with ``LoRALinear`` (frozen base + trainable low-rank A/B, B zero-initialised so
the wrapped model starts identical to the base).  Everything outside the blocks
(patch/time/text embeddings, output head) stays frozen.
"""

from __future__ import annotations

import math
from typing import Iterable

import torch
from torch import nn


TARGET_GROUPS = {
    "self_attn": ("self_attn.q", "self_attn.k", "self_attn.v", "self_attn.o"),
    "cross_attn": ("cross_attn.q", "cross_attn.k", "cross_attn.v", "cross_attn.o"),
    "ffn": ("ffn.0", "ffn.2"),
}


class LoRALinear(nn.Module):
    def __init__(
        self, base: nn.Linear, rank: int, alpha: float, dropout: float = 0.0
    ) -> None:
        super().__init__()
        if rank <= 0:
            raise ValueError(f"LoRA rank must be positive, got {rank}")
        self.base = base
        self.base.requires_grad_(False)
        self.rank = int(rank)
        self.alpha = float(alpha)
        self.scale = self.alpha / self.rank
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
        # LoRA parameters live in fp32 regardless of the base dtype; autocast
        # handles the matmul dtype at run time.
        self.lora_A = nn.Parameter(torch.empty(self.rank, base.in_features, device=base.weight.device))
        self.lora_B = nn.Parameter(torch.zeros(base.out_features, self.rank, device=base.weight.device))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self.base(x)
        update = nn.functional.linear(
            nn.functional.linear(self.dropout(x), self.lora_A.to(x.dtype)),
            self.lora_B.to(x.dtype),
        )
        return out + self.scale * update

    def extra_repr(self) -> str:
        return f"rank={self.rank}, alpha={self.alpha}, base={self.base.in_features}->{self.base.out_features}"


def resolve_targets(spec: str | Iterable[str]) -> tuple[str, ...]:
    """'self_attn,ffn' -> module-name suffixes; individual suffixes also accepted."""
    names = spec.split(",") if isinstance(spec, str) else list(spec)
    out: list[str] = []
    for name in (n.strip() for n in names if n.strip()):
        if name == "all":
            for group in TARGET_GROUPS.values():
                out.extend(group)
        elif name in TARGET_GROUPS:
            out.extend(TARGET_GROUPS[name])
        else:
            out.append(name)
    if not out:
        raise ValueError("No LoRA target modules selected")
    return tuple(dict.fromkeys(out))


def inject_lora(
    model: nn.Module,
    *,
    rank: int,
    alpha: float | None = None,
    targets: str | Iterable[str] = "all",
    dropout: float = 0.0,
) -> list[nn.Parameter]:
    """Freeze ``model`` and wrap the target Linears of every ``model.blocks[i]``.

    Returns the trainable LoRA parameters.
    """
    if not hasattr(model, "blocks"):
        raise ValueError("inject_lora expects a Wan model with a .blocks list")
    alpha = float(rank if alpha is None else alpha)
    suffixes = resolve_targets(targets)
    model.requires_grad_(False)
    trainable: list[nn.Parameter] = []
    replaced = 0
    for block in model.blocks:
        for name in suffixes:
            parent_name, _, attr = name.rpartition(".")
            parent = block.get_submodule(parent_name) if parent_name else block
            child = getattr(parent, attr) if not attr.isdigit() else parent[int(attr)]
            if not isinstance(child, nn.Linear):
                raise TypeError(f"{name} is {type(child).__name__}, not nn.Linear")
            wrapped = LoRALinear(child, rank, alpha, dropout)
            if attr.isdigit():
                parent[int(attr)] = wrapped
            else:
                setattr(parent, attr, wrapped)
            trainable.extend([wrapped.lora_A, wrapped.lora_B])
            replaced += 1
    if replaced == 0:
        raise RuntimeError("inject_lora replaced no layers")
    return trainable


def lora_state_dict(model: nn.Module) -> dict[str, torch.Tensor]:
    return {
        key: value.detach().cpu().contiguous()
        for key, value in model.state_dict().items()
        if ".lora_A" in key or ".lora_B" in key
    }


def load_lora_state_dict(model: nn.Module, state: dict[str, torch.Tensor]) -> None:
    expected = set(lora_state_dict(model))
    if set(state) != expected:
        missing = sorted(expected - set(state))
        unexpected = sorted(set(state) - expected)
        raise RuntimeError(f"LoRA state mismatch: missing={missing[:3]}..., unexpected={unexpected[:3]}...")
    incompatible = model.load_state_dict(state, strict=False)
    if incompatible.unexpected_keys:
        raise RuntimeError(f"Unexpected LoRA keys: {incompatible.unexpected_keys[:3]}")