# 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.
"""
Metrics related to the PPO trainer.
"""

import torch
from typing import Any, Dict, List, Callable
import numpy as np
from verl import DataProto
from collections import Counter, defaultdict


def reduce_metrics(metrics: Dict[str, List[Any]]) -> Dict[str, Any]:
    for key, val in metrics.items():
        metrics[key] = np.mean(val)
    return metrics


def _compute_response_info(batch: DataProto) -> Dict[str, Any]:
    response_length = batch.batch['responses'].shape[-1]

    prompt_mask = batch.batch['attention_mask'][:, :-response_length]
    response_mask = batch.batch['attention_mask'][:, -response_length:]

    prompt_length = prompt_mask.sum(-1).float()
    response_length = response_mask.sum(-1).float()  # (batch_size,)

    return dict(
        response_mask=response_mask,
        prompt_length=prompt_length,
        response_length=response_length,
    )


def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> Dict[str, Any]:
    # TODO: add response length
    sequence_score = batch.batch['token_level_scores'].sum(-1)
    sequence_reward = batch.batch['token_level_rewards'].sum(-1)

    advantages = batch.batch['advantages']
    returns = batch.batch['returns']

    max_response_length = batch.batch['responses'].shape[-1]

    prompt_mask = batch.batch['attention_mask'][:, :-max_response_length].bool()
    response_mask = batch.batch['attention_mask'][:, -max_response_length:].bool()

    max_prompt_length = prompt_mask.size(-1)

    response_info = _compute_response_info(batch)
    prompt_length = response_info['prompt_length']
    response_length = response_info['response_length']

    valid_adv = torch.masked_select(advantages, response_mask)
    valid_returns = torch.masked_select(returns, response_mask)

    if use_critic:
        values = batch.batch['values']
        valid_values = torch.masked_select(values, response_mask)
        return_diff_var = torch.var(valid_returns - valid_values)
        return_var = torch.var(valid_returns)

    metrics = {
        # score
        'critic/score/mean':
            torch.mean(sequence_score).detach().item(),
        'critic/score/max':
            torch.max(sequence_score).detach().item(),
        'critic/score/min':
            torch.min(sequence_score).detach().item(),
        # reward
        'critic/rewards/mean':
            torch.mean(sequence_reward).detach().item(),
        'critic/rewards/max':
            torch.max(sequence_reward).detach().item(),
        'critic/rewards/min':
            torch.min(sequence_reward).detach().item(),
        # adv
        'critic/advantages/mean':
            torch.mean(valid_adv).detach().item(),
        'critic/advantages/max':
            torch.max(valid_adv).detach().item(),
        'critic/advantages/min':
            torch.min(valid_adv).detach().item(),
        # returns
        'critic/returns/mean':
            torch.mean(valid_returns).detach().item(),
        'critic/returns/max':
            torch.max(valid_returns).detach().item(),
        'critic/returns/min':
            torch.min(valid_returns).detach().item(),
        **({
            # values
            'critic/values/mean': torch.mean(valid_values).detach().item(),
            'critic/values/max': torch.max(valid_values).detach().item(),
            'critic/values/min': torch.min(valid_values).detach().item(),
            # vf explained var
            'critic/vf_explained_var': (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
        } if use_critic else {}),

        # response length
        'response_length/mean':
            torch.mean(response_length).detach().item(),
        'response_length/max':
            torch.max(response_length).detach().item(),
        'response_length/min':
            torch.min(response_length).detach().item(),
        'response_length/clip_ratio':
            torch.mean(torch.eq(response_length, max_response_length).float()).detach().item(),
        # prompt length
        'prompt_length/mean':
            torch.mean(prompt_length).detach().item(),
        'prompt_length/max':
            torch.max(prompt_length).detach().item(),
        'prompt_length/min':
            torch.min(prompt_length).detach().item(),
        'prompt_length/clip_ratio':
            torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(),
    }
    return metrics


def compute_timing_metrics(batch: DataProto, timing_raw: Dict[str, float]) -> Dict[str, Any]:
    response_info = _compute_response_info(batch)
    num_prompt_tokens = torch.sum(response_info['prompt_length']).item()
    num_response_tokens = torch.sum(response_info['response_length']).item()
    num_overall_tokens = num_prompt_tokens + num_response_tokens

    num_tokens_of_section = {
        'gen': num_response_tokens,
        **{
            name: num_overall_tokens for name in ['ref', 'values', 'adv', 'update_critic', 'update_actor']
        },
    }

    return {
        **{
            f'timing_s/{name}': value for name, value in timing_raw.items()
        },
        **{
            f'timing_per_token_ms/{name}': timing_raw[name] * 1000 / num_tokens_of_section[name] for name in set(num_tokens_of_section.keys(
            )) & set(timing_raw.keys())
        },
    }


def compute_throughout_metrics(batch: DataProto, timing_raw: Dict[str, float], n_gpus: int) -> Dict[str, Any]:
    total_num_tokens = sum(batch.meta_info['global_token_num'])
    time = timing_raw['step']
    # estimated_flops, promised_flops = flops_function.estimate_flops(num_tokens, time)
    # f'Actual TFLOPs/s/GPU​': estimated_flops/(n_gpus),
    # f'Theoretical TFLOPs/s/GPU​': promised_flops,
    return {
        'perf/total_num_tokens': total_num_tokens,
        'perf/time_per_step': time,
        'perf/throughput': total_num_tokens / (time * n_gpus),
    }


import pandas as pd
import json
import multiprocessing
lock = multiprocessing.Lock()
def compute_reward_metrics(batch):
    data_source = batch.non_tensor_batch['data_source']
    index = batch.non_tensor_batch['uid']
    reward_tensor = batch.batch['token_level_scores'].sum(-1)
    reward_metrics = {}
    zero_correct_indices = []  # 用于存储正确率为 0 的 index

    # TODO: 未考虑混数据集的情况
    if data_source[0] == 'codev':
        # index = batch.non_tensor_batch['problem_id']
        reward_correct = 1.0
        reward_wrong_answer = 0.0
        reward_format_error = 0.0
        
        # exceed_length = torch.sum(reward_tensor == 0.0).float() / reward_tensor.numel()
        # reward_metrics["reward/exceed_length_ratio"] = exceed_length.detach().item()
    elif data_source[0] == "deepscaler":
        reward_correct = 1.0
        reward_wrong_answer = 0.0
        reward_format_error = 0.0
    else:
        return reward_metrics
        
    reward_metrics["reward/mean"] = torch.mean(reward_tensor).detach().item()
    all_correct = torch.sum(reward_tensor == reward_correct).float() / reward_tensor.numel()
    reward_metrics["reward/all_correct_ratio"] = all_correct.detach().item()
    wrong_answer = torch.sum(reward_tensor == reward_wrong_answer).float() / reward_tensor.numel()
    reward_metrics["reward/wrong_answer_ratio"] = wrong_answer.detach().item()
    if reward_format_error is not None and reward_format_error != reward_wrong_answer:
        format_error = torch.sum(reward_tensor == reward_format_error).float() / reward_tensor.numel()
        reward_metrics["reward/format_error_ratio"] = format_error.detach().item()

    correct = (reward_tensor == reward_correct).float()
    correct_series = pd.Series(correct.numpy())

    # 创建 DataFrame
    df = pd.DataFrame({
        'index': index,
        'correct': correct_series
    })

    # 按 index 分组并计算 correct 的平均值
    average_dict = df.groupby('index')['correct'].mean().to_dict()

    # 筛选出正确率为 0 的 index
    for idx, rate in average_dict.items():
        if rate == 0:
            zero_correct_indices.append(idx)

    correct_rates = list(average_dict.values())
    correct_rates_series = pd.Series(correct_rates)

    # 定义分类区间
    bins = [0, 0.00001, 0.5, 0.99999, 1.00001]
    labels = ['0%', '(0%, 50%)', '[50%, 100%)', '100%']

    # 对正确率进行分类
    categories = pd.cut(correct_rates_series, bins=bins, labels=labels, right=False)

    # 统计各个区间的数量
    proportions = categories.value_counts(normalize=True).to_dict()
    reward_metrics["reward/correct_0%_ratio"] = proportions['0%']
    reward_metrics["reward/correct_(0%,50%)_ratio"] = proportions['(0%, 50%)']
    reward_metrics["reward/correct_[50%,100%)_ratio"] = proportions['[50%, 100%)']
    reward_metrics["reward/correct_100%_ratio"] = proportions['100%']

    # 收集正确率为 0 的 batch 数据并保存到 JSONL 文件
    print('Correct proportions:', proportions)
    print('Zero correct indices:', zero_correct_indices)
    # global lock
    # with lock:
    #     with open('/nfs_global/S/zhuyaoyu/projects/logicRL/outputs/all_wrong_problems.jsonl', 'a') as f:
    #     # with open('/nfs_global/S/zhuyaoyu/projects/logicRL/outputs/problem_samples.jsonl', 'a') as f:
    #         for i, idx in enumerate(index):
    #             # if idx in zero_correct_indices:
    #                 data_item = batch[i]
    #                 # print(data_item)
    #                 sample_data = {
    #                     'problem_id': idx,
    #                     'question': data_item.non_tensor_batch['question'],
    #                     'responses': data_item.non_tensor_batch['sequence_str'],
    #                     'ground_truth': data_item.non_tensor_batch['reward_model']['ground_truth']
    #                 }
    #                 json_line = json.dumps(sample_data)
    #                 f.write(json_line + '\n')
    # raise SystemExit("Task completed and terminated")

    return reward_metrics


def bootstrap_metric(data: list[dict[str, Any]],
                     subset_size: int,
                     reduce_fns: list[Callable[[np.ndarray], float]],
                     n_bootstrap: int = 1000,
                     seed: int = 42) -> list[tuple[float, float]]:
    np.random.seed(seed)

    bootstrap_metric_lsts = [[] for _ in range(len(reduce_fns))]
    for _ in range(n_bootstrap):
        bootstrap_idxs = np.random.choice(len(data), size=subset_size, replace=True)
        bootstrap_data = [data[i] for i in bootstrap_idxs]
        for i, reduce_fn in enumerate(reduce_fns):
            bootstrap_metric_lsts[i].append(reduce_fn(bootstrap_data))
    return [(np.mean(lst), np.std(lst)) for lst in bootstrap_metric_lsts]


def calc_maj_val(data: list[dict[str, Any]], vote_key: str, val_key: str) -> float:
    """
    Calculate the majority voting metric
    """
    vote2vals = defaultdict(list)
    for d in data:
        vote2vals[d[vote_key]].append(d[val_key])

    vote2cnt = {k: len(v) for k, v in vote2vals.items()}
    maj_vote = max(vote2cnt, key=vote2cnt.get)

    maj_val = vote2vals[maj_vote][0]

    return maj_val
