"""Corza LoRA Loader (Clean) — LoRA loader with optional spectral cleanup.

Why: LoRAs (especially stacked, and especially on distilled/turbo models like
Krea-2-Turbo) often add blocky / crunchy artifacts with fuzzy edges. Two usual
culprits inside the LoRA file itself:
  1. noise — the low-energy tail of each layer's update is mostly training
     noise, and few-step turbo sampling amplifies it instead of averaging it out;
  2. hot layers — a handful of layers carry a far stronger update than the rest
     and push activations off the distilled model's narrow manifold.

This node is a drop-in LoRA loader (identical to the stock one with cleaning
off) that can rewrite the LoRA's low-rank factors on the fly:
  * keep_energy — per layer, SVD the update and keep only the strongest
    components adding up to this % of energy (same math as kohya's resize_lora
    dynamic rank; 90-98 is a sensible range).
  * max_rank    — hard cap on the per-layer rank after the energy cut.
  * tame_layers — compress outlier layers (Frobenius norm above the 90th
    percentile of this LoRA) back toward that percentile. 0 = off, 1 = fully
    clamped to the percentile.

The cleanup happens on the LoRA state dict *before* ComfyUI's own
load_lora_for_models call, so it is architecture-agnostic: anything ComfyUI can
apply a LoRA to (Krea 2, Flux, SDXL, ...) works unchanged. Exotic entries
(DoRA, tucker/conv mid matrices, reshape metadata, unknown keys) pass through
untouched. For cross-LoRA conflict resolution (TIES/DARE) see
ethanfel/ComfyUI-LoRA-Optimizer — this node deliberately stays complementary:
it cleans one LoRA well instead of merging many.
"""

import logging
import uuid

import torch

import comfy.sd
import comfy.utils
import folder_paths
from comfy_api.latest import io


# (up_suffix, down_suffix) naming schemes we know how to rebuild. Matches the
# common cases in comfy/weight_adapter/lora.py; anything else passes through.
_PAIR_SUFFIXES = [
    (".lora_up.weight", ".lora_down.weight"),   # kohya
    (".lora_B.weight", ".lora_A.weight"),       # peft / diffusers2
    ("_lora.up.weight", "_lora.down.weight"),   # diffusers
    (".lora.up.weight", ".lora.down.weight"),   # diffusers3
]


def _find_pairs(sd):
    """Yield (up_key, down_key, alpha_key_or_None, prefix) for cleanable pairs."""
    pairs = []
    for key in sd.keys():
        for up_sfx, down_sfx in _PAIR_SUFFIXES:
            if not key.endswith(up_sfx):
                continue
            prefix = key[: -len(up_sfx)]
            down_key = prefix + down_sfx
            if down_key not in sd:
                break
            # Skip variants whose patch math we'd break by refactoring:
            # tucker/conv mid matrix, DoRA magnitude, explicit reshape.
            if (prefix + ".lora_mid.weight") in sd:
                break
            if (prefix + ".dora_scale") in sd:
                break
            if (prefix + ".reshape_weight") in sd:
                break
            up, down = sd[key], sd[down_key]
            if up.ndim < 2 or down.ndim < 2:
                break
            # Conv LoRAs: up must be (out, r, 1, 1...) so a rank change is a
            # plain reshape; down (r, in, kh, kw) flattens per component.
            if any(d != 1 for d in up.shape[2:]):
                break
            alpha_key = prefix + ".alpha"
            pairs.append((key, down_key, alpha_key if alpha_key in sd else None, prefix))
            break
    return pairs


@torch.no_grad()
def _factorize(up, down, alpha):
    """Exact SVD of the low-rank delta via QR — cheap even for huge layers.

    delta = up @ down * (alpha / rank)  is rank <= r, so QR-reduce both factors
    and SVD only the tiny r x r core instead of the full out x in matrix.
    """
    rank = down.shape[0]
    scale = (float(alpha) / rank) if alpha is not None else 1.0
    u2 = up.reshape(up.shape[0], -1).float()      # out x r
    d2 = down.reshape(rank, -1).float()           # r x in
    q1, r1 = torch.linalg.qr(u2)                  # out x m1, m1 x r
    q2, r2 = torch.linalg.qr(d2.T)                # in x m2,  m2 x r
    core = (r1 @ r2.T) * scale                    # m1 x m2 (tiny)
    um, s, vmh = torch.linalg.svd(core)
    return {"q1": q1, "q2": q2, "um": um, "s": s, "vmh": vmh}


