#!/usr/bin/env python
"""Smoke test: random-weight full pipeline on MPS/CPU. Validates code paths before real weights arrive."""
import copy
import os
import sys
import time

import torch
from PIL import Image

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, AutoConfig
from generate_proposal import SimpleYOLOWorldDetector

MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "models", "wedetect-ref-2b")
IMG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "assets", "demo.jpeg")

device = "mps" if torch.backends.mps.is_available() else "cpu"
print("[smoke] device:", device)

# 1. detector with random weights (code-path validation only)
det = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100).to(device).eval()
with torch.no_grad():
    out = det([IMG])
props = out[0]["bboxes"].float().cpu().numpy()
print("[smoke] detector OK, proposals:", props.shape, "sample:", props[0].round(1).tolist())

# 2. qwen grounding model with random weights
cfg = AutoConfig.from_pretrained(MODEL_DIR)
t0 = time.time()
model = Qwen3VLGroundingForConditionalGeneration(cfg)
model = model.to(torch.bfloat16).to(device).eval()
print(f"[smoke] model built ({time.time()-t0:.0f}s), params={sum(p.numel() for p in model.parameters())/1e9:.2f}B")

processor = AutoProcessor.from_pretrained(MODEL_DIR)
object_token_index = processor.tokenizer.convert_tokens_to_ids("<object>")
model.model.object_token_id = object_token_index

image = Image.open(IMG).convert("RGB")
proposals = [torch.tensor(props.tolist()).to(device).to(model.dtype)]
proposal_str = "<object>" * len(proposals[0])
messages = [
    {"role": "user", "content": [{"type": "text", "text": 'Please detect the "trees and a river" in the image'}]},
    {"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)
print("[smoke] inputs ready, seq_len:", model_inputs["input_ids"].shape)

t1 = time.time()
with torch.inference_mode():
    pred = model(**model_inputs, bboxes=copy.deepcopy(proposals), ori_shapes=[image.size], bboxes_id=object_token_index)
print(f"[smoke] forward OK ({time.time()-t1:.1f}s), logits shape: {tuple(pred.logits.shape)}")
pp = model_inputs["input_ids"] == object_token_index
print("[smoke] object tokens in input:", int(pp.sum()), "== proposals:", len(proposals[0]))
scores = pred.logits.sigmoid()[pp].view(-1)
print("[smoke] random scores mean/std:", float(scores.mean()), float(scores.std()))
print("[smoke] ALL PATHS VERIFIED")