# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from collections import defaultdict
from typing import Any, Callable, Dict, Tuple, TypedDict

import torch
from transformers import PreTrainedTokenizer

from ...protocol import DataProto
from ...utils.reward_score import math_compute_score, r1v_compute_score, r1gui_compute_score


##################################################################################
from .verify import ground_only_positive,load_verify,load_tokenizer,load_processor
##################################################################################


import torch
import numpy as np
from PIL import Image



##################################################################################
import re
from typing import List, Tuple, Union

Number = Union[int, float]

def extract_points(text: str) -> List[Tuple[Number, Number]]:
    # 匹配 'point': [x, y] 或 "point": [x, y]
    num = r"[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?"
    pattern = re.compile(rf"""['"]point['"]\s*:\s*\[\s*({num})\s*,\s*({num})\s*\]""")
    pts: List[Tuple[Number, Number]] = []

    for x, y in pattern.findall(text):
        # 尽量转成 int，否则用 float
        def to_num(s: str) -> Number:
            return int(float(s)) if re.fullmatch(r"[+-]?\d+", s) else float(s)
        pts.append([to_num(x), to_num(y)])
    return pts
##################################################################################



class RewardScore(TypedDict):
    overall: float
    format: float
    accuracy: float


def point_in_bbox(point, bbox):
    if len(point) != 2 or len(bbox) != 4:
        raise ValueError("point must be [x, y] and bbox must be [x_min, y_min, x_max, y_max]")
    x, y = point
    x_min, y_min, x_max, y_max = bbox

    return (x_min <= x <= x_max) and (y_min <= y <= y_max)



class CustomRewardManager:
    def __init__(self, tokenizer: PreTrainedTokenizer,compute_score: str):
    # def __init__(self, tokenizer: PreTrainedTokenizer,verify_model,verify_tokenizer,verify_processor, compute_score: str):
        self.tokenizer = tokenizer
        if compute_score == "math":
            self.compute_score: Callable[[str, str], RewardScore] = math_compute_score
        elif compute_score == "r1v":
            self.compute_score: Callable[[str, str], RewardScore] = r1v_compute_score
        elif compute_score == "r1gui":
            self.compute_score: Callable[[str, str], RewardScore] = r1gui_compute_score
        else:
            raise NotImplementedError()
        
        
        ####################################################################################
        model_name_or_path = "/mnt/beegfs/Workspace/Models/GUI-Actor-Verifier-2B"
        self.verify_model = load_verify(model_name_or_path)
        self.tokenizer = load_tokenizer(model_name_or_path)
        self.processor = load_processor(model_name_or_path)
      
        
        # self.verify_model = verify_model
        # self.tokenizer = verify_tokenizer
        # self.processor = verify_processor
        ####################################################################################
        
        

    def __call__(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, Any]]:
        reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
        reward_metrics = defaultdict(list)
        
        
        for i in range(len(data)):
            data_item = data[i]  # DataProtoItem
            response_ids = data_item.batch["responses"]
            response_mask = data_item.batch["response_mask"]
            valid_response_length = response_mask.sum()
            valid_response_ids = response_ids[:valid_response_length]

            response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)
            ground_truth = data_item.non_tensor_batch["ground_truth"]
            
            ####################################################################################
            score_rule = self.compute_score(response_str, ground_truth)
            ####################################################################################
            from datetime import datetime
            # print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 1 ---------------------------------"))
            
            
            # data_item.non_tensor_batch['id']
            # data_item.non_tensor_batch['gt_action']
            # data_item.non_tensor_batch['gt_input_text']
            # data_item.non_tensor_batch['instruction']
            # data_item.non_tensor_batch['image']['bytes']
            # data_item.non_tensor_batch['history']
            # data_item.non_tensor_batch['ground_truth']
            # data_item.non_tensor_batch['multi_modal_inputs'] 
            # data_item.non_tensor_batch['multi_modal_inputs']['pixel_values']
            # data_item.non_tensor_batch['gt_bbox_ori']
            
            temp = data_item.non_tensor_batch
            ####################################################################################
            import base64
            image_bytes = data_item.non_tensor_batch['image']['bytes']
            # with open("output.png", "wb") as f:
            #     f.write(image_bytes)

            # print(response_str)
            # print(ground_truth)
            # print(data_item.non_tensor_batch['instruction'])
            ####################################################################################
            from PIL import Image
            import io
            ############################################
            image_file = io.BytesIO(image_bytes)
            image = Image.open(image_file)
            width, height = image.size
            instruction= data_item.non_tensor_batch['instruction']
            
            
            ############################################
            temp = extract_points(response_str)
            # print('Pred : ', response_str)
            
       
            if len(temp) < 1:
                temp.append([-0.1,-0.1])
                    
            point = temp[0]
            
            is_int_list = (
                isinstance(point, list)                     # 是列表
                and len(point) == 2                         # 长度为2（可选）
                and all(isinstance(x, int) for x in point)  # 所有元素都是 int
            )

            if is_int_list:
                point[0] = point[0] / width
                point[1] = point[1] / height
                
            ############################################
            
            # print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 2 ---------------------------------"))
            
            ###################################################################################
            _, probs = ground_only_positive(self.verify_model, self.tokenizer, self.processor, instruction, image_bytes, point)
            
            
            
            import json
            gt_dict = json.loads(ground_truth)
            print('Pred Point: ', point)
            print('GT : ', gt_dict['gt_bbox_ori'])
            
            if len(gt_dict['gt_bbox_ori']) > 2:
                if point_in_bbox(point, gt_dict['gt_bbox_ori']):
                    print("✅ The point is inside the bbox.")
                else:
                    print("❌ The point is outside the bbox.")
                    
            print('SSL Reward : ', probs)
            print('Rule Reward : ', score_rule)
            print('------------------------------------------')
            
            # 统一把 probs 变成 Python float
            if torch.is_tensor(probs):
                probs = probs.detach().float().cpu()
                probs = probs.mean().item() if probs.numel() > 1 else probs.item()
            else:
                probs = float(probs)
            ####################################################################################
            
            score = {
                "overall": 0.8 * probs + 0.2 * 1,
                "format": 1,
                "accuracy": probs,
            }
            
            # print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 5 ---------------------------------"))
            ############################################
            # import pdb; pdb.set_trace()
                    
            
            reward_tensor[i, valid_response_length - 1] = score["overall"]
            for key, value in score.items():
                reward_metrics[key].append(value)

        return reward_tensor, reward_metrics