def _pick_rank(s, keep_energy, max_rank):
    energy = s * s
    total = float(energy.sum())
    if total <= 0.0:
        return 1
    cum = torch.cumsum(energy, dim=0) / total
    k = int(torch.searchsorted(cum, keep_energy / 100.0).item()) + 1
    k = min(k, s.shape[0])
    if max_rank > 0:
        k = min(k, max_rank)
    return max(k, 1)


@torch.no_grad()
def _rebuild(f, k, gain, up, down):
    """Refactor the truncated SVD back into up/down with the scale baked in.

    Splitting sqrt(S) across both factors keeps their magnitudes balanced.
    Caller rewrites .alpha to the new rank, so the loader's alpha/rank scale
    becomes exactly 1 and the baked-in scale is the only one applied.
    """
    s = f["s"][:k] * gain
    sq = torch.sqrt(s)
    new_up = (f["q1"] @ f["um"][:, :k]) * sq.unsqueeze(0)        # out x k
    new_down = sq.unsqueeze(1) * (f["vmh"][:k] @ f["q2"].T)      # k x in
    out_dim = up.shape[0]
    new_up = new_up.reshape(out_dim, k, *up.shape[2:])
    new_down = new_down.reshape(k, *down.shape[1:])
    return new_up.to(up.dtype).contiguous(), new_down.to(down.dtype).contiguous()


def _is_gate(prefix):
    """True for multiplicative gate projections (gated attention output gate,
    SwiGLU MLP gate, etc.) — e.g. Krea 2's blocks.N.attn.gate / blocks.N.mlp.gate,
    or generic gate_proj / to_gate names."""
    return any("gate" in seg.lower() for seg in prefix.split("."))


@torch.no_grad()
def _clean_lora(sd, keep_energy, max_rank, tame_layers, star_rescale, gate_strength, lora_name):
    pairs = _find_pairs(sd)
    if not pairs:
        logging.info(f"[Corza LoRA] {lora_name}: no cleanable up/down pairs found, loading as-is")
        return sd

    # Pass 1: factorize every pair and pick its kept rank.
    facts = []
    for up_key, down_key, alpha_key, prefix in pairs:
        up, down = sd[up_key], sd[down_key]
        alpha = sd[alpha_key].item() if alpha_key is not None else None
        f = _factorize(up, down, alpha)
        k = _pick_rank(f["s"], keep_energy, max_rank)
        # STAR rescale (arXiv:2502.10339): after truncating, boost the kept
        # components so the layer's nuclear norm matches the original — trimming
        # removes the conflict-prone tail without weakening the LoRA's effect.
        gain = 1.0
        if star_rescale and k < f["s"].shape[0]:
            nuc_all = float(f["s"].sum())
            nuc_kept = float(f["s"][:k].sum())
            if nuc_kept > 0.0:
                gain = nuc_all / nuc_kept
        # Gate scaling: gates are multiplicative sigmoid controls, so a LoRA delta
        # there has outsized nonlinear effect (and stacked LoRAs' gate edits compound
        # rather than average) — a common artifact / "loras fighting" source that
        # norm-based taming misses. Scale their update down by gate_strength.
        is_gate = _is_gate(prefix)
        if is_gate:
            gain *= gate_strength
        norm = float(torch.sqrt((f["s"][:k] ** 2).sum())) * gain
        facts.append({"up_key": up_key, "down_key": down_key, "alpha_key": alpha_key,
                      "prefix": prefix, "f": f, "k": k, "norm": norm, "gain": gain,
                      "star": (gain if star_rescale and k < f["s"].shape[0] else None),
                      "is_gate": is_gate})

    # Pass 2: tame outlier layers — compress norms above the 90th percentile
    # back toward it. tame_layers interpolates between untouched (0) and fully
    # clamped to the percentile (1). Runs on post-rescale norms.
    if tame_layers > 0.0 and len(facts) >= 4:
        norms = torch.tensor([x["norm"] for x in facts])
        q = float(torch.quantile(norms, 0.90))
        if q > 0.0:
            for x in facts:
                if x["norm"] > q:
                    target = q + (x["norm"] - q) * (1.0 - tame_layers)
                    x["gain"] *= target / x["norm"]
                    x["tamed"] = True

    # Pass 3: rebuild the state dict (leaving genuinely-untouched layers as-is).
    out = dict(sd)
    rank_in, rank_out, tamed, gated = 0, 0, 0, 0
    for x in facts:
        up, down = sd[x["up_key"]], sd[x["down_key"]]
        rank_in += down.shape[0]
        # nothing to do: full rank kept and no gain change → leave original tensors
        if x["k"] == down.shape[0] and abs(x["gain"] - 1.0) < 1e-9:
            rank_out += x["k"]
            continue
        new_up, new_down = _rebuild(x["f"], x["k"], x["gain"], up, down)
        out[x["up_key"]] = new_up
        out[x["down_key"]] = new_down
        if x["alpha_key"] is not None:
            # scale is baked into the factors → make alpha/rank come out as 1
            out[x["alpha_key"]] = torch.tensor(float(x["k"]))
        rank_out += x["k"]
        if x.get("tamed"):
            tamed += 1
        if x["is_gate"] and gate_strength != 1.0:
            gated += 1

    hottest = sorted(facts, key=lambda x: x["norm"], reverse=True)[:5]
    hot_txt = ", ".join(f"{x['prefix'].split('.')[-1] or x['prefix']}={x['norm']:.3f}" for x in hottest)
    star_txt = ""
    if star_rescale:
        gains = [x["star"] for x in facts if x["star"] is not None]
        if gains:
            star_txt = f", STAR-rescaled {len(gains)} layers (avg x{sum(gains) / len(gains):.2f})"
    gate_txt = ""
    n_gate = sum(1 for x in facts if x["is_gate"])
    if n_gate:
        gate_txt = f", {n_gate} gate layers"
        if gate_strength != 1.0:
            gate_txt += f" scaled x{gate_strength:.2f}"
    logging.info(
        f"[Corza LoRA] {lora_name}: cleaned {len(facts)} layers, "
        f"total rank {rank_in} -> {rank_out}, tamed {tamed} hot layers{star_txt}{gate_txt} | hottest: {hot_txt}"
    )
    return out


