import os
import torch
from collections import Counter, defaultdict

ROOT = r"E:\Models SD XL\UNET\Minimax_H3"

H3_FILE = os.path.join(
    ROOT,
    "h3_bridge_hidden_480.pt"
)

SN_FILE = os.path.join(
    ROOT,
    "sensenova_bridge_hidden_480.pt"
)

BRIDGE_FILE = (
    r"D:\ComfyUI_Krea2\ComfyUI\models\bridge"
    r"\SN_L32_to_H3_L49_rank128.safetensors"
)

REPORT_FILE = os.path.join(
    ROOT,
    "distillation_source_inspection.txt"
)


# ============================================================
# HELPERS
# ============================================================

def tensor_info(t):
    return (
        f"shape={tuple(t.shape)} "
        f"dtype={t.dtype} "
        f"device={t.device}"
    )


def walk(obj, path="root", out=None, depth=0, max_depth=8):
    if out is None:
        out = []

    if depth > max_depth:
        return out

    if torch.is_tensor(obj):
        out.append({
            "path": path,
            "shape": tuple(obj.shape),
            "dtype": str(obj.dtype),
            "tensor": obj,
        })
        return out

    if isinstance(obj, dict):
        for key, value in obj.items():
            walk(
                value,
                f"{path}.{key}",
                out,
                depth + 1,
                max_depth,
            )

    elif isinstance(obj, (list, tuple)):
        for i, value in enumerate(obj):
            walk(
                value,
                f"{path}[{i}]",
                out,
                depth + 1,
                max_depth,
            )

    return out


def print_structure(obj, name, lines, depth=0, max_depth=4, max_items=4):
    indent = "  " * depth

    if depth > max_depth:
        lines.append(f"{indent}...")
        return

    if torch.is_tensor(obj):
        lines.append(
            f"{indent}{tensor_info(obj)}"
        )
        return

    if isinstance(obj, dict):
        lines.append(
            f"{indent}dict keys={len(obj)}"
        )

        for i, (key, value) in enumerate(obj.items()):
            if i >= max_items:
                lines.append(
                    f"{indent}  ... ({len(obj) - max_items} more)"
                )
                break

            lines.append(
                f"{indent}  KEY: {repr(key)}"
            )

            print_structure(
                value,
                name,
                lines,
                depth + 2,
                max_depth,
                max_items,
            )

    elif isinstance(obj, (list, tuple)):
        lines.append(
            f"{indent}{type(obj).__name__} len={len(obj)}"
        )

        for i, value in enumerate(obj[:max_items]):
            lines.append(
                f"{indent}  ITEM [{i}]"
            )

            print_structure(
                value,
                name,
                lines,
                depth + 2,
                max_depth,
                max_items,
            )

        if len(obj) > max_items:
            lines.append(
                f"{indent}  ... ({len(obj) - max_items} more)"
            )

    else:
        value_repr = repr(obj)

        if len(value_repr) > 180:
            value_repr = value_repr[:180] + "..."

        lines.append(
            f"{indent}{type(obj).__name__}: {value_repr}"
        )


