# 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

import torch

from meta_researcher.src.metric_utils import process_turns
from verl import DataProto
from meta_researcher.src.reward_score import _default_batch_compute_score_format_answer


class AgentRewardManager:
    """The reward manager."""

    def __init__(self, tokenizer, num_examine, compute_score=None, reward_fn_key="data_source", config=None) -> None:
        self.tokenizer = tokenizer
        self.num_examine = num_examine
        self.compute_score = compute_score or _default_batch_compute_score_format_answer
        self.reward_fn_key = reward_fn_key
        self.config = config


    def __call__(self, data: DataProto, return_dict=False):
        """We will expand this function gradually based on the available datasets"""

        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn
        if "rm_scores" in data.batch.keys():
            if return_dict:
                return {"reward_tensor": data.batch["rm_scores"]}
            else:
                return data.batch["rm_scores"]

        reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
        reward_extra_info = defaultdict(list)

        already_print_data_sources = {}
        
        data_source_list = []
        sequences_list = []
        ground_truth_list = []
        valid_response_length_list = []
        extra_info_list = []
        turn_list = []
        inference_turn_list = []

        group_turns = process_turns(data.batch['turns'], self.config.actor_rollout_ref.rollout.n_repeat)

        for i in range(len(data)):
            data_item = data[i]  # DataProtoItem
            turn = data_item.batch['turns']
            inference_turn = group_turns[i // self.config.actor_rollout_ref.rollout.n_repeat]

            prompt_ids = data_item.batch["prompts"]

            prompt_length = prompt_ids.shape[-1]

            valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum()
            valid_prompt_ids = prompt_ids[-valid_prompt_length:]

            response_ids = data_item.batch["responses"]
            valid_response_length = data_item.batch["attention_mask"][prompt_length:].sum()
            valid_response_ids = response_ids[:valid_response_length]

            sequences = torch.cat((valid_prompt_ids, valid_response_ids))

            sequences_str = self.tokenizer.decode(sequences, skip_special_tokens=False)
            pad_token_id = self.tokenizer.pad_token_id
            sequences_str = sequences_str.split(self.tokenizer.decode([pad_token_id]))[0]
            if not sequences_str.endswith(self.tokenizer.eos_token):
                sequences_str += self.tokenizer.eos_token

            ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"]

            data_source = data_item.non_tensor_batch[self.reward_fn_key]

            extra_info = data_item.non_tensor_batch.get("extra_info", None)
        
            data_source_list.append(data_source)
            sequences_list.append(sequences_str)
            ground_truth_list.append(ground_truth)
            valid_response_length_list.append(valid_response_length)
            extra_info_list.append(extra_info)
            turn_list.append(turn)
            inference_turn_list.append(inference_turn)
        
        print("Start of Batch Evaluation Results")
        score_list_dict = self.compute_score(
            data_source_list=data_source_list,
            solution_str_list=sequences_list, 
            ground_truth_list=ground_truth_list,
            extra_info_list=extra_info_list,
            turn_list=turn_list,
            inference_turn_list=inference_turn_list,
            think_len_threshold=self.config.actor_rollout_ref.rollout.think_len_threshold,
            think_len_belta=self.config.actor_rollout_ref.rollout.think_len_belta,
            max_step_num=self.config.tool.max_step_num,
        )
        print("End of Batch Evaluation Results")
     
        if isinstance(score_list_dict, dict):
            reward_list = score_list_dict["reward_list"]
            answer_reward_list = score_list_dict["answer_reward_list"]
            format_reward_list = score_list_dict["format_reward_list"]
            think_len_list = score_list_dict["think_len_list"]
            think_reward_list = score_list_dict["think_reward_list"]
            for i, (reward, answer_reward, format_reward, think_len, think_reward, data_source, valid_response_length, turn) in enumerate(zip(reward_list, answer_reward_list, format_reward_list, think_len_list, think_reward_list, data_source_list, valid_response_length_list, turn_list)):
                reward_extra_info["val_reward"].append(reward)
                reward_extra_info["answer_reward"].append(answer_reward)
                reward_extra_info["format_reward"].append(format_reward)
                reward_extra_info["think_len"].append(think_len)
                reward_extra_info["think_reward"].append(think_reward)
                reward_tensor[i, valid_response_length - 1] = reward
                if data_source not in already_print_data_sources:
                    already_print_data_sources[data_source] = 0
                
                if already_print_data_sources[data_source] < self.num_examine:
                    already_print_data_sources[data_source] += 1
                    print("[prompt+response]", sequences_list[i])
                    print("[ground_truth]", ground_truth_list[i])
                    print("[score]", reward)

        if return_dict:
            return {
                "reward_tensor": reward_tensor,
                "reward_extra_info": dict(reward_extra_info),
            }
        else:
            return reward_tensor
