#!/usr/bin/env python
"""WeDetect-Ref-2B inference runner (macOS MPS/CPU adapted from official infer_wedetect_ref.py).

Pipeline:
  1. WeDetect-Base-Uni (YOLO-World style) -> 100 proposals
  2. WeDetect-Ref-2B (Qwen3-VL grounding) -> score each <object> token in one forward pass
  3. top-1 (or thresholded) bbox -> visualization
"""
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",
    "/System/Library/Fonts/Supplemental/Arial Unicode.ttf",
]


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


def draw_results(image, boxes, scores, query, out_path, proposals=None, proposal_path=None):
    """Draw final prediction (box+score+label) and optionally all proposals."""
    img = image.copy()
    draw = ImageDraw.Draw(img)
    font = load_font()
    colors = ["#FF3B30", "#FF9500", "#34C759", "#007AFF", "#AF52DE", "#FF2D55"]
    if proposals is not None and proposal_path is not None:
        pimg = image.copy()
        pdraw = ImageDraw.Draw(pimg)
        for b in proposals:
            pdraw.rectangle(((b[0], b[1]), (b[2], b[3])), outline="#00C7BE", width=2)
        pimg.save(proposal_path)
        print(f"[vis] proposals ({len(proposals)}) -> {proposal_path}")
    for i, (b, s) in enumerate(zip(boxes, scores)):
        c = colors[i % len(colors)]
        x1, y1, x2, y2 = b
        draw.rectangle(((x1, y1), (x2, y2)), outline=c, width=5)
        label = f"{query} {float(s):.2f}"
        tb = draw.textbbox((0, 0), label, font=font)
        tw, th = tb[2] - tb[0], tb[3] - tb[1]
        ly = y1 - th - 10 if y1 - th - 10 > 0 else y1 + 2
        draw.rectangle(((x1, ly), (x1 + tw + 10, ly + th + 8)), fill=c)
        draw.text((x1 + 5, ly + 3), label, fill="white", font=font)
    img.save(out_path)
    print(f"[vis] prediction -> {out_path}")


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


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--wedetect_ref_checkpoint", type=str, default="../models/wedetect-ref-2b")
    parser.add_argument("--wedetect_uni_checkpoint", type=str, default="../models/wedetect-base-uni/wedetect_base_uni.pth")
    parser.add_argument("--image", type=str, default="../assets/demo.jpeg")
    parser.add_argument("--query", type=str, default="a photo of trees and a river")
    parser.add_argument("--score_thre", type=float, default=-1.0)
    parser.add_argument("--device", type=str, default="auto", choices=["auto", "mps", "cpu"])
    parser.add_argument("--dtype", type=str, default="auto", choices=["auto", "bf16", "fp32"])
    parser.add_argument("--out_dir", type=str, default="../outputs")
    args = parser.parse_args()

    device = args.device
    if device == "auto":
        device = "mps" if torch.backends.mps.is_available() else "cpu"
    if args.dtype == "auto":
        dtype = torch.bfloat16 if device == "mps" else torch.float32
    elif args.dtype == "bf16":
        dtype = torch.bfloat16
    else:
        dtype = torch.float32
    print(f"[env] device={device} dtype={dtype}")
    os.makedirs(args.out_dir, exist_ok=True)
    stem = os.path.splitext(os.path.basename(args.image))[0]

    # ---------- 1. proposal detector (WeDetect-Base-Uni) ----------
    t0 = time.time()
    det_model = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100)
    checkpoint = torch.load(args.wedetect_uni_checkpoint, map_location="cpu", weights_only=True)
    checkpoint = remap_uni_keys(checkpoint)
    det_model = det_model.to(device).eval()
    msg = det_model.load_state_dict(checkpoint, strict=False)
    n_missing = len(msg.missing_keys)
    n_unexpected = len(msg.unexpected_keys)
    print(f"[uni] loaded, missing={n_missing} unexpected={n_unexpected} ({time.time()-t0:.1f}s)")
    if n_missing > 0:
        print("  missing sample:", msg.missing_keys[:5])

    with torch.no_grad():
        det_out = det_model([args.image])
    proposals_np = det_out[0]["bboxes"].float().cpu().numpy()
    print(f"[uni] proposals: {proposals_np.shape}, took {time.time()-t0:.1f}s")

    # ---------- 2. WeDetect-Ref-2B ----------
    t1 = time.time()
    model_kwargs = dict(torch_dtype=dtype, attn_implementation="sdpa")
    model = Qwen3VLGroundingForConditionalGeneration.from_pretrained(args.wedetect_ref_checkpoint, **model_kwargs)
    processor = AutoProcessor.from_pretrained(args.wedetect_ref_checkpoint)
    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] model loaded on {device} ({time.time()-t1:.1f}s), object_token_id={object_token_index}")

    image = Image.open(args.image).convert("RGB")
    ori_shape = [image.size]  # [(W, H)]
    proposals = [torch.tensor(proposals_np.tolist()).to(device).to(model.dtype)]
    num_proposals = len(proposals[0])
    proposal_str = "<object>" * num_proposals

    messages = [
        {"role": "user", "content": [{"type": "text", "text": 'Please detect the "%s" in the image' % args.query}]},
        {"role": "assistant", "content": [{"type": "text", "text": proposal_str}]},
    ]
    messages[0]["content"] = [{"type": "image", "image": copy.deepcopy(image)}] + messages[0]["content"]
    image_inputs, video_inputs = process_vision_info(messages, image_patch_size=16)
    texts = [processor.apply_chat_template(messages, 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)

    t2 = time.time()
    with torch.inference_mode():
        pred = model(
            **model_inputs,
            bboxes=copy.deepcopy(proposals),
            ori_shapes=ori_shape,
            bboxes_id=object_token_index,
        )
    print(f"[ref] forward pass took {time.time()-t2:.1f}s, seq_len={model_inputs['input_ids'].shape}")

    proposal_positions = model_inputs["input_ids"] == object_token_index
    pred_scores_all = pred.logits.sigmoid()[proposal_positions].view(-1)
    pred_bboxes = proposals[0].clone().float()

    if args.score_thre < 0:
        topk_values, topk_indexes = torch.topk(pred_scores_all.view(-1), 1, dim=0)
        pred_bboxes = pred_bboxes[topk_indexes]
        pred_scores = topk_values
    else:
        mask = pred_scores_all > args.score_thre
        pred_bboxes = pred_bboxes[mask]
        pred_scores = pred_scores_all[mask]

    boxes_xyxy = pred_bboxes.cpu().tolist()
    scores = pred_scores.cpu().tolist()
    print(f"[result] query={args.query!r} -> boxes={boxes_xyxy} scores={[round(float(s),4) for s in scores]}")

    # top-5 proposal ranking for verification insight
    k = min(5, num_proposals)
    tv, ti = torch.topk(pred_scores_all.view(-1), k, dim=0)
    print("[result] top-%d proposal scores:" % k, [round(float(v), 4) for v in tv])
    print("[result] top-%d proposal boxes :" % k, [ [round(x,1) for x in proposals[0][i].float().cpu().tolist()] for i in ti.cpu().tolist()])

    # ---------- 3. visualization ----------
    draw_results(
        image, boxes_xyxy, scores, args.query,
        os.path.join(args.out_dir, f"pred_{stem}.png"),
        proposals=proposals_np.tolist(),
        proposal_path=os.path.join(args.out_dir, f"proposals_{stem}.png"),
    )
    with open(os.path.join(args.out_dir, f"result_{stem}.json"), "w") as f:
        json.dump({
            "image": os.path.abspath(args.image), "query": args.query,
            "boxes_xyxy": boxes_xyxy, "scores": [float(s) for s in scores],
            "num_proposals": num_proposals,
            "top5_scores": [float(v) for v in tv],
            "device": device, "dtype": str(dtype),
        }, f, ensure_ascii=False, indent=2)
    print("[done] all outputs saved to", os.path.abspath(args.out_dir))


if __name__ == "__main__":
    main()