#!/usr/bin/env python3
"""Run GUI grounding with a local directory or Hugging Face model ID."""

import argparse
import re

import torch
from PIL import Image
from transformers import AutoModelForCausalLM, AutoProcessor


POINT_PATTERN = re.compile(r"<loc_(\d+)>,<loc_(\d+)>")


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=".")
    parser.add_argument("--image", required=True)
    parser.add_argument("--prompt", required=True)
    return parser.parse_args()


def main():
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("This example requires a CUDA GPU")

    model = AutoModelForCausalLM.from_pretrained(
        args.model,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        attn_implementation="sdpa",
    ).cuda().eval()
    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
    image = Image.open(args.image).convert("RGB")
    inputs = processor(images=image, text=args.prompt, return_tensors="pt").to(
        "cuda", dtype=torch.bfloat16
    )

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            do_sample=False,
            max_new_tokens=16,
        )
    text = processor.tokenizer.batch_decode(
        output_ids, skip_special_tokens=False
    )[0]
    match = POINT_PATTERN.search(text)
    if match is None:
        print(text)
        raise RuntimeError("the model output did not contain location tokens")

    normalized = tuple(map(int, match.groups()))
    pixels = (
        normalized[0] / 999 * image.width,
        normalized[1] / 999 * image.height,
    )
    print(f"raw_output={text}")
    print(f"normalized_point={normalized}")
    print(f"pixel_point=({pixels[0]:.2f}, {pixels[1]:.2f})")


if __name__ == "__main__":
    main()