# 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.

import math
import os
from collections import defaultdict
from io import BytesIO
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch
from datasets import load_dataset
from PIL import Image
from PIL.Image import Image as ImageObject
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer, ProcessorMixin

from ..models.transformers.qwen2_vl import get_rope_index
from . import torch_functional as VF
import json

def collate_fn(features: List[Dict[str, Any]]) -> Dict[str, Any]:
    tensors = defaultdict(list)
    non_tensors = defaultdict(list)
    for feature in features:
        for key, value in feature.items():
            if isinstance(value, torch.Tensor):
                tensors[key].append(value)
            else:
                non_tensors[key].append(value)

    for key, value in tensors.items():
        tensors[key] = torch.stack(value, dim=0)

    for key, value in non_tensors.items():
        non_tensors[key] = np.array(value, dtype=object)

    return {**tensors, **non_tensors}


def process_image(image: Union[Dict[str, Any], ImageObject], max_pixels: int, min_pixels: int) -> ImageObject:
    if isinstance(image, dict):
        image = Image.open(BytesIO(image["bytes"]))
    if isinstance(image, str):
        image = Image.open(image)
    if (image.width * image.height) > max_pixels:
        resize_factor = math.sqrt(max_pixels / (image.width * image.height))
        width, height = int(image.width * resize_factor), int(image.height * resize_factor)
        image = image.resize((width, height))

    if (image.width * image.height) < min_pixels:
        resize_factor = math.sqrt(min_pixels / (image.width * image.height))
        width, height = int(image.width * resize_factor), int(image.height * resize_factor)
        image = image.resize((width, height))

    if image.mode != "RGB":
        image = image.convert("RGB")

    return image



def bbox_center(gt_bbox):
    x1, y1, x2, y2 = gt_bbox
    cx = (x1 + x2) / 2
    cy = (y1 + y2) / 2
    return [cx, cy]