def _adapter_alpha_value(alpha):
    if alpha is None:
        return None
    if hasattr(alpha, "item"):
        return alpha.item()
    return alpha


def _cleanable_lora_adapter(adapter):
    if getattr(adapter, "name", None) != "lora":
        return None
    weights = getattr(adapter, "weights", None)
    if not isinstance(weights, (list, tuple)) or len(weights) < 6:
        return None

    up, down, alpha, mid, dora_scale, reshape = weights[:6]
    if mid is not None or dora_scale is not None or reshape is not None:
        return None
    if not torch.is_tensor(up) or not torch.is_tensor(down):
        return None
    if up.ndim < 2 or down.ndim < 2:
        return None
    if any(d != 1 for d in up.shape[2:]):
        return None
    return up, down, _adapter_alpha_value(alpha)


@torch.no_grad()
def _clean_model_lora_patches(model, keep_energy, max_rank, tame_layers, star_rescale, gate_strength):
    facts = []
    for patch_key, patch_list in getattr(model, "patches", {}).items():
        for patch_index, patch in enumerate(patch_list):
            if len(patch) < 2:
                continue
            adapter = patch[1]
            cleanable = _cleanable_lora_adapter(adapter)
            if cleanable is None:
                continue

            up, down, alpha = cleanable
            f = _factorize(up, down, alpha)
            k = _pick_rank(f["s"], keep_energy, max_rank)
            gain = 1.0
            if star_rescale and k < f["s"].shape[0]:
                nuc_all = float(f["s"].sum())
                nuc_kept = float(f["s"][:k].sum())
                if nuc_kept > 0.0:
                    gain = nuc_all / nuc_kept
            # gates: multiplicative sigmoid controls, outsized/compounding effect
            is_gate = _is_gate(patch_key)
            if is_gate:
                gain *= gate_strength
            norm = float(torch.sqrt((f["s"][:k] ** 2).sum())) * gain
            facts.append({
                "patch_key": patch_key,
                "patch_index": patch_index,
                "adapter": adapter,
                "up": up,
                "down": down,
                "alpha": alpha,
                "f": f,
                "k": k,
                "norm": norm,
                "gain": gain,
                "star": (gain if star_rescale and k < f["s"].shape[0] else None),
                "is_gate": is_gate,
            })

    if not facts:
        logging.info("[Corza LoRA] applied patch cleaner: no cleanable LoRA patches found")
        return model

    if tame_layers > 0.0 and len(facts) >= 4:
        norms = torch.tensor([x["norm"] for x in facts])
        q = float(torch.quantile(norms, 0.90))
        if q > 0.0:
            for x in facts:
                if x["norm"] > q:
                    target = q + (x["norm"] - q) * (1.0 - tame_layers)
                    x["gain"] *= target / x["norm"]
                    x["tamed"] = True

    out_model = model.clone()
    rank_in, rank_out, tamed = 0, 0, 0
    for x in facts:
        rank_in += x["down"].shape[0]
        rank_out += x["k"]
        # untouched (full rank, no gain change) → leave the original patch as-is
        if x["k"] == x["down"].shape[0] and abs(x["gain"] - 1.0) < 1e-9:
            continue
        new_up, new_down = _rebuild(x["f"], x["k"], x["gain"], x["up"], x["down"])
        new_alpha = float(x["k"]) if x["alpha"] is not None else None
        new_weights = (new_up, new_down, new_alpha, None, None, None)
        loaded_keys = getattr(x["adapter"], "loaded_keys", set())
        new_adapter = type(x["adapter"])(loaded_keys, new_weights)

        patch_list = out_model.patches[x["patch_key"]]
        patch = patch_list[x["patch_index"]]
        patch_list[x["patch_index"]] = (patch[0], new_adapter, *patch[2:])

        if x.get("tamed"):
            tamed += 1

    out_model.patches_uuid = uuid.uuid4()
    hottest = sorted(facts, key=lambda x: x["norm"], reverse=True)[:5]
    hot_txt = ", ".join(f"{x['patch_key']}[{x['patch_index']}]={x['norm']:.3f}" for x in hottest)
    star_txt = ""
    if star_rescale:
        gains = [x["star"] for x in facts if x["star"] is not None]
        if gains:
            star_txt = f", STAR-rescaled {len(gains)} patches (avg x{sum(gains) / len(gains):.2f})"
    gate_txt = ""
    n_gate = sum(1 for x in facts if x["is_gate"])
    if n_gate:
        gate_txt = f", {n_gate} gate patches"
        if gate_strength != 1.0:
            gate_txt += f" scaled x{gate_strength:.2f}"
    logging.info(
        f"[Corza LoRA] applied patch cleaner: cleaned {len(facts)} patches, "
        f"total rank {rank_in} -> {rank_out}, tamed {tamed} hot patches{star_txt}{gate_txt} | hottest: {hot_txt}"
    )
    return out_model


