"""WeDetect-Ref inference adapted for CPU/MPS (Apple Silicon).

Two-stage pipeline:
  1. WeDetect-Uni (YOLO-World style) generates up to 100 object proposals.
  2. WeDetect-Ref (Qwen3-VL based) classifies proposals in a single forward pass.

Usage:
  python infer_ref_mac.py --ref_ckpt <path-to-WeDetect-Ref-4B> \
      --uni_ckpt <path-to-wedetect_base_uni.pth> \
      --image assets/demo.jpeg --query "the bed" --device cpu --out pred.png
"""
import argparse
import copy
import time

import torch
from PIL import Image
from transformers import AutoProcessor

from wedetect_ref.models.qwen3vl_referring import Qwen3VLGroundingForConditionalGeneration
from wedetect_ref.models.vision_process import process_vision_info
from generate_proposal import SimpleYOLOWorldDetector
from vis import plot_bounding_boxes


def remap_uni_checkpoint(checkpoint):
    """Same key remapping as the official infer_wedetect_ref.py."""
    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('--ref_ckpt', type=str, required=True)
    parser.add_argument('--uni_ckpt', type=str, required=True)
    parser.add_argument('--image', type=str, required=True)
    parser.add_argument('--query', type=str, required=True)
    parser.add_argument('--device', type=str, default='cpu', choices=['cpu', 'mps'])
    parser.add_argument('--score_thre', type=float, default=-1.0)
    parser.add_argument('--out', type=str, default='pred.png')
    args = parser.parse_args()

    device = torch.device(args.device)
    torch.set_grad_enabled(False)

    # ---------------- stage 1: proposals from WeDetect-Uni ----------------
    t0 = time.time()
    model_size = 'base' if 'base' in args.uni_ckpt else 'large'
    det_model = SimpleYOLOWorldDetector(
        backbone_size=model_size, prompt_dim=768, num_prompts=256, num_proposals=100)
    checkpoint = torch.load(args.uni_ckpt, map_location='cpu')
    checkpoint = remap_uni_checkpoint(checkpoint)
    msg = det_model.load_state_dict(checkpoint, strict=False)
    print(f'[uni] load_state_dict: {msg}', flush=True)
    det_model = det_model.float().to(device).eval()

    outputs = det_model([args.image])
    boxes = outputs[0]['bboxes'].float().cpu()
    print(f'[uni] {len(boxes)} proposals in {time.time() - t0:.1f}s', flush=True)

    # ---------------- stage 2: WeDetect-Ref classification ----------------
    t0 = time.time()
    model = Qwen3VLGroundingForConditionalGeneration.from_pretrained(
        args.ref_ckpt,
        torch_dtype=torch.float16,
        attn_implementation="sdpa",
    )
    processor = AutoProcessor.from_pretrained(args.ref_ckpt)
    object_token_index = processor.tokenizer.convert_tokens_to_ids("<object>")
    model.model.object_token_id = object_token_index
    model = model.to(device).eval()
    # RoI glue modules run on CPU fp32 (see NOTE in qwen3vl_referring.py)
    for name in ['first_scale_conv1', 'first_scale_norm', 'first_scale_conv2',
                 'second_scale_conv', 'merge', 'image_pos_projector',
                 'object_vision_projector', 'object_pos_projector']:
        getattr(model.model, name).to('cpu', torch.float32)
    print(f'[ref] model loaded on {args.device} in {time.time() - t0:.1f}s', flush=True)

    image = Image.open(args.image).convert('RGB')
    ori_shape = [image.size]
    proposals = [boxes.to(device).to(model.dtype)]

    num_proposals = len(boxes)
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": copy.deepcopy(image)},
                {"type": "text", "text": 'Please detect the "%s" in the image' % args.query},
            ],
        },
        {
            "role": "assistant",
            "content": [{"type": "text", "text": "<object>" * num_proposals}],
        },
    ]
    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(device)

    t0 = time.time()
    with torch.inference_mode():
        pred = model(
            **model_inputs,
            bboxes=copy.deepcopy(proposals),
            ori_shapes=ori_shape,
            bboxes_id=object_token_index,
            image_inputs=image_inputs,
        )
    print(f'[ref] forward pass in {time.time() - t0:.1f}s', flush=True)

    proposal_positions = model_inputs['input_ids'] == object_token_index
    pred_scores = pred.logits.sigmoid()[proposal_positions].view(-1).float().cpu()
    pred_bboxes = boxes.clone()

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

    print(f'[result] query="{args.query}" -> {len(pred_bboxes)} box(es)')
    for b, s in zip(pred_bboxes.tolist(), pred_scores.tolist()):
        print(f'  score={s:.4f}  box=[{b[0]:.1f}, {b[1]:.1f}, {b[2]:.1f}, {b[3]:.1f}]')

    pred_image = plot_bounding_boxes(copy.deepcopy(image), pred_bboxes.tolist())
    pred_image.save(args.out)
    print(f'[vis] saved to {args.out}')


if __name__ == '__main__':
    main()