"""Which of a LoRA's weights land in this model, and which cannot.

ComfyUI converts the naming dialects (lora_up / lora_A / PEFT); it does not unwrap a file whose
keys sit under a training wrapper. v4 measured it: an H3 adapter keyed `base_model.model.dit.*`
matched 0 of 532 keys and did nothing, silently. Every known wrapper is tried and the one matching
the most weights wins. A pair that matches by name but not by shape is dropped: ComfyUI reshapes
any delta with the right element count, so a wrong-shaped one merges scrambled (all-NaN mid-render),
and a full-width LoRA on a curve-form H3 checkpoint churned ~25 GB failing per key.
"""

import re

import torch

import comfy.lora
import comfy.lora_convert

# (prefix, replacement). Replacing, not only stripping: ComfyUI's key map has `diffusion_model.<path>`
# and no bare form, so `base_model.model.dit.blocks.0...` must become `diffusion_model.blocks.0...`.
REKEY = (
    ("transformer.", ""), ("diffusion_model.", ""), ("model.diffusion_model.", ""),
    ("base_model.model.", ""), ("base_model.model.dit.", ""), ("base_model.model.diffusion_model.", ""),
    ("lora_model.", ""), ("net.", ""),
    ("base_model.model.dit.", "diffusion_model."), ("base_model.model.", "diffusion_model."),
    ("dit.", "diffusion_model."), ("model.dit.", "diffusion_model."), ("transformer.", "diffusion_model."),
)


_DIFFUSERS = re.compile(r"(^|\.)(transformer_blocks|refiner_blocks)\.\d+\.")
_WRAPPER = re.compile(r"^.*?(?=(?:token_refiner\.)?(?:transformer_blocks|refiner_blocks)\.\d+\.)")
_PART = re.compile(r"^(?P<mod>.+?)(?P<tail>\.(?:lora_(?:[AB]|down|up)\.weight|lora_(?:[AB]|down|up)|alpha))$")
_SUFFIX = (("attn.to_out.0", "attn.out_proj"), ("ff.net.0.proj", "mlp.fc1"), ("ff.net.2", "mlp.fc2"))
_QKV = ("to_q", "to_k", "to_v")


_BLOCKS = re.compile(r"(?:^|\.)(?:token_refiner\.)?(?:transformer_blocks|refiner_blocks)\.\d+\.")      # a Diffusers module under any training wrapper