class CorzaLoRAClean(io.ComfyNode):
    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="CorzaLoRAClean",
            display_name="Corza LoRA Loader (Clean)",
            category="corza/lora",
            description="LoRA loader with optional spectral cleanup: drops the noisy "
                        "low-energy tail of each layer (kohya resize math) and tames "
                        "outlier layers. Helps with blocky/crunchy artifacts, "
                        "especially on distilled/turbo models and stacked LoRAs. "
                        "With cleaning off it behaves exactly like the stock loader.",
            inputs=[
                io.Model.Input("model"),
                io.Clip.Input("clip", optional=True),
                io.Combo.Input("lora_name", options=folder_paths.get_filename_list("loras")),
                io.Float.Input("strength_model", default=1.0, min=-100.0, max=100.0, step=0.01),
                io.Float.Input("strength_clip", default=1.0, min=-100.0, max=100.0, step=0.01),
                io.Float.Input(
                    "keep_energy", default=100.0, min=50.0, max=100.0, step=0.5,
                    tooltip="Per layer, keep only the strongest SVD components adding up to "
                            "this % of the update's energy; the rest is mostly training "
                            "noise. 100 = off. Try 95, then 90 if artifacts persist.",
                ),
                io.Int.Input(
                    "max_rank", default=0, min=0, max=1024,
                    tooltip="Hard cap on each layer's rank after the energy cut. 0 = off.",
                ),
                io.Float.Input(
                    "tame_layers", default=0.0, min=0.0, max=1.0, step=0.05,
                    tooltip="Compress layers whose update is much stronger than the rest "
                            "(above the LoRA's 90th percentile) back toward the pack. "
                            "0 = off, 1 = fully clamped. Try 0.5 for crunchy edges.",
                ),
                io.Boolean.Input(
                    "star_rescale", default=False,
                    tooltip="STAR (arXiv:2502.10339): after keep_energy trims a layer, boost "
                            "the kept components so the layer's total strength matches the "
                            "original — trims the conflict-prone tail without weakening the "
                            "LoRA. Only does something when keep_energy < 100.",
                ),
                io.Float.Input(
                    "gate_strength", default=1.0, min=0.0, max=1.0, step=0.05,
                    tooltip="How much of the LoRA's effect reaches 'gate' layers only (Krea 2's "
                            "gated attention + SwiGLU gates, etc.). Gates are multiplicative "
                            "sigmoid controls, so LoRA edits there have outsized, compounding "
                            "effect — a big source of artifacts and of stacked LoRAs fighting, "
                            "which the norm-based tame_layers misses. 1 = full effect (default), "
                            "0 = strip the LoRA from gates (they stay at base). Try 0.5 if "
                            "stacked LoRAs deform. No effect on LoRAs without gate layers.",
                ),
            ],
            outputs=[
                io.Model.Output("model"),
                io.Clip.Output("clip"),
            ],
        )

    @classmethod
    def execute(
        cls,
        model,
        lora_name: str,
        strength_model: float = 1.0,
        strength_clip: float = 1.0,
        keep_energy: float = 100.0,
        max_rank: int = 0,
        tame_layers: float = 0.0,
        star_rescale: bool = False,
        gate_strength: float = 1.0,
        clip=None,
    ) -> io.NodeOutput:
        if strength_model == 0 and (clip is None or strength_clip == 0):
            return io.NodeOutput(model, clip)

        lora_path = folder_paths.get_full_path_or_raise("loras", lora_name)
        sd = comfy.utils.load_torch_file(lora_path, safe_load=True)

        cleaning = keep_energy < 100.0 or max_rank > 0 or tame_layers > 0.0 or gate_strength != 1.0
        if cleaning:
            sd = _clean_lora(sd, keep_energy, max_rank, tame_layers, star_rescale, gate_strength, lora_name)

        model_lora, clip_lora = comfy.sd.load_lora_for_models(
            model, clip, sd, strength_model, strength_clip
        )
        return io.NodeOutput(model_lora, clip_lora)