# class CustomRewardManager:
#     def __init__(self, tokenizer: PreTrainedTokenizer, compute_score: str):
#         self.tokenizer = tokenizer
#         if compute_score == "math":
#             self.compute_score: Callable[[str, str], RewardScore] = math_compute_score
#         elif compute_score == "r1v":
#             self.compute_score: Callable[[str, str], RewardScore] = r1v_compute_score
#         elif compute_score == "r1gui":
#             self.compute_score: Callable[[str, str], RewardScore] = r1gui_compute_score
#         else:
#             raise NotImplementedError()
        
        
#         ####################################################################################
#         # model_name_or_path = "/scratch/Renee/Model/GUI-Actor-Verifier-2B"
#         # self.verify_model = load_verify(model_name_or_path)
#         # self.tokenizer = load_tokenizer(model_name_or_path)
#         # self.processor = load_processor(model_name_or_path)
#         # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#         # self.verify_model = self.verify_model.to(device)
#         # self.verify_model.eval() 
#         ####################################################################################
        
        

#     def __call__(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, Any]]:
#         reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
#         reward_metrics = defaultdict(list)
        
        
#         for i in range(len(data)):
#             data_item = data[i]  # DataProtoItem
#             response_ids = data_item.batch["responses"]
#             response_mask = data_item.batch["response_mask"]
#             valid_response_length = response_mask.sum()
#             valid_response_ids = response_ids[:valid_response_length]

#             response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)
#             ground_truth = data_item.non_tensor_batch["ground_truth"]
            
#             ####################################################################################
#             # score = self.compute_score(response_str, ground_truth)
#             ####################################################################################
#             from datetime import datetime
#             print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 1 ---------------------------------"))
            
            
#             # data_item.non_tensor_batch['id']
#             # data_item.non_tensor_batch['gt_action']
#             # data_item.non_tensor_batch['gt_input_text']
#             # data_item.non_tensor_batch['instruction']
#             # data_item.non_tensor_batch['image']['bytes']
#             # data_item.non_tensor_batch['history']
#             # data_item.non_tensor_batch['ground_truth']
#             # data_item.non_tensor_batch['multi_modal_inputs'] 
#             # data_item.non_tensor_batch['multi_modal_inputs']['pixel_values']
#             # data_item.non_tensor_batch['gt_bbox_ori']
            
#             temp = data_item.non_tensor_batch
#             ####################################################################################
#             import base64
#             image_bytes = data_item.non_tensor_batch['image']['bytes']
#             # with open("output.png", "wb") as f:
#             #     f.write(image_bytes)

#             # print(response_str)
#             # print(ground_truth)
#             # print(data_item.non_tensor_batch['instruction'])
#             ####################################################################################
#             from PIL import Image
#             import io
#             ############################################
#             image_file = io.BytesIO(image_bytes)
#             image = Image.open(image_file)
#             instruction= data_item.non_tensor_batch['instruction']
#             point = extract_points(response_str)[0]
            
#             print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 2 ---------------------------------"))
            
#             ###################################################################################
#             _, probs = ground_only_positive(self.verify_model, self.tokenizer, self.processor, instruction, image_bytes, point)
#             ####################################################################################
#             score = {
#                 "overall": 0.8 * probs + 0.2 * 1,
#                 "format": 1,
#                 "accuracy": probs,
#             }
            
#             print(datetime.now().strftime("[%Y-%m-%d %H:%M:%S] ------------------------------- 5 ---------------------------------"))
#             ############################################
#             # import pdb; pdb.set_trace()
                    
            
#             reward_tensor[i, valid_response_length - 1] = score["overall"]
#             for key, value in score.items():
#                 reward_metrics[key].append(value)

#         return reward_tensor, reward_metrics





# (Runner pid=3842623) TensorDict(
# (Runner pid=3842623)     fields={
# (Runner pid=3842623)         attention_mask: Tensor(shape=torch.Size([2560]), device=cpu, dtype=torch.int64, is_shared=False),
# (Runner pid=3842623)         input_ids: Tensor(shape=torch.Size([2560]), device=cpu, dtype=torch.int64, is_shared=False),
# (Runner pid=3842623)         position_ids: Tensor(shape=torch.Size([3, 2560]), device=cpu, dtype=torch.int64, is_shared=False),
# (Runner pid=3842623)         prompts: Tensor(shape=torch.Size([2048]), device=cpu, dtype=torch.int64, is_shared=False),
# (Runner pid=3842623)         response_mask: Tensor(shape=torch.Size([512]), device=cpu, dtype=torch.int64, is_shared=False),
# (Runner pid=3842623)         responses: Tensor(shape=torch.Size([512]), device=cpu, dtype=torch.int64, is_shared=False)},
# (Runner pid=3842623)     batch_size=torch.Size([]),
# (Runner pid=3842623)     device=None,
# (Runner pid=3842623)     is_shared=False)