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


def ground_only_positive(model, tokenizer, processor, instruction, image_bytes, point):
 
    # assert os.path.exists(image_path) and os.path.isfile(image_path), "Invalid input image path."
    # image = Image.open(image_path)
    
    from PIL import Image
    import io
    ############################################
    image_file = io.BytesIO(image_bytes)
    image = Image.open(image_file)
    ############################################
 
    width, height = image.size
    point = [int(point[0] * width), int(point[1] * height)]
    # width, height = image.size
    image = draw_annotations(image, point, None, output_path=None, size=height/1000 * 1.2)
    
    
    ############################################
    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": full_prompt},
            ],
        }
    ]
    # 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")
    ####################################
 
    # print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 3.5 ---------------------------------"))
    # ####################################
    # import pdb; pdb.set_trace()
    # ####################################


    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 
        ####################################
    )
    
    # print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 4 ---------------------------------"))

    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}")
    ########################################################################
    
    return token_str, max_prob


model_name_or_path = "/mnt/beegfs/Workspace/Models/GUI-Actor-Verifier-2B"
# model_name_or_path = "/mnt/task_runtime/sh_workspace/new_temp/Models/GUI-Actor-Verifier-2B"



# def load_verify(model_name_or_path):
#     model = Qwen2VLForConditionalGeneration.from_pretrained(
#                 model_name_or_path, 
#                 device_map="cuda:0", 
#                 trust_remote_code=True, 
#                 torch_dtype=torch.bfloat16,
#                 attn_implementation="flash_attention_2"
#             ).eval()
#     output_len = 1
#     return model



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

    # 允许更快的 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,
    )
    # 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






# instruction = 'close this window'
# point = [0.97,0.15]
# image = '/scratch/Renee/Dataset/GUI-Actor-Data/AndroidControl/tfrecord/images/android_control_episode_[11]_9.png'
# _, max_prob = ground_only_positive(model, tokenizer, processor, instruction, image, point) # output True or False

# max_prob