"""Verify that all tipsv2 and tipsv2-dpt repos load via the three loading APIs.

Each repo is loaded twice per API: once with trust_remote_code=True (remote
custom code) and once with trust_remote_code=False (native transformers
integration). Both must pass.

In addition to pass/fail, this version reports the first few unexpected and
missing weights for each load via output_loading_info=True.
"""

from transformers import AutoModel, AutoBackbone, Tipsv2VisionModel, Tipsv2TextModel, AutoModelForDepthEstimation, AutoModelForNormalEstimation, AutoModelForSemanticSegmentation, Tipsv2DptModel, Tipsv2DptForDepthEstimation, Tipsv2DptForNormalEstimation, Tipsv2DptForSemanticSegmentation

REPOS = [
    "guarin/tipsv2-b14",
    "guarin/tipsv2-l14",
    "guarin/tipsv2-so400m14",
    "guarin/tipsv2-g14",
    "guarin/tipsv2-b14-dpt",
    "guarin/tipsv2-l14-dpt",
    "guarin/tipsv2-so400m14-dpt",
    "guarin/tipsv2-g14-dpt",
]

LOADERS = [
    ("AutoModel", AutoModel),
    ("AutoBackbone", AutoBackbone),
    ("Tipsv2VisionModel", Tipsv2VisionModel),
    ("Tipsv2TextModel", Tipsv2TextModel),
]

REMOTE_CODE_MODES = [True, False]

results = {}
for repo in REPOS:
    for name, loader in LOADERS:
        for trust_remote_code in REMOTE_CODE_MODES:
            # AutoBackbone is never wired into the repos' auto_map, so it can only
            # resolve through the native integration (trust_remote_code=False).
            if name == "AutoBackbone" and (trust_remote_code or "dpt" in name):
                continue
            try:
                model, loading_info = loader.from_pretrained(
                    repo,
                    trust_remote_code=trust_remote_code,
                    output_loading_info=True,
                )
                # These may come back as sets (native path) or lists (remote
                # path); normalise to a sorted list so we can slice deterministically.
                unexpected = sorted(loading_info.get("unexpected_keys") or [])
                missing = sorted(loading_info.get("missing_keys") or [])
                results[(repo, name, trust_remote_code)] = (
                    "OK",
                    type(model).__name__,
                    unexpected,
                    missing,
                )
                del model
            except Exception as error:
                results[(repo, name, trust_remote_code)] = (
                    "FAIL",
                    f"{type(error).__name__}: {error}",
                    [],
                    [],
                )

print("\n==================== DETAILS ====================")
for (repo, name, trust_remote_code), (status, info, unexpected, missing) in results.items():
    print(f"[{status}] {repo} via {name} (remote_code={trust_remote_code}) -> {info}")
    if unexpected:
        print(f"        unexpected (first 3): {unexpected[:3]}")
    if missing:
        print(f"        missing (first 3):    {missing[:3]}")

print("\n\n==================== SUMMARY ====================")
print("Each cell shows [remote_code=True / remote_code=False].")
print("Suffix flags: U = has unexpected weights, M = has missing weights.")
header = f"{'repo':<28}{'AutoModel':<20}{'AutoBackbone':<20}{'Tipsv2VisionModel':<20}{'Tipsv2TextModel':<20}"
print(header)
print("-" * len(header))


def format_cell(repo, name):
    """Build the [remote/native] status cell with U/M flags for one loader."""
    parts = []
    for trust_remote_code in REMOTE_CODE_MODES:
        key = (repo, name, trust_remote_code)
        if key not in results:
            parts.append("-")
            continue
        status, _, unexpected, missing = results[key]
        flags = ""
        if unexpected:
            flags += "U"
        if missing:
            flags += "M"
        parts.append(f"{status}{flags}")
    return "/".join(parts)


for repo in REPOS:
    cells = [format_cell(repo, name) for name, _ in LOADERS]
    print(f"{repo:<28}{cells[0]:<20}{cells[1]:<20}{cells[2]:<20}{cells[3]:<20}")

all_ok = all(status == "OK" for (status, *_) in results.values())
print(f"\nAll loads passed: {all_ok}")