import torch
from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
from transformers.generation import GenerationConfig
import json
import re
import os
import numpy as np
from PIL import Image, ImageDraw
from qwen_vl_utils import process_vision_info



def draw_annotations(img, point_in_pixel, bbox, output_path='test.png', color='red', size=1):
    draw = ImageDraw.Draw(img)
    
    # Draw the ground truth bounding box in green
    if bbox:
        # Assuming bbox format is [x1, y1, x2, y2]
        draw.rectangle(bbox, outline="yellow", width=4)
    
    # Draw a small circle around the predicted point in red
    if point_in_pixel:
        # Create a small rectangle around the point (5 pixels in each direction)
        radius = np.ceil(8 * size).astype(int)
        circle_bbox = [
            point_in_pixel[0] - radius,  # x1
            point_in_pixel[1] - radius,  # y1
            point_in_pixel[0] + radius,  # x2
            point_in_pixel[1] + radius   # y2
        ]
        draw.ellipse(circle_bbox, outline=color, width=np.ceil(4 * size).astype(int))
    
    return img


# ---- 关键：在当前作用域导入 torch ----
import torch
import io
import os
import time
import random
import string
from PIL import Image
import io

def ask_llm(model, tokenizer, processor, instruction, image):
    

    # ############################################
    # save_dir = "/mnt/beegfs/Workspace/ICLR_2026/GUI-R1_SSL_V2/annotated_images"
    # os.makedirs(save_dir, exist_ok=True)   
    # timestamp = time.strftime("%Y%m%d_%H%M%S")
    # rand_str = ''.join(random.choices(string.ascii_lowercase + string.digits, k=4))
    # filename = f"annotated_{timestamp}_{rand_str}.png"
    # output_path = os.path.join(save_dir, filename)
    # image.save(output_path)
    # ############################################
    
    # from datetime import datetime
    # # print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 3 ---------------------------------"))        
 
    # prompt_origin = "Please observe the screenshot and exame whether the hollow red circle accurately placed on the intended position in the image: '{}'. Answer True or False."
    # full_prompt = prompt_origin.format(instruction)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": image,
                },
                {"type": "text", "text": instruction},
            ],
        }
    ]
    
    
    # Preparation for inference
    text_input = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text_input],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    
    ####################################
    inputs = inputs.to("cuda:0")
    ####################################

    generated_ids = model.generate(
        **inputs,  
        max_new_tokens=1,
        do_sample=False,
        temperature=0.0,
        ####################################
        return_dict_in_generate=True,  # 必须加
        output_scores=True,              # 必须加
        output_logits=True 
        ####################################
    )
    
    import torch
    ########################################################################
    for step, logits in enumerate(generated_ids.scores):
        probs = torch.softmax(logits, dim=-1)
        max_prob, max_token_id = torch.max(probs, dim=-1)
        token_str = tokenizer.decode([max_token_id.item()])
        print(f"Step {step}: max prob = {max_prob.item():.4f}, token = {token_str}")

        if 'False' in token_str:
            max_prob = max_prob*-1
    ########################################################################
    
    return token_str, max_prob


def load_verify(model_name_or_path):
    import torch
    from transformers import Qwen2VLForConditionalGeneration,Qwen2_5_VLForConditionalGeneration

    # 允许更快的 matmul
    if torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = True

    # 检测 flash-attn 是否可用（已安装且有 CUDA）
    has_cuda = torch.cuda.is_available()
    try:
        import flash_attn  # type: ignore
        has_fa2 = has_cuda  # 只要有 CUDA 且包存在，就认为可用
    except Exception:
        has_fa2 = False

    attn_impl = "flash_attention_2" if has_fa2 else "sdpa"

    # dtype 选择
    use_bf16 = has_cuda and torch.cuda.is_bf16_supported()
    dtype = torch.bfloat16 if use_bf16 else (torch.float16 if has_cuda else torch.float32)

    # # **关键**：device_map="auto" 让模型自动上 GPU（可多卡分片）
    # model = Qwen2VLForConditionalGeneration.from_pretrained(
    #     model_name_or_path,
    #     device_map="auto" if has_cuda else None,
    #     torch_dtype=dtype,
    #     attn_implementation=attn_impl,   # 若没有 FA2，就用 SDPA，避免你遇到的报错
    #     low_cpu_mem_usage=True,
    # )
    
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        model_name_or_path,
        torch_dtype=torch.bfloat16,
        # attn_implementation="flash_attention_2",
        attn_implementation=attn_impl,   # 若没有 FA2，就用 SDPA，避免你遇到的报错
        device_map="auto" if has_cuda else None,
    )
    
    # output_len = 1\
    
    # ####################################
    # import pdb; pdb.set_trace()
    # ####################################

    return model


def load_tokenizer(model_name_or_path):
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
    return tokenizer
    
def load_processor(model_name_or_path):
    processor = AutoProcessor.from_pretrained(model_name_or_path)
    return processor