class RLHFDataset(Dataset):
    """
    We assume the dataset contains a column that contains prompts and other information
    """

    def __init__(
        self,
        data_path: str,
        tokenizer: PreTrainedTokenizer,
        processor: Optional[ProcessorMixin],
        prompt_key: str = "prompt",
        answer_key: str = "answer",
        image_key: str = "images",
        max_prompt_length: int = 1024,
        truncation: str = "error",
        system_prompt: str = None,
        max_pixels: int = None,
        min_pixels: int = None,
    ):
        self.tokenizer = tokenizer
        self.processor = processor
        self.prompt_key = prompt_key
        self.answer_key = answer_key
        self.image_key = image_key
        self.max_prompt_length = max_prompt_length
        self.truncation = truncation
        self.system_prompt = system_prompt
        self.max_pixels = max_pixels
        self.min_pixels = min_pixels

        if "@" in data_path:
            data_path, data_split = data_path.split("@")
        else:
            data_split = "train"
            # print(data_path)

        if os.path.isdir(data_path):
            self.dataset = load_dataset("parquet", data_dir=data_path, split="train")
        elif os.path.isfile(data_path):
            self.dataset = load_dataset("parquet", data_files=data_path, split="train")
        else:  # remote dataset
            self.dataset = load_dataset(data_path, split=data_split)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        row_dict: dict = self.dataset[index]

        # prompt_str: str = row_dict[self.prompt_key]
        text=row_dict['instruction']
        history=row_dict['history'] 
        task_type=row_dict['task_type']
        row_dict.pop('verify_bbox', None)
        row_dict.pop('success_rate', None)
        row_dict.pop('scale', None)
        images=[row_dict['image']]
      
        if task_type=='high':
            # prompt_str=  (
            #     f"You are GUI-R1, a reasoning GUI Agent Assistant. In this UI screenshot <image>, I want you to continue executing the command '{text}', with the action history being '{history}'.\n"
            #     "Please provide the action to perform (enumerate from ['complete', 'close/delete', 'press_home', 'click', 'press_back', 'type', 'select', 'scroll', 'enter']), the point where the cursor is moved to (integer) if a click is performed, and any input text required to complete the action.\n"
            #     "Output the thinking process in <think> </think> tags, and the final answer in <answer> </answer> tags as follows:\n"
            #     "<think> ... </think> <answer>[{'action': enum['complete', 'close/delete', 'press_home', 'click', 'press_back', 'type', 'select', 'scroll', 'enter'], 'point': [x, y], 'input_text': 'no input text [default]'}]</answer>\n"
            #     "Note:\n specific input text (no default) is necessary for actions enum['type', 'select', 'scroll'] \n Example:\n"
            #     "[{'action': enum['complete', 'close/delete', 'press_home', 'press_back', 'enter'], 'point': [-100, -100], 'input_text': 'no input text'}]\n"
            #     "[{'action': enum['click'], 'point': [123, 300], 'input_text': 'no input text'}]\n"
            #     "[{'action': enum['type', 'select'], 'point': [-100, -100], 'input_text': 'shanghai shopping mall'}]\n"
            #     "[{'action': enum['scroll'], 'point': [-100, -100], 'input_text': enum['up', 'left', 'right', 'down']}]"
            # )
            
            prompt_str=(
            f"You are GUI-R1, a reasoning GUI Agent Assistant. In this UI screenshot <image>, I want you to continue executing the command '{text}', with the action history being '{history}'.\n"
            "Please provide the action to perform (enumerate from ['wait', 'long_press', 'click', 'press_back', 'type', 'open_app', 'scroll']), the point where the cursor is moved to (integer) if a click is performed, and any input text required to complete the action.\n"
            "Output the final answer ONLY in <answer> </answer> tags as follows:\n"
            "<answer>[{'action': enum['wait', 'long_press', 'click', 'press_back', 'type', 'open_app', 'scroll'], 'point': [x, y], 'input_text': 'no input text [default]'}]</answer>\n"
            "Note:\n"
            "1. Specific input text (no default) is necessary for actions enum['type', 'open_app', 'scroll'].\n"
            "2. ONLY ONE action is allowed. Output exactly one action.\n"
            "Example:\n"
            "<answer>[{'action': enum['wait'], 'point': [-100, -100], 'input_text': 'no input text'}]</answer>\n"
            "<answer>[{'action': enum['click'], 'point': [123, 300], 'input_text': 'no input text'}]</answer>\n"
            "<answer>[{'action': enum['type'], 'point': [-100, -100], 'input_text': 'shanghai shopping mall'}]</answer>\n"
            "<answer>[{'action': enum['scroll'], 'point': [-100, -100], 'input_text': 'down'}]</answer>"
        )
            
            
        else:
            # prompt_str=(
            #     f"You are GUI-R1, a reasoning GUI Agent Assistant. In this UI screenshot <image>, I want you to continue executing the command '{text}', with the action history being '{history}'.\n"
            #     "Please provide the action to perform (enumerate from ['click']), the point where the cursor is moved to (integer) if a click is performed, and any input text required to complete the action.\n"
            #     "Output the thinking process in <think> </think> tags, and the final answer in <answer> </answer> tags as follows:\n"
            #     "<think> ... </think> <answer>[{'action': enum[ 'click'], 'point': [x, y], 'input_text': 'no input text'}]</answer>\n"
            #     "Example:\n"
            #     "[{'action': enum['click'], 'point': [123, 300], 'input_text': 'no input text'}]\n"
            # )
            
            prompt_str=(
                f"You are GUI-R1, a reasoning GUI Agent Assistant. In this UI screenshot <image>, I want you to continue executing the command '{text}', with the action history being '{history}'.\n"
                "Please provide the action to perform (enumerate from ['click']), the point where the cursor is moved to (integer) if a click is performed, and any input text required to complete the action.\n"
                "Output the final answer ONLY in <answer> </answer> tags as follows: <answer>[{'action': enum['click'], 'point': [x, y], 'input_text': 'no input text'}]</answer>\n"
                "Example:\n"
                "<answer>[{'action': enum['click'], 'point': [123, 300], 'input_text': 'no input text'}]</answer>\n"
            )
            
            
        # prompt_str = (
        #     f"You are GUI-R1, a reasoning GUI Agent Assistant. In this UI screenshot <image>, "
        #     f"I want you to continue executing the command '{text}', with the action history being '{history}'.\n"
        #     "Please describe in natural language what specific action should be performed next "
        #     "(for example: 'click the Close button', 'scroll down', 'type the username', 'press back').\n"
        #     "Output your reasoning process inside <think></think> tags, and provide the final concise action description "
        #     "inside <answer></answer> tags.\n"
        #     "Format example:\n"
        #     "<think> The next step is to close the review panel because the instruction says to finish the process. </think>\n"
        #     "<answer>Click the Close button</answer>\n"
        #     "Another example:\n"
        #     "<think> The user needs to input a location name to continue searching. </think>\n"
        #     "<answer>Type the location name in the search box</answer>\n"
        #     "Note:\n"
        #     "- Your final <answer> must be a short natural-language phrase describing the action, not coordinates or code.\n"
        #     "- Keep the answer concise and actionable (e.g., 'Click Confirm', 'Press Back', 'Select Settings').\n"
        # )
        
        
        
        # ############# Open 标准版 #############
        # prompt_str = (
        #     f"I want you to continue executing the command '{text}', with the action history being '{history}'.\n" 
        # )
            
            

        messages = [{"role": "user", "content": prompt_str}]
        images=[process_image(image, self.max_pixels, self.min_pixels) for image in images]

        scalex,scaley=images[0].size
        gt_bbox=row_dict['gt_bbox']
        gt_bbox[0]*=scalex
        gt_bbox[1]*=scaley
        if len(gt_bbox)>2:
            gt_bbox[2]*=scalex
            gt_bbox[3]*=scaley

	
        ############## 这里 要及时切换 看数据 ##############
        ##############################################################################################################################
        # gt={'action': row_dict['gt_action'],'gt_bbox': gt_bbox,'input_text': row_dict['gt_input_text']}
        # gt={'action': row_dict['gt_action'],'gt_bbox': gt_bbox,'input_text': row_dict['gt_input_text'],'gt_open_ans': row_dict['gt_action_open']}
        
        if len(gt_bbox)>2:
            gt={'action': row_dict['gt_action'],'gt_bbox': bbox_center(gt_bbox),'input_text': row_dict['gt_input_text'],'gt_open_ans': ' ' }
        else:
            gt={'action': row_dict['gt_action'],'gt_bbox': gt_bbox,'input_text': row_dict['gt_input_text'],'gt_open_ans': ' ' }
        # if self.system_prompt:
        #     messages.insert(0, {"role": "system", "content": self.system_prompt})
        ##############################################################################################################################
        
        
        
        prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)

        # if self.image_key in row_dict:
        prompt = prompt.replace("<image>", "<|vision_start|><|image_pad|><|vision_end|>")
        row_dict["multi_modal_data"] = {
            "image": images
        }
        model_inputs = self.processor(row_dict["multi_modal_data"]["image"], prompt, return_tensors="pt")
        input_ids = model_inputs.pop("input_ids")[0]
        attention_mask = model_inputs.pop("attention_mask")[0]
        row_dict["multi_modal_inputs"] = dict(model_inputs)
        position_ids = get_rope_index(
            self.processor,
            input_ids=input_ids,
            image_grid_thw=model_inputs["image_grid_thw"],
            attention_mask=attention_mask,
        )  # (3, seq_length)
        # else:
        #     model_inputs = self.tokenizer([prompt], add_special_tokens=False, return_tensors="pt")
        #     input_ids = model_inputs.pop("input_ids")[0]
        #     attention_mask = model_inputs.pop("attention_mask")[0]
        #     position_ids = torch.clip(attention_mask.cumsum(dim=0) - 1, min=0, max=None)  # (seq_length,)

        input_ids, attention_mask, position_ids = VF.postprocess_data(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            max_length=self.max_prompt_length,
            pad_token_id=self.tokenizer.pad_token_id,
            left_pad=True,
            truncation=self.truncation,
        )
        row_dict["input_ids"] = input_ids
        row_dict["attention_mask"] = attention_mask
        row_dict["position_ids"] = position_ids
        row_dict["raw_prompt_ids"] = self.tokenizer.encode(prompt, add_special_tokens=False)
        row_dict["ground_truth"] = json.dumps(gt)
        return row_dict
