# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# 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 os
from typing import Any, Callable, Dict, List, Optional, Tuple

import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from paddle.io import DataLoader, Dataset

from ..data import DataCollator
from ..datasets import MapDataset
from ..losses import RDropLoss
from ..trainer import Trainer, TrainerCallback
from ..trainer.trainer_utils import EvalPrediction, get_scheduler
from ..transformers import PretrainedTokenizer, export_model
from ..utils.log import logger
from .prompt_args import PromptTuningArguments
from .prompt_utils import PromptDataCollatorWithPadding
from .template import AutoTemplate
from .verbalizer import SoftVerbalizer


class PromptTrainer(Trainer):
    """
    PromptTrainer is a feature-complete training and eval loop for PaddleNLP
    on prompt-tuning.
    """

    def __init__(
        self,
        model: nn.Layer,
        tokenizer: PretrainedTokenizer,
        criterion: Optional[nn.Layer] = None,
        args: Optional[PromptTuningArguments] = None,
        data_collator: Optional[DataCollator] = None,
        train_dataset: Optional[MapDataset] = None,
        eval_dataset: Optional[MapDataset] = None,
        compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
        callbacks: Optional[List[TrainerCallback]] = None,
        optimizers: Tuple[paddle.optimizer.Optimizer, paddle.optimizer.lr.LRScheduler] = (None, None),
    ):
        if args is None:
            output_dir = "tmp_trainer"
            logger.info(
                "No `TrainingArguments` passed, initialized with " "output_dir={} by default.".format(output_dir)
            )
            args = PromptTuningArguments(output_dir=output_dir)

        if data_collator is None:
            data_collator = PromptDataCollatorWithPadding(tokenizer, padding=True, return_tensors="pd")

        if criterion is None and (args.use_rgl or args.use_rdrop):
            raise Exception("'To use 'use_rgl', 'use_rdrop', 'criterion' must be specified")

        super(PromptTrainer, self).__init__(
            model=model,
            criterion=criterion,
            args=args,
            data_collator=data_collator,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            tokenizer=tokenizer,
            compute_metrics=compute_metrics,
            callbacks=callbacks,
            optimizers=optimizers,
        )

        self._load_from_checkpoint(args.resume_from_checkpoint)

        self.train_dataset = self._map_dataset(self.train_dataset)
        self.eval_dataset = self._map_dataset(self.eval_dataset)

        if self.args.use_rdrop:
            self.rdrop_criterion = RDropLoss()

    def _get_model(self):
        model = self.model
        if isinstance(model, paddle.DataParallel):
            model = model._layers
        return model

    @property
    def template(self):
        return self._get_model().template

    @template.setter
    def template(self, template):
        self._get_model().template = template

    @property
    def verbalizer(self):
        return self._get_model().verbalizer

    @verbalizer.setter
    def verbalizer(self, verbalizer):
        self._get_model().verbalizer = verbalizer

    @property
    def pretrained_model(self):
        return self._get_model().plm

    @pretrained_model.setter
    def pretrained_model(self, model):
        setattr(self._get_model(), "plm", model)

    def _map_dataset(self, dataset: MapDataset):
        if dataset is None:
            return None
        if not isinstance(dataset, MapDataset):
            raise ValueError("Expected `MapDataset` but received {}.".format(type(dataset)))

        def encode_with_template(example):
            return self.template(example)

        return dataset.map(encode_with_template)

    def _prepare_input(self, inputs: Dict):
        return inputs

    def _save(
        self,
        output_dir: Optional[str] = None,
        state_dict: Dict[str, Any] = None,
        merge_tensor_parallel: Optional[bool] = True,
    ):
        super(PromptTrainer, self)._save(output_dir, state_dict, merge_tensor_parallel)
        output_dir = output_dir if output_dir is not None else self.args.output_dir
        if self.template:
            self.template.save(output_dir)
        if self.verbalizer is not None:
            self.verbalizer.save(output_dir)
        if self.args.save_plm:
            plm_output_dir = os.path.join(output_dir, "plm")
            os.makedirs(plm_output_dir, exist_ok=True)
            self.pretrained_model.save_pretrained(plm_output_dir)

    def _load_from_checkpoint(self, resume_from_checkpoint: os.PathLike = None):
        if resume_from_checkpoint is not None:
            self.template = AutoTemplate.load_from(
                resume_from_checkpoint, self.tokenizer, self.args.max_seq_length, self._get_model().plm
            )
        super(PromptTrainer, self)._load_from_checkpoint(resume_from_checkpoint)

    def get_test_dataloader(self, test_dataset):
        test_dataset = self._map_dataset(test_dataset)
        return super(PromptTrainer, self).get_test_dataloader(test_dataset)

    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
        if eval_dataset is not None:
            eval_dataset = self._map_dataset(eval_dataset)
        return super(PromptTrainer, self).get_eval_dataloader(eval_dataset)

    def create_optimizer(self, lr_scheduler=None):
        """
        Setup the optimizer for both model and prompt parameters.
        """
        if self.optimizer is None:
            optim_cls, optim_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)

            plm_parameters = []
            if not self.args.freeze_plm:
                plm_parameters.extend([p for p in self._get_model().plm.parameters() if not p.stop_gradient])

            ppt_parameters = []
            if self.template is not None:
                ppt_parameters.extend([x for n, x in self.template.named_parameters() if not x.stop_gradient])
            if self.verbalizer is not None:
                if isinstance(self.verbalizer, SoftVerbalizer):
                    if not self.args.freeze_plm:
                        plm_parameters.extend(
                            [p for n, p in self.verbalizer.non_head_parameters() if not p.stop_gradient]
                        )
                    ppt_parameters.extend([p for n, p in self.verbalizer.head_parameters()])
                else:
                    ppt_parameters.extend([p for n, p in self.verbalizer.parameters()])

            decay_parameters = [
                p.name for n, p in self._get_model().named_parameters() if not any(nd in n for nd in ["bias", "norm"])
            ]

            if len(plm_parameters) > 0:
                ppt_lr = self.args.ppt_learning_rate / self.args.learning_rate
                lr = self.lr_scheduler if lr_scheduler is None else lr_scheduler
                if len(ppt_parameters) > 0:
                    params = [
                        {"params": plm_parameters},
                        {
                            "params": ppt_parameters,
                            "learning_rate": ppt_lr,
                            "weight_decay": self.args.ppt_weight_decay,
                            "beta1": self.args.ppt_adam_beta1,
                            "beta2": self.args.ppt_adam_beta2,
                            "epsilon": self.args.ppt_adam_epsilon,
                        },
                    ]
                else:
                    params = plm_parameters
            else:
                if self.args.max_steps > 0:
                    max_steps = self.args.max_steps
                else:
                    raise ValueError("Please use `max_steps` to set the maximum training steps.")
                warmup = (
                    self.args.warmup_steps if self.args.warmup_steps > 0 else int(self.args.warmup_ratio * max_steps)
                )
                self.lr_scheduler = get_scheduler(
                    self.args.lr_scheduler_type,
                    learning_rate=self.args.ppt_learning_rate,
                    num_warmup_steps=warmup,
                    num_training_steps=max_steps,
                )
                lr = self.lr_scheduler
                params = ppt_parameters

            self.optimizer = optim_cls(
                learning_rate=lr,
                apply_decay_param_fun=lambda x: x in decay_parameters,
                parameters=params,
                weight_decay=self.args.weight_decay,
                grad_clip=nn.ClipGradByGlobalNorm(self.args.max_grad_norm),
                **optim_kwargs,
            )

        return self.optimizer

    def compute_loss(self, model, inputs, return_outputs=False):
        """
        Compute the total loss for every batch.
        """
        if "labels" not in inputs:
            raise ValueError("Fail to compute loss as `labels` not in {}.".format(inputs))
        labels = inputs["labels"]

        input_dict = inputs.copy()

        if self.criterion is not None:
            # pop labels to move loss computation out of the model
            input_dict.pop("labels")
            input_dict["return_hidden_states"] = True
            logits, hidden_states = model(**input_dict)
            loss = self.criterion(logits, labels)

            if self.args.use_rdrop:
                loss = self._compute_rdrop_loss(model, input_dict, labels, logits, loss)

            if self.args.use_rgl:
                loss += self._compute_rgl_loss(hidden_states, labels)
        else:
            loss, logits = model(**input_dict)

        outputs = (loss, logits)

        return (loss, outputs) if return_outputs else loss

    def _compute_rdrop_loss(self, model, input_dict, labels, outputs, loss):
        re_outputs, _ = model(**input_dict)
        ce_loss = (self.criterion(re_outputs, labels) + loss) * 0.5
        kl_loss = self.rdrop_criterion(outputs, re_outputs)
        loss = ce_loss + self.args.alpha_rdrop * kl_loss
        return loss

    def _compute_rgl_loss(self, embeddings, labels, equal_type="raw"):
        """
        Compute the label consistency loss of sentence embeddings per batch.
        Please refer to https://aclanthology.org/2022.findings-naacl.81/
        for more details.
        """

        def _max_equal(x, y):
            return int(paddle.argmax(x, axis=0) == paddle.argmax(y, axis=0))

        def _raw_equal(x, y):
            return int(x == y)

        if equal_type == "raw":
            equals = _raw_equal
        elif equal_type == "max":
            equals = _max_equal
        else:
            raise ValueError("Unsupported equal type {}.".format(equal_type))
        batch_size = embeddings.shape[0]
        loss = 0
        for i in range(batch_size):
            for j in range(batch_size):
                score = F.cosine_similarity(embeddings[i], embeddings[j], axis=0)
                score = score.unsqueeze(0)
                logits = paddle.concat([(1 - score) * 50, (1 + score) * 50], axis=-1)
                label = paddle.to_tensor([equals(labels[i], labels[j])])
                logits = logits.reshape([-1, logits.shape[-1]])
                loss += F.cross_entropy(logits, label.unsqueeze(0))
        loss = loss / (batch_size * (batch_size - 1))
        loss = loss / 100 * self.args.alpha_rgl

        return loss

    def export_model(self, export_path, input_spec=None, export_type="paddle"):
        os.makedirs(export_path, exist_ok=True)
        self.template.save(export_path)
        if self.verbalizer is not None:
            self.verbalizer.save(export_path)
        if input_spec is None:
            input_spec = self.model.get_input_spec()
        export_model(self.model, input_spec, export_path, export_type)
