"""10Eros-Max (TenStrip) beta_5 support: replace the diffusers transformer's weights with the Eros checkpoint.

`TenStrip/10Eros-Max` ships full bf16 checkpoints in the *reference* (ComfyUI) module tree, not LoRAs, so the update is
a straight weight replacement rather than a fold. Each tensor is renamed and row-transformed exactly like the
`larryvrh` LoRA deltas in `h3_lora` (fused-QKV thirds onto `attn.to_q/k/v`, the `SwiGLU` gate/value swap onto
`ff.net.0.proj`, `blocks.` -> `transformer_blocks.`, ...). The mapping was checked against the base weights: block 0
q/k/v/out_proj/fc1 all sit within ~1% cosine of MiniMax-H3 after that transform, and fc1 does not match without the swap.

AdaLN: the reference file stores the projections rank-8 compressed over the 1025-step timestep grid — `adaln_t_table[i] =
adaln_basis @ (silu(temb(i / 1024)) - adaln_mean)` and `*.adaln_proj.linear` of shape `[out, 8]`. Diffusers computes the
full `linear(silu(temb))`, so each projection is rebuilt as `W = W_ref @ basis`, `b = b_ref - W @ mean`. Checked against
the base model (time embedder from MiniMax-H3, which the file does not carry): the table matches at cos 0.9998 and the
rebuilt output matches the reference form at cos 1.0000 (0.9993-0.9999 vs. base weights). Exact on the sampler's
timestep grid, a rank-8 approximation off it; the time embedder itself stays base. `rope.inv_freq` is skipped.

Tensors are streamed one at a time with `safe_open`, so peak CPU memory is one tensor rather than the 40 GB file.

`H3_EROS=off` skips it, `H3_EROS_REPO` / `H3_EROS_FILE` pick the checkpoint. A `TURBO` file has the 4-step delta baked
in, which is why `app.py` then defaults the LoRA selector to `off` — stacking larry/lightx on top double-applies it.
"""

from __future__ import annotations

import os

import torch

EROS_REPO = os.environ.get("H3_EROS_REPO", "TenStrip/10Eros-Max")
EROS_FILE = os.environ.get("H3_EROS_FILE", "10Eros_Max_h3_TURBO-hybrid_beta5.safetensors")

_SKIP_EXACT = {"rope.inv_freq", "adaln_basis", "adaln_mean", "adaln_t_table"}
_DIRECT = {
    "video_patch_proj": "proj_in",
    "audio_patch_proj": "audio_proj_in",
    "condition_proj": "context_embedder",
    "final_layer.video_out": "proj_out",
    "final_layer.audio_out": "audio_proj_out",
    "final_layer.norm": "norm_out.norm",
}


def enabled() -> bool:
    return EROS_FILE.lower() not in ("", "off", "none") and os.environ.get("H3_EROS", "on").lower() not in ("", "off", "none")


def turbo_baked() -> bool:
    return enabled() and "turbo" in EROS_FILE.lower()


def adaln_target(name: str) -> str | None:
    """Diffusers module prefix for a compressed reference AdaLN projection tensor, else `None`."""
    if name.startswith("final_layer.adaln_proj.linear."):
        return "norm_out.linear"
    if name.startswith("blocks.") and ".adaln_proj.linear." in name:
        return name.replace("blocks.", "transformer_blocks.", 1).rsplit(".", 1)[0]
    return None


def rebuild_adaln(w_ref: torch.Tensor, b_ref: torch.Tensor, basis: torch.Tensor, mean: torch.Tensor):
    """Rank-8 reference projection -> full `[out, time_embed_dim]` weight and bias (float32)."""
    w = w_ref.float() @ basis.float()
    return w, b_ref.float() - w @ mean.float()


