"""A+2: Run SAM / DepthAnything / DINOv2 on a fixed test image; record feature stats.

Goal: detect silent expert-checkpoint or preprocessing mismatch BEFORE 22h training.
This is codex xhigh's "Top single risk" — must not be skipped.
"""
import os, sys, json, hashlib
import numpy as np
import torch
from PIL import Image
ROOT = "/root/autodl-tmp/covt_repro"
DATA = f"{ROOT}/covt_data"
OUT  = f"{ROOT}/diagnostics_a_plus/A2_expert_features.json"

assert torch.cuda.is_available(), "A+2 requires GPU"
device = "cuda"
torch.set_grad_enabled(False)

# Fixed test image — use first CoVT-Dataset parquet sample for reproducibility
import pyarrow.parquet as pq, glob, io
parq = sorted(glob.glob(f"{DATA}/dataset/CoVT-Dataset/part1/*.parquet"))
assert parq, "no parquet found"
tbl = pq.read_table(parq[0])
sample = tbl.slice(0,1).to_pylist()[0]
img_bytes = sample.get("image", {}).get("bytes") if isinstance(sample.get("image"), dict) else None
if img_bytes is None:
    # fall back to test image from CoVT assets
    for cand in glob.glob(f"{ROOT}/CoVT/assets/*.jpg") + glob.glob(f"{ROOT}/CoVT/assets/*.png"):
        with open(cand,"rb") as f: img_bytes = f.read()
        break
assert img_bytes, "no test image found"
img_sha = hashlib.sha256(img_bytes).hexdigest()[:16]
img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
print(f"Test image SHA-16: {img_sha}  size: {img.size}")

results = {"image_sha16": img_sha, "image_size": list(img.size)}

# --- DINOv2 ---
print("\n[DINOv2]")
try:
    from transformers import AutoImageProcessor, AutoModel
    proc = AutoImageProcessor.from_pretrained(f"{DATA}/models/dinov2-large")
    mdl  = AutoModel.from_pretrained(f"{DATA}/models/dinov2-large", torch_dtype=torch.bfloat16).to(device).eval()
    inp = proc(images=img, return_tensors="pt").to(device, dtype=torch.bfloat16)
    out = mdl(**inp).last_hidden_state  # (1, N+1, D)
    f = out.float().cpu().numpy()
    results["dinov2"] = {
        "shape": list(f.shape),
        "mean": float(np.mean(f)),
        "std":  float(np.std(f)),
        "l2_per_token_mean": float(np.linalg.norm(f, axis=-1).mean()),
        "cls_l2": float(np.linalg.norm(f[0,0])),
    }
    print(json.dumps(results["dinov2"], indent=2))
    del mdl; torch.cuda.empty_cache()
except Exception as e:
    results["dinov2"] = {"error": str(e)}; print("dinov2 FAIL:", e)

# --- Depth-Anything V2 ---
print("\n[DepthAnything V2]")
try:
    from transformers import AutoImageProcessor as P2, AutoModelForDepthEstimation
    proc = P2.from_pretrained(f"{DATA}/models/Depth-Anything-V2-Large")
    mdl  = AutoModelForDepthEstimation.from_pretrained(f"{DATA}/models/Depth-Anything-V2-Large", torch_dtype=torch.bfloat16).to(device).eval()
    inp = proc(images=img, return_tensors="pt").to(device, dtype=torch.bfloat16)
    out = mdl(**inp).predicted_depth  # (1, H, W)
    d = out.float().cpu().numpy()
    results["depth_anything"] = {
        "shape": list(d.shape),
        "mean": float(np.mean(d)),
        "std":  float(np.std(d)),
        "min":  float(np.min(d)),
        "max":  float(np.max(d)),
    }
    print(json.dumps(results["depth_anything"], indent=2))
    del mdl; torch.cuda.empty_cache()
except Exception as e:
    results["depth_anything"] = {"error": str(e)}; print("depth FAIL:", e)

# --- SAM ViT-H ---
print("\n[SAM ViT-H]")
try:
    # CoVT uses segment-anything package; if unavailable use raw checkpoint load + dim check
    ckpt = f"{DATA}/models/sam/sam_vit_h_4b8939.pth"
    sd = torch.load(ckpt, map_location="cpu", weights_only=False)
    keys = list(sd.keys()) if isinstance(sd, dict) else []
    img_emb_keys = [k for k in keys if "image_encoder" in k]
    results["sam"] = {
        "ckpt_path": ckpt,
        "ckpt_size_bytes": os.path.getsize(ckpt),
        "ckpt_sha256_16": hashlib.sha256(open(ckpt,"rb").read(1024*1024*64)).hexdigest()[:16],
        "total_keys": len(keys),
        "image_encoder_keys": len(img_emb_keys),
        "sample_keys": img_emb_keys[:5],
    }
    print(json.dumps(results["sam"], indent=2))
except Exception as e:
    results["sam"] = {"error": str(e)}; print("sam FAIL:", e)

with open(OUT,"w") as f: json.dump(results, f, indent=2)
print(f"\nWrote {OUT}")