"""Differentiability + cost probe for the v0.2.1 Tier-1 candidate towers.

Each tower is probed in isolation (own try/except) so one failure doesn't sink the
rest. For each we check: loads, forward produces a pooled feature, input gradient is
finite and nonzero (pixels -> feature), plus VRAM and fwd/bwd latency. Results are
written to runs/tier1_probe.json.

Run on a CUDA GPU host:
  HF_HOME=~/workspace/hf-cache HF_TOKEN=... python ensemble/towers/probe_tier1.py
"""

from __future__ import annotations

import json
import time
import traceback
from pathlib import Path

import torch
import torch.nn.functional as F

DEV = "cuda"
OUT = Path("runs/tier1_probe.json")
results: dict[str, dict] = {}


def _grad_ok(feat: torch.Tensor, x: torch.Tensor) -> bool:
    g = torch.autograd.grad(feat.float().pow(2).sum(), x, retain_graph=False)[0]
    return bool(torch.isfinite(g).all()) and float(g.abs().sum()) > 0


def record(name, **kw):
    results[name] = kw
    OUT.write_text(json.dumps(results, indent=2))
    print(f"[{name}] " + " ".join(f"{k}={v}" for k, v in kw.items()), flush=True)


def probe_openclip(name, arch, tag, res, force_qgelu=False):
    import open_clip
    t0 = time.time()
    torch.cuda.reset_peak_memory_stats()
    m, _, _ = open_clip.create_model_and_transforms(arch, pretrained=tag,
                                                     force_quick_gelu=force_qgelu)
    m = m.to(DEV, torch.float16).eval()
    for p in m.parameters():
        p.requires_grad_(False)
    load_s = time.time() - t0
    x = torch.rand(1, 3, res, res, device=DEV, dtype=torch.float16, requires_grad=True)
    t1 = time.time()
    f = m.encode_image(x)
    ok = _grad_ok(f, x)
    lat = time.time() - t1
    has_text = hasattr(m, "encode_text")
    tdim = None
    if has_text:
        tok = open_clip.get_tokenizer(arch)
        tf = m.encode_text(tok(["a dog", "a church"]).to(DEV))
        tdim = tuple(tf.shape)
    record(name, backend="openclip", kind="contrastive", feat=tuple(f.shape),
           text_feat=tdim, grad_ok=ok, load_s=round(load_s), lat_s=round(lat, 2),
           vram_gb=round(torch.cuda.max_memory_allocated() / 1e9, 1))
    del m
    torch.cuda.empty_cache()


def probe_timm(name, model_id, res):
    import timm
    t0 = time.time()
    torch.cuda.reset_peak_memory_stats()
    m = timm.create_model(model_id, pretrained=True, num_classes=0).to(DEV, torch.float16).eval()
    for p in m.parameters():
        p.requires_grad_(False)
    load_s = time.time() - t0
    x = torch.rand(1, 3, res, res, device=DEV, dtype=torch.float16, requires_grad=True)
    t1 = time.time()
    f = m(x)
    if f.dim() == 3:
        f = f.mean(1)
    ok = _grad_ok(f, x)
    record(name, backend="timm", kind="feature", feat=tuple(f.shape), grad_ok=ok,
           load_s=round(load_s), lat_s=round(time.time() - t1, 2),
           vram_gb=round(torch.cuda.max_memory_allocated() / 1e9, 1))
    del m
    torch.cuda.empty_cache()


def probe_cradio(name, repo, res):
    from transformers import AutoModel
    t0 = time.time()
    torch.cuda.reset_peak_memory_stats()
    m = AutoModel.from_pretrained(repo, trust_remote_code=True).eval().to(DEV)
    for p in m.parameters():
        p.requires_grad_(False)
    load_s = time.time() - t0
    x = torch.rand(1, 3, res, res, device=DEV, requires_grad=True)
    t1 = time.time()
    out = m(x)
    summ = out.summary if hasattr(out, "summary") else (out[0] if isinstance(out, (tuple, list)) else out)
    ok = _grad_ok(summ, x)
    record(name, backend="cradio", kind="feature", feat=tuple(summ.shape), grad_ok=ok,
           load_s=round(load_s), lat_s=round(time.time() - t1, 2),
           vram_gb=round(torch.cuda.max_memory_allocated() / 1e9, 1))
    del m
    torch.cuda.empty_cache()


def probe_qwen3vl(name, repo):
    """Differentiable path for Qwen3-VL's NaViT vision tower from a plain (1,3,H,W)."""
    from transformers import Qwen3VLVisionModel, AutoConfig
    t0 = time.time()
    torch.cuda.reset_peak_memory_stats()
    cfg = AutoConfig.from_pretrained(repo)
    vcfg = cfg.vision_config
    m = Qwen3VLVisionModel.from_pretrained(repo, dtype=torch.float16,
                                           attn_implementation="sdpa").eval().to(DEV)
    for p in m.parameters():
        p.requires_grad_(False)
    load_s = time.time() - t0
    ps = vcfg.patch_size
    merge = getattr(vcfg, "spatial_merge_size", 2)
    # pick H=W that's a multiple of patch_size*merge
    grid = merge * 8                      # 8 merged tokens per side
    hpx = wpx = grid * ps
    x = torch.rand(1, 3, hpx, wpx, device=DEV, dtype=torch.float16, requires_grad=True)
    gh, gw = hpx // ps, wpx // ps
    # Build the flattened patch sequence Qwen3VL expects: (t*gh*gw, C*ps*ps), differentiable.
    t1 = time.time()
    xp = x.reshape(1, 3, gh, ps, gw, ps).permute(0, 2, 4, 1, 3, 5).reshape(gh * gw, 3 * ps * ps)
    grid_thw = torch.tensor([[1, gh, gw]], device=DEV)
    out = m(hidden_states=xp, grid_thw=grid_thw)
    feat = out[0] if isinstance(out, (tuple, list)) else (out.last_hidden_state if hasattr(out, "last_hidden_state") else out)
    if feat.dim() == 2:
        pooled = feat.mean(0, keepdim=True)
    else:
        pooled = feat.reshape(-1, feat.shape[-1]).mean(0, keepdim=True)
    ok = _grad_ok(pooled, x)
    record(name, backend="qwen3vl", kind="feature", feat=tuple(pooled.shape),
           grad_ok=ok, load_s=round(load_s), lat_s=round(time.time() - t1, 2),
           vram_gb=round(torch.cuda.max_memory_allocated() / 1e9, 1),
           patch_size=ps, merge=merge)
    del m
    torch.cuda.empty_cache()


PROBES = [
    ("siglip2-giant", lambda: probe_openclip("siglip2-giant", "ViT-gopt-16-SigLIP2-384", "webli", 384)),
    ("metaclip2-H-worldwide", lambda: probe_openclip("metaclip2-H-worldwide", "ViT-H-14-worldwide-378", "metaclip2_worldwide", 378)),
    ("c-radio-v3-h", lambda: probe_cradio("c-radio-v3-h", "nvidia/C-RADIOv3-H", 384)),
    ("qwen3-vl-vision", lambda: probe_qwen3vl("qwen3-vl-vision", "Qwen/Qwen3-VL-8B-Instruct")),
]


def main():
    OUT.parent.mkdir(parents=True, exist_ok=True)
    for name, fn in PROBES:
        try:
            fn()
        except Exception as e:
            record(name, error=str(e)[:200])
            traceback.print_exc()
    print("PROBE_TIER1_DONE", flush=True)


if __name__ == "__main__":
    main()