def analyze_file(label, path, lines):
    lines.append("")
    lines.append("=" * 100)
    lines.append(label)
    lines.append("=" * 100)
    lines.append(f"Path: {path}")
    lines.append(
        f"Size: {os.path.getsize(path) / (1024 ** 2):.2f} MiB"
    )

    print(f"\nLoading {label}...")
    data = torch.load(
        path,
        map_location="cpu",
        weights_only=False,
    )

    print(f"{label} loaded.")

    lines.append("")
    lines.append("TOP-LEVEL STRUCTURE")
    lines.append("-" * 100)

    print_structure(
        data,
        label,
        lines,
    )

    tensors = walk(data)

    lines.append("")
    lines.append(
        f"Total tensors found recursively: {len(tensors)}"
    )

    shape_counter = Counter(
        item["shape"]
        for item in tensors
    )

    lines.append("")
    lines.append("MOST COMMON TENSOR SHAPES")
    lines.append("-" * 100)

    for shape, count in shape_counter.most_common(20):
        lines.append(
            f"{count:6d} x {shape}"
        )

    # ========================================================
    # HIDDEN-STATE CANDIDATES
    # ========================================================

    candidates_4096 = []
    candidates_5120 = []

    for item in tensors:
        shape = item["shape"]

        if len(shape) >= 1:
            if shape[-1] == 4096:
                candidates_4096.append(item)

            if shape[-1] == 5120:
                candidates_5120.append(item)

    lines.append("")
    lines.append("4096-DIM CANDIDATES")
    lines.append("-" * 100)
    lines.append(
        f"Count: {len(candidates_4096)}"
    )

    for item in candidates_4096[:30]:
        lines.append(
            f"{item['path']} | "
            f"{item['shape']} | "
            f"{item['dtype']}"
        )

    if len(candidates_4096) > 30:
        lines.append(
            f"... {len(candidates_4096) - 30} more"
        )

    lines.append("")
    lines.append("5120-DIM CANDIDATES")
    lines.append("-" * 100)
    lines.append(
        f"Count: {len(candidates_5120)}"
    )

    for item in candidates_5120[:30]:
        lines.append(
            f"{item['path']} | "
            f"{item['shape']} | "
            f"{item['dtype']}"
        )

    if len(candidates_5120) > 30:
        lines.append(
            f"... {len(candidates_5120) - 30} more"
        )

    # ========================================================
    # PATHS THAT LOOK LIKE OUR TARGET LAYERS
    # ========================================================

    lines.append("")
    lines.append("PATHS CONTAINING LAYER 32 / L32")
    lines.append("-" * 100)

    layer32 = [
        item
        for item in tensors
        if (
            "32" in item["path"].lower()
            or "l32" in item["path"].lower()
            or "layer_32" in item["path"].lower()
            or "layer32" in item["path"].lower()
        )
    ]

    lines.append(
        f"Count: {len(layer32)}"
    )

    for item in layer32[:40]:
        lines.append(
            f"{item['path']} | "
            f"{item['shape']}"
        )

    lines.append("")
    lines.append("PATHS CONTAINING LAYER 49 / L49")
    lines.append("-" * 100)

    layer49 = [
        item
        for item in tensors
        if (
            "49" in item["path"].lower()
            or "l49" in item["path"].lower()
            or "layer_49" in item["path"].lower()
            or "layer49" in item["path"].lower()
        )
    ]

    lines.append(
        f"Count: {len(layer49)}"
    )

    for item in layer49[:40]:
        lines.append(
            f"{item['path']} | "
            f"{item['shape']}"
        )

    return data, tensors


# ============================================================
# MAIN
# ============================================================

print("=" * 100)
print("DISTILLATION SOURCE INSPECTION")
print("=" * 100)

for path in [
    H3_FILE,
    SN_FILE,
    BRIDGE_FILE,
]:
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"Missing required file:\n{path}"
        )

print("H3:")
print(H3_FILE)

print("\nSenseNova:")
print(SN_FILE)

print("\nBridge:")
print(BRIDGE_FILE)

lines = []

lines.append(
    "SenseNova -> MiniMax H3 Distillation Source Inspection"
)

lines.append(
    "=" * 100
)

h3_data, h3_tensors = analyze_file(
    "MINIMAX H3 DATA",
    H3_FILE,
    lines,
)

sn_data, sn_tensors = analyze_file(
    "SENSENOVA DATA",
    SN_FILE,
    lines,
)


# ============================================================
# BRIDGE
# ============================================================

lines.append("")
lines.append("=" * 100)
lines.append("BRIDGE FILE")
lines.append("=" * 100)
lines.append(f"Path: {BRIDGE_FILE}")

from safetensors.torch import load_file

bridge = load_file(
    BRIDGE_FILE,
    device="cpu",
)

for key, tensor in bridge.items():
    lines.append(
        f"{key:40s} "
        f"{tuple(tensor.shape)} "
        f"{tensor.dtype}"
    )


# ============================================================
# SUMMARY COUNTS
# ============================================================

h3_5120 = [
    x for x in h3_tensors
    if (
        len(x["shape"]) > 0
        and x["shape"][-1] == 5120
    )
]

sn_4096 = [
    x for x in sn_tensors
    if (
        len(x["shape"]) > 0
        and x["shape"][-1] == 4096
    )
]

lines.append("")
lines.append("=" * 100)
lines.append("FINAL SUMMARY")
lines.append("=" * 100)

lines.append(
    f"H3 5120-dim tensors found:       {len(h3_5120)}"
)

lines.append(
    f"SenseNova 4096-dim tensors found: {len(sn_4096)}"
)

lines.append("")
lines.append(
    "The next script will use this report to identify the exact "
    "H3 input representation and SenseNova teacher target without "
    "guessing the .pt layout."
)


# ============================================================
# SAVE
# ============================================================

with open(
    REPORT_FILE,
    "w",
    encoding="utf-8",
) as f:
    f.write(
        "\n".join(lines)
    )

print()
print("=" * 100)
print("SUCCESS")
print("=" * 100)

print("Report saved:")
print(REPORT_FILE)

print()
print(
    "H3 5120-dim tensors:",
    len(h3_5120)
)

print(
    "SenseNova 4096-dim tensors:",
    len(sn_4096)
)

print()
print(
    "Send me the contents of "
    "distillation_source_inspection.txt"
)