# Copyright (c) ModelScope Contributors. All rights reserved.
import pandas as pd
import torch
import torch.nn as nn
import trl
from accelerate.utils import gather_object
from collections import defaultdict
from contextlib import nullcontext
from packaging import version
from transformers import PreTrainedModel
from trl import RewardTrainer as HFRewardTrainer
from typing import Any, Dict, Tuple, Union

from swift.trainers import SwiftMixin
from swift.utils import get_logger, swanlab_get_run
from .rlhf_mixin import RLHFTrainerMixin

try:
    from trl.trainer.utils import print_rich_table
except ImportError:
    from trl.experimental.ppo.ppo_trainer import print_rich_table

del HFRewardTrainer.__init__

logger = get_logger()


class RewardTrainer(RLHFTrainerMixin, SwiftMixin, HFRewardTrainer):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._metrics = {'train': defaultdict(list), 'eval': defaultdict(list)}
        if version.parse(trl.__version__) >= version.parse('0.24'):
            # During evaluation, Trainer calls compute_loss() only if can_return_loss is True and label_names is empty.
            self.can_return_loss = True
            self.label_names = []

    def compute_loss(self,
                     model: Union[PreTrainedModel, nn.Module],
                     inputs: Dict[str, Union[torch.Tensor, Any]],
                     return_outputs=False,
                     num_items_in_batch=None) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
        margin = inputs.pop('margin', None)
        attention_mask = inputs['attention_mask']
        batch_size = attention_mask.shape[0] // 2
        rewards = model(**inputs).logits
        rewards_chosen, rewards_rejected = torch.split(rewards, batch_size, dim=0)
        if margin is not None:
            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - margin).mean()
        else:
            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
        mode = 'train' if self.model.training else 'eval'
        if self.args.center_rewards_coefficient is not None:
            center_rewards_loss = self.args.center_rewards_coefficient * torch.mean(
                (rewards_chosen + rewards_rejected)**2)
            loss += center_rewards_loss
            self.custom_metrics[mode]['center_rewards_loss'].update(center_rewards_loss.detach())
        # metrics
        rewards_chosen, rewards_rejected = rewards_chosen.detach(), rewards_rejected.detach()
        self.custom_metrics[mode]['rewards/chosen'].update(rewards_chosen.mean())
        self.custom_metrics[mode]['rewards/rejected'].update(rewards_rejected.mean())
        self.custom_metrics[mode]['rewards/accuracies'].update((rewards_chosen > rewards_rejected).float().mean())
        self.custom_metrics[mode]['rewards/margins'].update((rewards_chosen - rewards_rejected).mean())
        # compat transformers>=4.46.*
        if num_items_in_batch is not None and self.model_accepts_loss_kwargs:
            loss = loss / self.args.gradient_accumulation_steps
        if return_outputs:
            return loss, {
                'rewards_chosen': rewards_chosen,
                'rewards_rejected': rewards_rejected,
            }
        return loss

    def visualize_samples(self, num_print_samples: int):
        """
        Visualize the reward model logits prediction

        Args:
            num_print_samples (`int`, defaults to `4`):
                The number of samples to print. Set to `-1` to print all samples.
        """
        eval_dataloader = self.get_eval_dataloader()
        table = defaultdict(list)
        for _, inputs in enumerate(eval_dataloader):
            _, logits, _ = self.prediction_step(self.model, inputs, prediction_loss_only=False)
            input_ids = inputs['input_ids']
            attention_mask = inputs['attention_mask']
            sequence_lengths = ((torch.eq(attention_mask, 0).int().argmax(-1) - 1) % attention_mask.shape[1]).tolist()
            text = [self.template.safe_decode(tokens[:sequence_lengths[i]]) for i, tokens in enumerate(input_ids)]
            batch_size = input_ids.shape[0] // 2
            chosen_text, rejected_text = text[:batch_size], text[batch_size:]
            table['chosen_text'].extend(gather_object(chosen_text))
            table['rejected_text'].extend(gather_object(rejected_text))
            table['logits'].extend(
                gather_object([[round(inner_item, 4) for inner_item in item] for item in logits.tolist()]))
            if 0 <= num_print_samples <= len(table['chosen_text']):
                break
        df = pd.DataFrame(table)
        if self.accelerator.process_index == 0:
            try:
                print_rich_table(df[:num_print_samples])
            except Exception as e:
                logger.error(e)
            if 'wandb' in self.args.report_to:
                import wandb

                if wandb.run is not None:
                    wandb.log({'completions': wandb.Table(dataframe=df)})

            if 'swanlab' in self.args.report_to:
                import swanlab
                if swanlab_get_run() is not None:
                    swanlab_table = swanlab.echarts.Table()
                    swanlab_table.add(headers=df.columns.tolist(), rows=df.values.tolist())
                    swanlab.log({'completions': swanlab_table})
