# Copyright (c) 2026. LoRA adapters for MagiHuman / DAVINCI DiT.
"""Low-rank adapters for DiT linear layers (BaseLinear, NativeMoELinear, nn.Linear).

Scope: the middle DiT only (adapter embedders + TransformerBlock + final heads).
Video VAE / audio VAE / text encoder stay outside training (precomputed latents).
"""

from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

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

from .dit_module import BaseLinear, NativeMoELinear


# Module-name suffixes / exact leaf names we wrap when targets="all".
ATTN_MLP_SUFFIXES: Tuple[str, ...] = (
    "linear_qkv",
    "linear_proj",
    "up_gate_proj",
    "down_proj",
)
PLAIN_LINEAR_NAMES: Tuple[str, ...] = (
    "video_embedder",
    "text_embedder",
    "audio_embedder",
    "final_linear_video",
    "final_linear_audio",
)


@dataclass
class LoraConfig:
    r: int = 16
    alpha: float = 16.0
    dropout: float = 0.0
    targets: str = "all"  # all | attn | attn_mlp | embed_final
    """Which DiT linears to wrap. See TRAIN.md in v8/LoRA."""

    def scaling(self) -> float:
        return float(self.alpha) / float(max(self.r, 1))

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


def _match_target(name: str, targets: str) -> bool:
    leaf = name.rsplit(".", 1)[-1]
    t = (targets or "all").strip().lower()
    if t == "all":
        return leaf in ATTN_MLP_SUFFIXES or leaf in PLAIN_LINEAR_NAMES
    if t == "attn":
        return leaf in ("linear_qkv", "linear_proj")
    if t == "attn_mlp":
        return leaf in ATTN_MLP_SUFFIXES
    if t in ("embed_final", "heads"):
        return leaf in PLAIN_LINEAR_NAMES
    raise ValueError(f"Unknown lora targets={targets!r}; use all|attn|attn_mlp|embed_final")


def _init_lora_ab(
    in_features: int, out_features: int, r: int, device: torch.device, dtype: torch.dtype
) -> Tuple[nn.Parameter, nn.Parameter]:
    # A: (r, in), B: (out, r); B=0 so delta starts at 0.
    a = nn.Parameter(torch.empty(r, in_features, device=device, dtype=dtype))
    b = nn.Parameter(torch.zeros(out_features, r, device=device, dtype=dtype))
    nn.init.kaiming_uniform_(a, a=5**0.5)
    return a, b


