#!/usr/bin/env python
"""Verify WeDetect-Base-Uni (proposal generator) on its own end-to-end on MPS."""
import os, sys, time, math
import torch
from PIL import Image, ImageDraw, ImageFont
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from generate_proposal import SimpleYOLOWorldDetector

CKP = "/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-base-uni/wedetect_base_uni.pth"
IMG = "/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/assets/demo.jpeg"
OUT = "/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/outputs"


def remap_uni_keys(ck):
    keys = list(ck.keys())
    for k in keys:
        if "backbone" in k:
            ck[k.replace("backbone.image_model.model.", "backbone.")] = ck.pop(k)
    keys = list(ck.keys())
    for k in keys:
        if "bbox_head" in k:
            nk = k.replace("bbox_head.head_module.", "bbox_head.")
            for a, b in [("0.2.","0.6."),("1.2.","1.6."),("2.2.","2.6."),("1.bn","4"),("1.conv","3"),("0.bn","1"),("0.conv","0")]:
                nk = nk.replace(a, b)
            ck[nk] = ck.pop(k)
    return ck


def main():
    os.makedirs(OUT, exist_ok=True)
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"[uni] device={device}")
    t0 = time.time()
    det = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100)
    ck = torch.load(CKP, map_location="cpu", weights_only=True)
    remap_uni_keys(ck)
    det = det.to(device).eval()
    msg = det.load_state_dict(ck, strict=False)
    print(f"[uni] loaded (missing={len(msg.missing_keys)} unexpected={len(msg.unexpected_keys)}) in {time.time()-t0:.1f}s")
    if msg.missing_keys[:3]: print("  missing sample:", msg.missing_keys[:3])

    t1 = time.time()
    with torch.no_grad():
        out = det([IMG])
    props = out[0]["bboxes"].float().cpu().numpy()
    scores = out[0].get("scores", None)
    if scores is not None: scores = scores.float().cpu().numpy()
    print(f"[uni] 100 proposals in {time.time()-t1:.1f}s; shape={props.shape}; score range={scores.min() if scores is not None else 'N/A'}..{scores.max() if scores is not None else 'N/A'}")

    # visualize all 100 proposals
    img = Image.open(IMG).convert("RGB")
    overlay = img.copy()
    d = ImageDraw.Draw(overlay)
    for i, (b, s) in enumerate(zip(props.tolist(), scores.tolist() if scores is not None else [None]*len(props))):
        d.rectangle(((b[0], b[1]), (b[2], b[3])), outline="#00C7BE", width=2)
    overlay.save(os.path.join(OUT, "uni_proposals.png"))
    print(f"[vis] all 100 proposals -> {OUT}/uni_proposals.png")

    # top 10 visualization
    order = scores.argsort()[::-1] if scores is not None else range(len(props))
    top10 = order[:10]
    print("[uni] top 10 proposals:")
    for i, idx in enumerate(top10):
        b = props[idx].tolist()
        s = scores[idx] if scores is not None else None
        area = (b[2]-b[0]) * (b[3]-b[1])
        print(f"  #{i+1}: box=[{b[0]:.0f},{b[1]:.0f},{b[2]:.0f},{b[3]:.0f}] area={area:.0f}{' score='+(f'{s:.3f}' if s is not None else '')}")

    # Save proposals.txt
    out_path = os.path.join(OUT, "uni_proposals.txt")
    with open(out_path, "w") as f:
        f.write(f"image: {IMG}\n")
        f.write(f"proposals: {len(props)} (xyxy in original image coordinates, image size = {img.size})\n")
        if scores is not None: f.write("--sorted by score (desc)--\n")
        if scores is not None:
            for i in order:
                b = props[i].tolist()
                f.write(f"  {i:3d}  box=[{b[0]:6.1f},{b[1]:6.1f},{b[2]:6.1f},{b[3]:6.1f}]  score={scores[i]:.4f}  area={int((b[2]-b[0])*(b[3]-b[1]))}\n")
        else:
            for i, b in enumerate(props.tolist()):
                f.write(f"  {i:3d}  box=[{b[0]:6.1f},{b[1]:6.1f},{b[2]:6.1f},{b[3]:6.1f}]  area={int((b[2]-b[0])*(b[3]-b[1]))}\n")
    print(f"[done] saved to {out_path}")


if __name__ == "__main__":
    main()