# 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 ask_llm,load_verify,load_tokenizer,load_processor
##################################################################################


import torch
import numpy as np
from PIL import Image




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



import base64
from PIL import Image
import io
import json
import numpy as np


####################################################################
def bbox_to_center(bbox):
    bbox = np.array(bbox, dtype=float)
    if bbox.ndim == 1:
        bbox = bbox[None, :]  # (1, 4)

    if np.all(bbox[:, 2] > bbox[:, 0]) and np.all(bbox[:, 3] > bbox[:, 1]):
        # 形如 [x_min, y_min, x_max, y_max]
        cx = (bbox[:, 0] + bbox[:, 2]) / 2
        cy = (bbox[:, 1] + bbox[:, 3]) / 2
    else:
        # 形如 [x, y, w, h]
        cx = bbox[:, 0] + bbox[:, 2] / 2
        cy = bbox[:, 1] + bbox[:, 3] / 2

    centers = np.stack([cx, cy], axis=1)
    return centers[0].tolist()

def fix_bbox_shape(lst):
    fixed = []
    for item in lst:
        if isinstance(item, (list, tuple)):
            if len(item) == 2:
                # 把 [x, y] 复制成 [x, y, x, y]
                x, y = item
                fixed.append([x, y, x, y])
            elif len(item) == 4:
                fixed.append(item)
            else:
                print(f"⚠️ 跳过长度异常项: {item}")
        else:
            print(f"⚠️ 非list/tuple项: {item}")
    return fixed
####################################################################


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 = "/mnt/beegfs/Workspace/Models/GUI-Actor-Verifier-2B"
        model_name_or_path = "/mdr0/quantaalpha/wubiao/Models/Qwen2.5-VL-3B-Instruct"
        # model_name_or_path = "/home/omnisky/data2/data/xk/Models/Qwen2.5-VL-3B-Instruct"
        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)
        ####################################################################################

    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)
        
        ############################################
        all_info_list_Pred_Point = []
        all_info_list_GT_BBox_Ori = []
        all_info_list_GT_Point = []
        all_info_list_Score = []
        
        # print(len(data)) # 2560 / 4*5*128
        ############################################
        
        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)
            # print('------------------------------------------------------')
            # print(response_str)
            # print(ground_truth)
            # print('------------------------------------------------------')
            ####################################################################################
            
            ####################################################################################
            image_bytes = data_item.non_tensor_batch['image']['bytes']
            ############################################
            image_file = io.BytesIO(image_bytes)
            image = Image.open(image_file) 
            instruction= data_item.non_tensor_batch['instruction']
            gt_dict = json.loads(ground_truth)
            gt_open_ans = gt_dict['gt_open_ans']
            new_prediction = response_str.split('<answer>')[-1].split('</answer>')[0]
 
            
            # All_instruction_1 = ' Sentence 1: ' + gt_open_ans + ' Sentence 2: ' + new_prediction
            All_instruction_1 = ' Sentence 1: ' + str(gt_dict) + ' Sentence 2: ' + new_prediction
            # prompt_origin = "'{}' ###### Do you think the two sentences have the same meaning? .  Only response Answer True or False."
            prompt_origin = """'{}'
                You are judging whether a predicted UI action matches the ground-truth action
                at the intent level.

                Focus on whether they represent the SAME USER INTENT.

                Only consider:
                - action type
                - input_text meaning

                Ignore:
                - coordinates, bounding boxes, point differences
                - formatting differences
                - enum vs string differences
                - minor numerical deviations

                Guidelines:
                - If action types are the same AND the input_text has the same meaning
                (including equivalent phrases, empty text, or "no input text"),
                answer True.
                - If the user intent is clearly different, answer False.

                Do NOT require exact matching. Be lenient and intent-focused.
                Only output: True or False.
                """

            all_instruction = prompt_origin.format(All_instruction_1)
            
            
            point = [-0.1,-0.1] 
            gt_dict['gt_bbox_ori'] = [-0.1,-0.1] 
            gt_point = [-0.1,-0.1] 
            all_info_list_Pred_Point.append(point)
            all_info_list_GT_BBox_Ori.append(gt_dict['gt_bbox_ori'])
            all_info_list_GT_Point.append(gt_point)
            ####################################################################################
            _, probs = ask_llm(self.verify_model, self.tokenizer, self.processor, all_instruction, image)
            # 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)
            ####################################################################################
            
            
            print('++++++++++++++++++++++++++++++++++++++++++++++++++++++')
            # print('All_Instruction : ', all_instruction) 
            print('GT : ', gt_dict)
            print('Pred: ', new_prediction)
            print('Yes Logits : ', probs)
            print('Rule Score : ', score_rule['overall'])
            print('++++++++++++++++++++++++++++++++++++++++++++++++++++++')
            print()
            
            
            score = {
                "overall": 0.5 * probs + 0.5 * score_rule['overall'] ,
                "format": 1,
                "accuracy": probs,
            }
            
            # all_info['score'] = probs
            all_info_list_Score.append(probs)
            ####################################################################################
            
            
            reward_tensor[i, valid_response_length - 1] = score["overall"]
            for key, value in score.items():
                reward_metrics[key].append(value)

        ############################################
        # all_info_list.append(all_info)
        
        pred_point_tensor = torch.tensor(all_info_list_Pred_Point, dtype=torch.float32)
        fixed_bbox_list = fix_bbox_shape(all_info_list_GT_BBox_Ori)
        gt_bbox_tensor = torch.tensor(fixed_bbox_list, dtype=torch.float32)
        gt_point_tensor = torch.tensor(all_info_list_GT_Point, dtype=torch.float32)
        score_tensor = torch.tensor(all_info_list_Score, dtype=torch.float32)
        
        ############################################
        
        # return reward_tensor, reward_metrics
        return reward_tensor, reward_metrics, pred_point_tensor,gt_bbox_tensor,gt_point_tensor,score_tensor