class CorzaCleanAppliedLoRAs(io.ComfyNode):
    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="CorzaCleanAppliedLoRAs",
            display_name="Corza Clean Applied LoRAs",
            category="corza/lora",
            description="Cleans standard LoRA patches already applied to the incoming "
                        "MODEL. Place it after upstream LoRA loaders; downstream LoRA "
                        "loaders remain untouched.",
            inputs=[
                io.Model.Input("model"),
                io.Float.Input(
                    "keep_energy", default=100.0, min=50.0, max=100.0, step=0.5,
                    tooltip="Per applied LoRA patch, keep only the strongest SVD "
                            "components adding up to this % of the update's energy. "
                            "100 = off.",
                ),
                io.Int.Input(
                    "max_rank", default=0, min=0, max=1024,
                    tooltip="Hard cap on each applied LoRA patch's rank after the "
                            "energy cut. 0 = off.",
                ),
                io.Float.Input(
                    "tame_layers", default=0.0, min=0.0, max=1.0, step=0.05,
                    tooltip="Compress applied LoRA patches whose update is much "
                            "stronger than the rest of the upstream LoRA stack. "
                            "0 = off, 1 = fully clamped.",
                ),
                io.Boolean.Input(
                    "star_rescale", default=False,
                    tooltip="STAR rescale after truncation. Only does something when "
                            "keep_energy trims a patch or max_rank caps it.",
                ),
                io.Float.Input(
                    "gate_strength", default=1.0, min=0.0, max=1.0, step=0.05,
                    tooltip="How much of the applied LoRA patches' effect reaches 'gate' layers "
                            "only (Krea 2 gated-attention + SwiGLU gates, etc.). Gates are "
                            "multiplicative sigmoid controls whose LoRA edits have outsized, "
                            "compounding effect — a big source of artifacts and of stacked "
                            "LoRAs fighting. 1 = full effect (default), 0 = strip the LoRA from "
                            "gates (they stay at base). Try 0.5 if stacked LoRAs deform.",
                ),
            ],
            outputs=[
                io.Model.Output("model"),
            ],
        )

    @classmethod
    def execute(
        cls,
        model,
        keep_energy: float = 100.0,
        max_rank: int = 0,
        tame_layers: float = 0.0,
        star_rescale: bool = False,
        gate_strength: float = 1.0,
    ) -> io.NodeOutput:
        cleaning = keep_energy < 100.0 or max_rank > 0 or tame_layers > 0.0 or gate_strength != 1.0
        if not cleaning:
            return io.NodeOutput(model)
        model_lora = _clean_model_lora_patches(
            model,
            keep_energy=keep_energy,
            max_rank=max_rank,
            tame_layers=tame_layers,
            star_rescale=star_rescale,
            gate_strength=gate_strength,
        )
        return io.NodeOutput(model_lora)


NODE_CLASS_MAPPINGS = {
    "CorzaLoRAClean": CorzaLoRAClean,
    "CorzaCleanAppliedLoRAs": CorzaCleanAppliedLoRAs,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CorzaLoRAClean": "Corza LoRA Loader (Clean)",
    "CorzaCleanAppliedLoRAs": "Corza Clean Applied LoRAs",
}