def map_tensor(name: str, tensor: torch.Tensor, inner_dim: int) -> list[tuple[str, torch.Tensor]]:
    """Reference-tree tensor -> [(diffusers key, tensor)]. Returns [] for tensors with no diffusers counterpart."""
    if name in _SKIP_EXACT or adaln_target(name) is not None:
        return []
    stem, _, leaf = name.rpartition(".")  # leaf is `weight` / `bias`
    if stem in _DIRECT:
        return [(f"{_DIRECT[stem]}.{leaf}", tensor)]

    if name.startswith("token_refiner.blocks."):
        target = name.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1)
    elif name.startswith("blocks."):
        target = name.replace("blocks.", "transformer_blocks.", 1)
    else:
        target = name  # token_refiner.final_norm.weight

    if target.endswith(".attn.qkv_proj.weight"):
        prefix = target.removesuffix("qkv_proj.weight")
        return [
            (f"{prefix}to_{kind}.weight", part.contiguous())
            for kind, part in zip(("q", "k", "v"), tensor.split(inner_dim, dim=0))
        ]
    if target.endswith(".mlp.fc1.weight"):
        gate, value = tensor.chunk(2, dim=0)
        return [(target.replace(".mlp.fc1.weight", ".ff.net.0.proj.weight"), torch.cat([value, gate]).contiguous())]
    if target.endswith(".mlp.fc2.weight"):
        return [(target.replace(".mlp.fc2.weight", ".ff.net.2.weight"), tensor)]
    if target.endswith(".attn.out_proj.weight"):
        return [(target.replace(".attn.out_proj.weight", ".attn.to_out.0.weight"), tensor)]
    if target.endswith(".attn.q_norm.weight"):
        return [(target.replace(".attn.q_norm.weight", ".attn.norm_q.weight"), tensor)]
    if target.endswith(".attn.k_norm.weight"):
        return [(target.replace(".attn.k_norm.weight", ".attn.norm_k.weight"), tensor)]
    return [(target, tensor)]  # norm1 / norm2 / token_refiner.final_norm


def apply_eros(transformer) -> str | None:
    """Replace the transformer's weights in place with the Eros checkpoint. Returns a status line, or `None` if off."""
    if not enabled():
        return None
    from huggingface_hub import hf_hub_download
    from safetensors import safe_open

    inner_dim = transformer.config.num_attention_heads * transformer.config.attention_head_dim
    params = dict(transformer.named_parameters())
    path = hf_hub_download(EROS_REPO, EROS_FILE)

    replaced, skipped = 0, []  # skipped: reference-only tensors (basis / mean / table / rope)
    with safe_open(path, framework="pt", device="cpu") as f:
        basis, mean = f.get_tensor("adaln_basis"), f.get_tensor("adaln_mean")
        for name in f.keys():
            if name.endswith("adaln_proj.linear.weight") and (prefix := adaln_target(name)):
                w, b = rebuild_adaln(f.get_tensor(name), f.get_tensor(name.replace(".weight", ".bias")), basis, mean)
                for key, tensor in ((f"{prefix}.weight", w), (f"{prefix}.bias", b)):
                    if key not in params or tuple(params[key].shape) != tuple(tensor.shape):
                        raise ValueError(f"AdaLN `{key}`: transformer vs Eros {tuple(tensor.shape)} mismatch")
                    params[key].data.copy_(tensor.to(params[key].dtype))
                    replaced += 1
                continue
            targets = map_tensor(name, f.get_tensor(name), inner_dim)
            if not targets:
                skipped.append(name)
                continue
            for key, tensor in targets:
                param = params.get(key)
                if param is None:
                    raise KeyError(f"Eros tensor `{name}` -> `{key}` not found in the transformer")
                if tuple(param.shape) != tuple(tensor.shape):
                    raise ValueError(f"`{key}`: transformer {tuple(param.shape)} vs Eros {tuple(tensor.shape)}")
                param.data.copy_(tensor.to(param.dtype))
                replaced += 1

    return (
        f"Eros `{EROS_REPO}/{EROS_FILE}`: {replaced} weights replaced, {len(skipped)} reference-only tensors skipped "
        f"(AdaLN rebuilt from rank-8){' · turbo baked in' if turbo_baked() else ''}"
    )