#!/usr/bin/env python3
"""Gradient-flow probe for v2 vision-tower loaders."""
from __future__ import annotations

import gc
import os

import torch

os.environ.setdefault("HF_HOME", os.path.expanduser("~/workspace/hf-cache"))

from internvit_loader import image_feat_internvit, load_internvit
from minimax_loader import image_feat_minimax, load_minimax
from moonvit_loader import image_feat_moonvit, load_moonvit


def _probe(name: str, load_fn, feat_fn, h: int, w: int) -> dict:
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    gc.collect()

    model = load_fn()
    x = torch.rand(1, 3, h, w, device="cuda", requires_grad=True)
    feat = feat_fn(model, x)
    loss = feat.float().pow(2).sum()
    grad = torch.autograd.grad(loss, x)[0]
    grad_l2 = grad.float().norm().item()
    vram_gb = torch.cuda.max_memory_allocated() / (1024**3)
    feat_shape = tuple(feat.shape)

    ok = bool(torch.isfinite(torch.tensor(grad_l2)).item() and grad_l2 > 0)
    print(
        f"{name}: feat.shape={feat_shape} grad_l2={grad_l2:.6f} "
        f"peak_vram_gb={vram_gb:.3f} PASS={ok}"
    )

    del model, x, feat, grad, loss
    gc.collect()
    torch.cuda.empty_cache()

    return {"name": name, "feat_shape": feat_shape, "grad_l2": grad_l2, "vram_gb": vram_gb, "pass": ok}


def main() -> None:
    assert torch.cuda.is_available(), "CUDA required"
    print(f"torch {torch.__version__} cuda={torch.cuda.is_available()}")

    results = [
        _probe("moonvit", load_moonvit, image_feat_moonvit, 448, 448),
        _probe("internvit", load_internvit, image_feat_internvit, 448, 448),
        _probe("minimax", load_minimax, image_feat_minimax, 672, 672),
    ]

    failed = [r["name"] for r in results if not r["pass"]]
    if failed:
        raise SystemExit(f"FAILED towers: {failed}")
    print("All towers passed gradient probe.")


if __name__ == "__main__":
    main()