#!/usr/bin/env python
"""Final verify: run multiple WeDetect-Ref-2B queries and save side-by-side panel."""
import argparse
import copy
import json
import os
import sys
import time

import torch
from PIL import Image, ImageDraw, ImageFont

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from wedetect_ref.models.qwen3vl_referring import Qwen3VLGroundingForConditionalGeneration
from wedetect_ref.models.vision_process import process_vision_info
from transformers import AutoProcessor
from generate_proposal import SimpleYOLOWorldDetector

FONT_CANDIDATES = [
    "/System/Library/Fonts/PingFang.ttc",
    "/System/Library/Fonts/Hiragino Sans GB.ttc",
    "/System/Library/Fonts/STHeiti Medium.ttc",
]


def load_font(size):
    for p in FONT_CANDIDATES:
        if os.path.exists(p):
            try: return ImageFont.truetype(p, size)
            except: continue
    return ImageFont.load_default()


def remap_uni_keys(checkpoint):
    keys = list(checkpoint.keys())
    for key in keys:
        if "backbone" in key:
            checkpoint[key.replace("backbone.image_model.model.", "backbone.")] = checkpoint.pop(key)
    keys = list(checkpoint.keys())
    for key in keys:
        if "bbox_head" in key:
            nk = key.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)
            checkpoint[nk] = checkpoint.pop(key)
    return checkpoint


def detect(det_model, image_path, device):
    with torch.no_grad():
        out = det_model([image_path])
    return out[0]["bboxes"].float().cpu().numpy()


def run_ref(model, processor, image, query, proposals_np, device, object_token_index, orig_w, orig_h):
    proposals = [torch.tensor(proposals_np.tolist()).to(device).to(model.dtype)]
    proposal_str = "<object>" * len(proposals[0])
    msgs = [
        {"role":"user","content":[{"type":"text","text":f'Please detect the "{query}" in the image'}]},
        {"role":"assistant","content":[{"type":"text","text":proposal_str}]},
    ]
    msgs[0]["content"] = [{"type":"image","image":copy.deepcopy(image)}] + msgs[0]["content"]
    image_inputs, video_inputs = process_vision_info(msgs, image_patch_size=16)
    texts = [processor.apply_chat_template(msgs, tokenize=False)]
    model_inputs = processor(text=texts, images=image_inputs, videos=video_inputs, return_tensors="pt", padding=True, do_resize=False)
    model_inputs = model_inputs.to(model.device)
    t0 = time.time()
    with torch.inference_mode():
        pred = model(**model_inputs, bboxes=copy.deepcopy(proposals), ori_shapes=[(orig_w, orig_h)], bboxes_id=object_token_index)
    pos = model_inputs["input_ids"] == object_token_index
    scores = pred.logits.sigmoid()[pos].view(-1)
    return scores, time.time() - t0


def draw_panel(items, image, out_path):
    """items: list of (query, box, score)"""
    img = image.copy()
    draw = ImageDraw.Draw(img)
    font = load_font(20)
    colors = ["#FF3B30", "#FF9500", "#34C759", "#007AFF", "#AF52DE"]
    for i, (q, b, s) in enumerate(items):
        c = colors[i % len(colors)]
        x1, y1, x2, y2 = b
        draw.rectangle(((x1, y1), (x2, y2)), outline=c, width=5)
        lbl = f"Q{i+1}: {q}  score={float(s):.2f}"
        tb = draw.textbbox((0,0), lbl, font=font)
        ly = max(0, y1 - (tb[3]-tb[1]) - 8)
        draw.rectangle(((x1, ly), (x1 + (tb[2]-tb[0]) + 10, ly + (tb[3]-tb[1]) + 6)), fill=c)
        draw.text((x1+5, ly+2), lbl, fill="white", font=font)
    img.save(out_path)
    print(f"[panel] {len(items)} queries -> {out_path}")


def draw_proposals(image, props, out_path):
    img = image.copy()
    d = ImageDraw.Draw(img)
    for b in props:
        d.rectangle(((b[0], b[1]), (b[2], b[3])), outline="#00C7BE", width=2)
    img.save(out_path)
    print(f"[proposals] {len(props)} -> {out_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ref_dir", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-ref-2b")
    ap.add_argument("--uni_ckpt", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-base-uni/wedetect_base_uni.pth")
    ap.add_argument("--image", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/assets/demo.jpeg")
    ap.add_argument("--queries", nargs="+", default=[
        "yellow pillows on the bed",
        "the round alarm clock",
        "paintings on the wall",
        "the arc floor lamp",
        "flowers in the vase",
    ])
    ap.add_argument("--out_dir", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/outputs")
    ap.add_argument("--device", default="auto")
    ap.add_argument("--dtype", default="auto")
    args = ap.parse_args()

    device = args.device
    if device == "auto": device = "mps" if torch.backends.mps.is_available() else "cpu"
    dtype = torch.bfloat16 if (args.dtype == "auto" and device == "mps") else torch.float32
    print(f"[env] device={device} dtype={dtype}")
    os.makedirs(args.out_dir, exist_ok=True)

    # sanity: weights complete
    import safetensors.torch
    ref_index = json.load(open(os.path.join(args.ref_dir, "model.safetensors.index.json")))
    expected = set()
    for f, keys in ref_index.get("weight_map", {}).items():
        expected.update([f"/{k}" for k in keys])  # not used, just info
    print(f"[ref] {len(ref_index.get('weight_map', {}))} tensor entries declared across shards")

    # 1) detector
    t0 = time.time()
    det = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100)
    ck = torch.load(args.uni_ckpt, 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")
    image = Image.open(args.image).convert("RGB")
    W, H = image.size
    t1 = time.time()
    props_np = detect(det, args.image, device)
    print(f"[uni] {len(props_np)} proposals in {time.time()-t1:.1f}s")
    draw_proposals(image, props_np.tolist(), os.path.join(args.out_dir, "proposals.png"))

    # 2) ref model
    t2 = time.time()
    model = Qwen3VLGroundingForConditionalGeneration.from_pretrained(args.ref_dir, torch_dtype=dtype, attn_implementation="sdpa")
    processor = AutoProcessor.from_pretrained(args.ref_dir)
    object_token_index = processor.tokenizer.convert_tokens_to_ids("<object>")
    model.model.object_token_id = object_token_index
    model = model.to(device).eval()
    print(f"[ref] loaded in {time.time()-t2:.1f}s, params={sum(p.numel() for p in model.parameters())/1e9:.2f}B")

    results = []
    for q in args.queries:
        t3 = time.time()
        scores, dt = run_ref(model, processor, image, q, props_np, device, object_token_index, W, H)
        top_v, top_i = torch.topk(scores, 1, dim=0)
        box = props_np[top_i.item()].tolist()
        results.append({"query": q, "box_xyxy": box, "score": float(top_v), "forward_s": dt})
        print(f"[ref] {q!r} -> box={[round(x,1) for x in box]} score={float(top_v):.3f} ({dt:.1f}s)")

    # 3) save per-image annotated image & json
    stem = os.path.splitext(os.path.basename(args.image))[0]
    draw_panel([(r["query"], r["box_xyxy"], r["score"]) for r in results], image, os.path.join(args.out_dir, f"results_{stem}.png"))
    with open(os.path.join(args.out_dir, f"results_{stem}.json"), "w") as f:
        json.dump({"image": os.path.abspath(args.image), "device": device, "dtype": str(dtype), "results": results, "n_proposals": len(props_np)}, f, ensure_ascii=False, indent=2)
    print(f"[done] all outputs in {args.out_dir}")


if __name__ == "__main__":
    main()