"""Lossless ai-toolkit Qwen Image 2.1 fused-MLP LoRA conversion.

ai-toolkit's Qwen Image 2.1 transformer stores ``gate_up = [gate; up]``;
Diffusers stores ``gate_layer`` and ``proj`` separately. For a fused update
``B @ A``, duplicate A and split B's output rows. Rank and scale do not change.

Sources:
https://github.com/ostris/ai-toolkit/blob/main/extensions_built_in/diffusion_models/qwen_image_2/src/transformer.py
https://github.com/huggingface/diffusers/blob/85c9fa4caee33d1c5b473395784471d2d7957c1b/src/diffusers/models/transformers/transformer_qwenimage21.py

This deliberately rejects unknown tensors or unmatched modules instead of
silently loading only the attention part of a checkpoint.
"""

from __future__ import annotations

import math
from collections.abc import Mapping
from typing import Any


def convert_angles_state_dict(state_dict: Mapping[str, Any], transformer: Any) -> dict[str, Any]:
    """Convert A/B-only weights and validate every target against the model.

    The published v2 checkpoint uses PEFT A/B weights without alpha tensors.
    Its rank/alpha are both 64, hence the effective update is B @ A. An alpha
    tensor is rejected here so a differently scaled format cannot be mistaken
    for this one. No adapter tensors are discarded or rescaled.
    """
    pairs: dict[str, dict[str, Any]] = {}
    for key, tensor in state_dict.items():
        name = key
        for prefix in ("diffusion_model.", "transformer."):
            if name.startswith(prefix):
                name = name[len(prefix):]
                break
        suffix = next((s for s in (".lora_A.weight", ".lora_B.weight") if name.endswith(s)), None)
        if suffix is None or tensor.ndim != 2:
            raise ValueError(f"Unsupported LoRA tensor: {key}")
        module_name = name[:-len(suffix)]
        kind = suffix[6]
        pair = pairs.setdefault(module_name, {})
        if kind in pair:
            raise ValueError(f"Duplicate LoRA tensor: {key}")
        pair[kind] = tensor
    if not pairs:
        raise ValueError("The LoRA checkpoint has no adapter tensors.")

    converted: dict[str, Any] = {}

    def add_pair(module_name: str, a: Any, b: Any) -> None:
        try:
            module = transformer.get_submodule(module_name)
            # Also works when another adapter has already wrapped this Linear.
            base = module.get_base_layer() if hasattr(module, "get_base_layer") else module
            expected = tuple(base.weight.shape)
        except (AttributeError, KeyError) as exc:
            raise ValueError(f"LoRA targets a missing model module: {module_name}") from exc
        if len(expected) != 2 or tuple((b.shape[0], a.shape[1])) != expected or b.shape[1] != a.shape[0]:
            raise ValueError(f"LoRA shape mismatch for {module_name}: A={tuple(a.shape)}, B={tuple(b.shape)}, base={expected}")
        for kind, tensor in (("A", a), ("B", b)):
            key = f"transformer.{module_name}.lora_{kind}.weight"
            if key in converted:
                raise ValueError(f"Overlapping fused and split LoRA tensors: {key}")
            converted[key] = tensor

    for name, pair in pairs.items():
        if set(pair) != {"A", "B"}:
            raise ValueError(f"Incomplete A/B pair: {name}")
        a, b = pair["A"], pair["B"]
        if name.endswith(".img_mlp.gate_up"):
            if b.shape[0] % 2:
                raise ValueError(f"Fused gate/up output width must be even: {name}")
            halfway = b.shape[0] // 2
            prefix = name[:-len("gate_up")]
            add_pair(prefix + "gate_layer", a, b[:halfway].contiguous())
            add_pair(prefix + "proj", a, b[halfway:].contiguous())
        else:
            add_pair(name, a, b)
    return converted


def load_angles_adapter(pipe: Any, checkpoint_path: str, *, adapter_name: str = "angles", strength: float = 1.0) -> dict[str, Any]:
    """Load the complete adapter and verify that every A/B pair was installed.

    Use a fresh adapter name. Conversion and shape checks happen before model
    mutation. If installation validation fails, remove this adapter and raise.
    ``strength`` is the sole additional multiplier; rank is never halved when
    splitting gate/up and no sqrt-rank scaling is introduced.
    """
    from safetensors.torch import load_file

    if not math.isfinite(strength) or strength < 0:
        raise ValueError("Adapter strength must be finite and nonnegative.")
    if adapter_name in getattr(pipe.transformer, "peft_config", {}):
        raise ValueError(f"Adapter already exists: {adapter_name}")
    source = load_file(str(checkpoint_path), device="cpu")
    weights = convert_angles_state_dict(source, pipe.transformer)
    expected = {
        key.removeprefix("transformer.").removesuffix(".lora_A.weight"): tuple(tensor.shape)
        for key, tensor in weights.items() if key.endswith(".lora_A.weight")
    }
    fused_count = sum(key.endswith(".img_mlp.gate_up.lora_A.weight") for key in source)
    try:
        pipe.load_lora_weights(weights, adapter_name=adapter_name)
        for name, a_shape in expected.items():
            module = pipe.transformer.get_submodule(name)
            for kind in ("A", "B"):
                bank = getattr(module, f"lora_{kind}", {})
                wanted = weights[f"transformer.{name}.lora_{kind}.weight"]
                if adapter_name not in bank or tuple(bank[adapter_name].weight.shape) != tuple(wanted.shape):
                    raise RuntimeError(f"Adapter weight was not installed completely: {name}.lora_{kind}")
            # PEFT must use alpha/rank = 1 before the requested user strength.
            if module.r[adapter_name] != a_shape[0] or not math.isclose(float(module.scaling[adapter_name]), 1.0):
                raise RuntimeError(f"Unexpected rank or alpha scaling: {name}")
        pipe.set_adapters([adapter_name], adapter_weights=[strength])
    except Exception:
        if adapter_name in getattr(pipe.transformer, "peft_config", {}):
            pipe.delete_adapters(adapter_name)
        raise
    return {
        "input_tensors": len(source),
        "loaded_tensors": len(weights),
        "fused_mlp_blocks": fused_count,
        "loaded_modules": len(expected),
        "ranks": sorted({shape[0] for shape in expected.values()}),
        "strength": strength,
    }