def from_diffusers(lora, fused_rows=None):
    """-> (lora, extra key map). A LoRA written against the Diffusers layout of the H3 transformer, renamed to the one
    ComfyUI loads, or the LoRA unchanged with {} when it is not in that layout.

    Three things differ. The names (`transformer_blocks.N` / `refiner_blocks.N` are `blocks.N` / `token_refiner.blocks.N`;
    `to_out.0`, `ff.net.0.proj`, `ff.net.2` are `out_proj`, `mlp.fc1`, `mlp.fc2`). Q, K and V are three Linears there and
    ONE fused `qkv_proj` here, so each is aimed at its own rows of that weight (an offset, which costs no extra rank).
    And the gated MLP: Diffusers stores the up half before the gate half, ComfyUI the gate half first, so the rows of
    the `fc1` up-matrix swap halves. Without the swap the file loads, matches, and quietly does the wrong thing.

    A file may train only some of Q, K and V. The offset of each needs every part's height, so `fused_rows(base)` (the
    row count of the model's own fused weight) fills in the parts the file does not hold, as equal thirds; a file whose
    parts are not thirds of that is left unrouted as before."""
    if not any(_DIFFUSERS.search(k) for k in lora):
        return lora, {}
    out, extra, sizes = {}, {}, {}
    for k, v in lora.items():
        m = _PART.match(_WRAPPER.sub("", k, count=1))                      # whatever training wrapper sits in front of the block names
        if not m:
            out[k] = v
            continue
        mod, tail = m.group("mod"), m.group("tail")
        tail = tail.replace("lora_down", "lora_A").replace("lora_up", "lora_B")        # the older spelling of the same two matrices
        mod = mod.replace("transformer_blocks.", "blocks.").replace("token_refiner.refiner_blocks.", "token_refiner.blocks.")
        for old, new in _SUFFIX:
            if mod.endswith(old):
                mod = mod[: -len(old)] + new
                break
        if mod.endswith(".mlp.fc1") and tail.replace(".weight", "") == ".lora_B":
            half = v.shape[0] // 2
            v = torch.cat((v[half:], v[:half]), 0)
        part = next((q for q in _QKV if mod.endswith(f".attn.{q}")), None)
        if part is not None:
            base = mod[: -len(part)] + "qkv_proj"
            mod = f"{base}@{part}"
            if tail.endswith("lora_B.weight") or tail == ".lora_B":
                sizes[mod] = int(v.shape[0])
        out[mod + tail] = v
    # Offsets need every part's height: q starts at 0, k after q, v after k.
    for mod in list(sizes):
        base, part = mod.split("@")
        heights = [sizes.get(f"{base}@{q}") for q in _QKV]
        if None in heights:
            rows = fused_rows(base) if fused_rows else None
            if not rows or rows % 3 or any(h not in (None, rows // 3) for h in heights):
                continue
            heights = [rows // 3] * 3
        start = sum(heights[: _QKV.index(part)])
        extra[mod] = (f"diffusion_model.{base}.weight", (0, start, sizes[mod]))
    return out, extra


def _rekey(lora, prefix, replacement):
    if not any(k.startswith(prefix) for k in lora):
        return None
    return {(replacement + k[len(prefix):] if k.startswith(prefix) else k): v for k, v in lora.items()}


def _absent(sd, key):
    """True when a patch aims at a weight of the diffusion model that this model does not have (a block number past its
    depth, a renamed layer): ComfyUI skips such a key without a word, so it must not be counted as applied."""
    name = key[0] if isinstance(key, tuple) else key
    return str(name).startswith("diffusion_model.") and name not in sd


def _misfit(sd, key, adapter):
    """True when a plain LoRA pair cannot make the weight it names (both dimensions, not the count)."""
    if getattr(adapter, "name", "lora") != "lora":                 # LoKr, LoHa, ... keep other factors in those slots: not an up/down pair
        return False
    w = getattr(adapter, "weights", None)
    window = key[1] if isinstance(key, tuple) else None            # (dim, start, length): a LoRA aimed at some rows of a fused weight (Q, K, V of qkv_proj)
    target = sd.get(key[0] if isinstance(key, tuple) else key)
    if not w or len(w) < 2 or target is None or target.dim() < 2:
        return False
    if (len(w) > 3 and w[3] is not None) or (len(w) > 5 and w[5] is not None):   # locon mid / padded target: not modelled
        return False
    try:
        got = (int(w[0].flatten(start_dim=1).shape[0]), int(w[1].flatten(start_dim=1).shape[-1]))
    except Exception:  # noqa: BLE001 -- not a plain pair (LoKr, LoHa, diff): left to ComfyUI
        return False
    if window is not None:
        dim, start, length = window
        return got != (int(length), int(target.shape[1:].numel())) or start + length > int(target.shape[dim])
    return got != (int(target.shape[0]), int(target.shape[1:].numel()))


def match(model, clip, lora):
    """-> (patches, how, dropped, unplaced): the best keying's patches, which keying won, how many were dropped for shape,
    and (Diffusers layout only) how many of the file's transformer modules found no place in this model at all."""
    key_map = comfy.lora.model_lora_keys_unet(model.model, {})
    if clip is not None:
        key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
    extra = {}
    modules = 0
    sd = model.model.state_dict()
    if any(".attn.qkv_proj." in str(v) for v in key_map.values()):         # only the H3 transformer is renamed: other models use `transformer_blocks` as they are
        modules = len({k.split(".lora_")[0].split(".alpha")[0] for k in lora if (k.startswith("transformer.") or _BLOCKS.search(k)) and (".lora_" in k or k.endswith(".alpha"))})
        lora, extra = from_diffusers(lora, lambda base: (lambda w: int(w.shape[0]) if w is not None and w.dim() else None)(sd.get(f"diffusion_model.{base}.weight")))
    key_map.update(extra)
    converted = comfy.lora_convert.convert_lora(lora)
    best, how = comfy.lora.load_lora(converted, key_map, log_missing=False), "as-is"
    for prefix, replacement in REKEY:
        variant = _rekey(converted, prefix, replacement)
        if variant is not None:
            got = comfy.lora.load_lora(variant, key_map, log_missing=False)
            if len(got) > len(best):
                best, how = got, f"{prefix} -> {replacement or '(none)'}"
    bad = [k for k, a in best.items() if _absent(sd, k) or _misfit(sd, k, a)]
    for k in bad:
        best.pop(k)
    unplaced = max(0, modules - len(best) - len(bad)) if modules else 0      # modules ComfyUI never mapped: skipped without a word
    return best, how, len(bad), unplaced