# Copyright (c) ModelScope Contributors. All rights reserved.
import torch
from collections import namedtuple
from functools import partial
from megatron.core import mpu
from trl import KTOTrainer
from typing import Any, Dict

from swift.utils import get_current_device, get_logger
from .rlhf_mixin import MegatronRLHFTrainer

logger = get_logger()


class DummyKTOTrainer(KTOTrainer):
    # For reusing the kto_loss function in TRL.

    def gather_for_metrics(self, input_data, *args, **kwargs):
        output_tensors = torch.empty(
            mpu.get_data_parallel_world_size() * input_data.numel(),
            dtype=input_data.dtype,
            device=input_data.device,
        )
        torch.distributed.all_gather_into_tensor(output_tensors, input_data, group=mpu.get_data_parallel_group())
        return output_tensors

    def __init__(self, args):
        self.accelerator = namedtuple('Accelerator', ['device', 'gather_for_metrics'])(
            device=get_current_device(), gather_for_metrics=self.gather_for_metrics)
        self.loss_type = args.loss_type
        self.beta = args.beta
        self.desirable_weight = args.desirable_weight
        self.undesirable_weight = args.undesirable_weight
        self.calculate_KL = args.calculate_KL


class MegatronKTOTrainer(MegatronRLHFTrainer):

    def __init__(self, args, template):
        super().__init__(args, template)
        self.dummy_kto_trainer = DummyKTOTrainer(args)

    def _kto_get_logps(self, output_tensor, data, is_KL: bool, is_ref: bool, length: int):
        output = self._get_input_tensor(output_tensor, is_KL, is_ref, length, dim=1)
        return self.get_logps(output, data['labels'], data.get('packed_seq_params'))

    def _get_kto_length(self, data: Dict[str, Any]) -> int:
        if 'packed_seq_params' in data:
            return data['packed_seq_params'].cu_seqlens_q[-1] // self.args.context_parallel_size
        else:
            return data['position_ids'].shape[-1]

    def loss_func(self, output_tensor, *, data, kl_data, label):
        length = self._get_kto_length(data)
        policy_logps = self._kto_get_logps(output_tensor, data, False, False, length)
        ref_logps = self._kto_get_logps(output_tensor, data, False, True, length)
        if self.args.calculate_KL:
            policy_KL_logps = self._kto_get_logps(output_tensor, kl_data, True, False, length)
            ref_KL_logps = self._kto_get_logps(output_tensor, kl_data, True, True, length)
        else:
            policy_KL_logps, ref_KL_logps = None, None
        label = output_tensor.new_tensor(label, dtype=torch.bool)
        policy_chosen_logps = policy_logps[label]
        policy_rejected_logps = policy_logps[~label]
        ref_chosen_logps = ref_logps[label]
        ref_rejected_logps = ref_logps[~label]

        loss, chosen_rewards, rejected_rewards, kl = self.dummy_kto_trainer.kto_loss(
            policy_chosen_logps,
            policy_rejected_logps,
            policy_KL_logps,
            ref_chosen_logps,
            ref_rejected_logps,
            ref_KL_logps,
        )

        loss = loss.mean()
        mean_metric = {
            'loss': loss.detach().clone(),
            'kl': kl.squeeze().detach(),
        }
        metric = self._all_reduce_metric(mean_metric)
        chosen_count = chosen_rewards.shape[0]
        rejected_count = rejected_rewards.shape[0]
        sum_metric = {
            'logps/chosen': loss.new_tensor([policy_chosen_logps.detach().nansum(), chosen_count]),
            'logps/rejected': loss.new_tensor([policy_rejected_logps.detach().nansum(), rejected_count]),
            'rewards/chosen': loss.new_tensor([chosen_rewards.nansum(), chosen_count]),
            'rewards/rejected': loss.new_tensor([rejected_rewards.nansum(), rejected_count]),
        }
        metric.update(self._all_reduce_metric(sum_metric, torch.distributed.ReduceOp.SUM))
        # fix megatron-lm bug
        loss = loss / mpu.get_context_parallel_world_size()
        return loss, metric

    @staticmethod
    def _get_input_tensor(input_tensor, is_KL: bool, is_ref: bool, length: int, dim: int):
        # policy, ref, policy_KL, ref_KL
        total_length = input_tensor.shape[dim]
        KL_length = (total_length - 2 * length) // 2
        slice_list = [0, length, 2 * length, total_length - KL_length, total_length]
        idx = is_KL * 2 + is_ref
        slice_ = (slice(None), ) * dim + (slice(slice_list[idx], slice_list[idx + 1]), )
        res = input_tensor[slice_]
        if is_KL or is_ref:
            res = res.detach()
        return res

    def forward_step(self, data_iterator, model):
        unwrapped_model = model.module.module
        input_tensor = unwrapped_model.get_input_tensor()
        vp_stage = unwrapped_model.vp_stage
        # not support loss_scale
        data, kl_data = self.get_batch(data_iterator, vp_stage)
        label = data.pop('label')
        data.pop('loss_scale', None)
        kl_data.pop('loss_scale', None)

        length = self._get_kto_length(data)
        if self.args.sequence_parallel:
            length //= mpu.get_tensor_model_parallel_world_size()
        with torch.no_grad(), self.null_ref_context() as ref_models:
            ref_model = ref_models[vp_stage or 0]
            if self.args.calculate_KL:
                if input_tensor is not None:
                    ref_model.set_input_tensor(self._get_input_tensor(input_tensor, True, True, length, 0))
                ref_KL_output_tensor = ref_model(**kl_data)

            if input_tensor is not None:
                ref_model.set_input_tensor(self._get_input_tensor(input_tensor, False, True, length, 0))
            ref_output_tensor = ref_model(**data)

        if self.args.calculate_KL:
            with torch.no_grad():
                if input_tensor is not None:
                    unwrapped_model.set_input_tensor(self._get_input_tensor(input_tensor, True, False, length, 0))
                KL_output_tensor = model(**kl_data)

        if input_tensor is not None:
            unwrapped_model.set_input_tensor(self._get_input_tensor(input_tensor, False, False, length, 0))
        output_tensor = model(**data)
        is_pp_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage)
        dim = 1 if is_pp_last_stage else 0
        if self.args.calculate_KL:
            res = torch.concat([output_tensor, ref_output_tensor, KL_output_tensor, ref_KL_output_tensor], dim=dim)
        else:
            res = torch.concat([output_tensor, ref_output_tensor], dim=dim)
        return res, partial(self.loss_func, data=data, kl_data=kl_data, label=label)

    def _prepare_batch(self, data, vp_stage=None):
        res = []
        for key in ['completion_', 'KL_completion_']:
            _data = {k[len(key):]: v for k, v in data.items() if k.startswith(key)}
            if not self.args.calculate_KL and key == 'KL_completion_':
                _data = {}
            else:
                _data = super()._prepare_batch(_data, vp_stage)
            res.append(_data)
        res[0]['label'] = data['label']
        return res

    def _log_callback(self, logs, n_steps):
        super()._log_callback(logs, n_steps)
        if 'rewards/chosen' in logs and 'rewards/rejected' in logs:
            logs['rewards/margins'] = logs['rewards/chosen'] - logs['rewards/rejected']