class LoRAPlainLinear(nn.Module):
    """LoRA around a standard ``nn.Linear`` (adapter embedders / final heads)."""

    def __init__(self, base: nn.Linear, *, r: int, alpha: float, dropout: float = 0.0):
        super().__init__()
        if not isinstance(base, nn.Linear):
            raise TypeError(f"LoRAPlainLinear expects nn.Linear, got {type(base)}")
        self.base = base
        self.r = int(r)
        self.alpha = float(alpha)
        self.scaling = float(alpha) / float(max(self.r, 1))
        self.dropout_p = float(dropout)
        for p in self.base.parameters():
            p.requires_grad = False
        dtype = self.base.weight.dtype
        device = self.base.weight.device
        self.lora_A, self.lora_B = _init_lora_ab(
            self.base.in_features, self.base.out_features, self.r, device, dtype
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = self.base(x)
        x_d = F.dropout(x, p=self.dropout_p, training=self.training) if self.dropout_p > 0 else x
        # Cast LoRA matmul to input dtype for mixed precision.
        a = self.lora_A.to(dtype=x_d.dtype)
        b = self.lora_B.to(dtype=x_d.dtype)
        delta = (x_d @ a.t() @ b.t()) * self.scaling
        return y + delta.to(dtype=y.dtype)


class LoRABaseLinear(nn.Module):
    """LoRA around ``BaseLinear`` / ``NativeMoELinear`` (DiT attn + MLP).

    For ``num_experts==1``: one (A, B) pair.
    For ``NativeMoELinear`` (mm_layers): one (A, B) per modality expert, applied
    after ``modality_dispatcher.dispatch``, matching the base forward split.
    """

    def __init__(self, base: BaseLinear, *, r: int, alpha: float, dropout: float = 0.0):
        super().__init__()
        if not isinstance(base, BaseLinear):
            raise TypeError(f"LoRABaseLinear expects BaseLinear, got {type(base)}")
        self.base = base
        self.r = int(r)
        self.alpha = float(alpha)
        self.scaling = float(alpha) / float(max(self.r, 1))
        self.dropout_p = float(dropout)
        self.num_experts = int(base.num_experts)
        for p in self.base.parameters():
            p.requires_grad = False

        dtype = self.base.weight.dtype
        device = self.base.weight.device
        in_f = int(base.in_features)
        out_f = int(base.out_features)
        if self.num_experts <= 1:
            self.lora_A, self.lora_B = _init_lora_ab(in_f, out_f, self.r, device, dtype)
            self.lora_A_list = None
            self.lora_B_list = None
        else:
            self.lora_A = None
            self.lora_B = None
            self.lora_A_list = nn.ParameterList()
            self.lora_B_list = nn.ParameterList()
            for _ in range(self.num_experts):
                a, b = _init_lora_ab(in_f, out_f, self.r, device, dtype)
                self.lora_A_list.append(a)
                self.lora_B_list.append(b)

    def forward(
        self,
        input: torch.Tensor,
        output_dtype: Optional[torch.dtype] = None,
        modality_dispatcher: Optional[Any] = None,
    ) -> torch.Tensor:
        y = self.base(input, output_dtype=output_dtype, modality_dispatcher=modality_dispatcher)
        if self.num_experts <= 1:
            x_d = F.dropout(input, p=self.dropout_p, training=self.training) if self.dropout_p > 0 else input
            a = self.lora_A.to(dtype=x_d.dtype)
            b = self.lora_B.to(dtype=x_d.dtype)
            delta = (x_d @ a.t() @ b.t()) * self.scaling
            return y + delta.to(dtype=y.dtype)

        if modality_dispatcher is None:
            raise RuntimeError("NativeMoELinear LoRA requires modality_dispatcher")
        input_list = modality_dispatcher.dispatch(input)
        deltas: List[torch.Tensor] = []
        for i in range(self.num_experts):
            x_i = input_list[i]
            if x_i.numel() == 0:
                deltas.append(x_i.new_zeros(x_i.shape[0], self.base.out_features))
                continue
            x_d = F.dropout(x_i, p=self.dropout_p, training=self.training) if self.dropout_p > 0 else x_i
            a = self.lora_A_list[i].to(dtype=x_d.dtype)
            b = self.lora_B_list[i].to(dtype=x_d.dtype)
            deltas.append(((x_d @ a.t() @ b.t()) * self.scaling).to(dtype=y.dtype))
        delta = modality_dispatcher.undispatch(*deltas)
        return y + delta


def _wrap_module(module: nn.Module, cfg: LoraConfig) -> Optional[nn.Module]:
    if isinstance(module, (LoRAPlainLinear, LoRABaseLinear)):
        return None  # already wrapped
    if isinstance(module, BaseLinear):
        return LoRABaseLinear(module, r=cfg.r, alpha=cfg.alpha, dropout=cfg.dropout)
    if isinstance(module, nn.Linear):
        return LoRAPlainLinear(module, r=cfg.r, alpha=cfg.alpha, dropout=cfg.dropout)
    return None


def inject_lora(model: nn.Module, cfg: LoraConfig) -> List[str]:
    """Replace matching linear modules in-place. Returns list of wrapped paths."""
    if int(cfg.r) <= 0:
        raise ValueError(f"lora r must be > 0, got {cfg.r}")
    to_wrap: List[Tuple[str, nn.Module]] = []
    for name, module in model.named_modules():
        if not name:
            continue
        if not _match_target(name, cfg.targets):
            continue
        if isinstance(module, (LoRAPlainLinear, LoRABaseLinear)):
            continue
        if isinstance(module, (BaseLinear, nn.Linear)):
            to_wrap.append((name, module))

    wrapped: List[str] = []
    for name, module in to_wrap:
        new_mod = _wrap_module(module, cfg)
        if new_mod is None:
            continue
        parent = model
        parts = name.split(".")
        for p in parts[:-1]:
            parent = getattr(parent, p)
        setattr(parent, parts[-1], new_mod)
        wrapped.append(name)
    if not wrapped:
        raise RuntimeError(f"inject_lora: no modules matched targets={cfg.targets!r}")
    return wrapped


def freeze_non_lora_parameters(model: nn.Module) -> Tuple[int, int]:
    """Freeze everything except LoRA A/B. Returns (trainable_numel, frozen_numel)."""
    n_train, n_freeze = 0, 0
    for name, p in model.named_parameters():
        is_lora = ("lora_A" in name) or ("lora_B" in name) or ("lora_A_list" in name) or ("lora_B_list" in name)
        if is_lora:
            p.requires_grad = True
            n_train += p.numel()
        else:
            p.requires_grad = False
            n_freeze += p.numel()
    return n_train, n_freeze


def lora_state_dict(model: nn.Module) -> Dict[str, torch.Tensor]:
    out: Dict[str, torch.Tensor] = {}
    for name, p in model.named_parameters():
        if ("lora_A" in name) or ("lora_B" in name):
            out[name] = p.detach().cpu()
    return out


def load_lora_state_dict(model: nn.Module, state: Dict[str, torch.Tensor], strict: bool = True) -> None:
    missing, unexpected = [], []
    own = {n: p for n, p in model.named_parameters() if ("lora_A" in n) or ("lora_B" in n)}
    for k, v in state.items():
        if k not in own:
            unexpected.append(k)
            continue
        if tuple(own[k].shape) != tuple(v.shape):
            raise RuntimeError(f"LoRA shape mismatch for {k}: ckpt {tuple(v.shape)} vs model {tuple(own[k].shape)}")
        own[k].data.copy_(v.to(device=own[k].device, dtype=own[k].dtype))
    for k in own:
        if k not in state:
            missing.append(k)
    if strict and (missing or unexpected):
        raise RuntimeError(f"load_lora_state_dict strict fail missing={missing[:8]} unexpected={unexpected[:8]}")


def mark_lora_text_param_ids(model: nn.Module) -> set:
    """Parameter ids under adapter.text_embedder LoRA (for text_lr_mult group)."""
    ids: set = set()
    for name, p in model.named_parameters():
        if "adapter.text_embedder" in name and (("lora_A" in name) or ("lora_B" in name)):
            ids.add(id(p))
    return ids


def summarize_lora(model: nn.Module, wrapped: Sequence[str]) -> str:
    n_train, n_freeze = 0, 0
    for name, p in model.named_parameters():
        if p.requires_grad:
            n_train += p.numel()
        else:
            n_freeze += p.numel()
    return (
        f"lora_wrapped={len(wrapped)} trainable={n_train / 1e6:.3f}M "
        f"frozen={n_freeze / 1e6:.1f}M sample={list(wrapped)[:6]}"
    )