# Copyright 2020 The Forte 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.
"""
Example of building a reinforcement learning based,
data augmentation enhanced sentence classifier
based on pre-trained BERT model.
"""
import argparse
import functools
import logging
import os

import torch
import torch.nn.functional as F
import texar.torch as tx
from transformers import BertForMaskedLM

from examples.data_augmentation.reinforcement.config \
    import config_data, config_classifier
from examples.data_augmentation.reinforcement.utils import model_utils
from forte.models.da_rl import MetaAugmentationWrapper, TexarBertMetaModule


parser = argparse.ArgumentParser()
parser.add_argument(
    '--pretrained-model-name', type=str, default='bert-base-uncased',
    choices=tx.modules.BERTEncoder.available_checkpoints(),
    help="Name of the pre-trained downstream checkpoint to load.")
parser.add_argument(
    "--output-dir", default="output/",
    help="The output directory where the model checkpoints will be written.")
parser.add_argument(
    "--do-train", action="store_true", help="Whether to run training.")
parser.add_argument(
    "--do-eval", action="store_true",
    help="Whether to run eval on the dev set.")
parser.add_argument(
    "--do-test", action="store_true",
    help="Whether to run test on the test set.")
parser.add_argument(
    '--augmentation-model-name', type=str, default='bert-base-uncased',
    choices=tx.modules.BERTEncoder.available_checkpoints(),
    help="Name of the pre-trained augmentation model checkpoint to load.")
parser.add_argument(
    '--num-aug', type=int, default=4,
    help="number of augmentation samples when fine-tuning aug model")
parser.add_argument(
    '--classifier-pretrain-epoch', type=int, default=10,
    help="number of epochs to pretrain the classifier")
args = parser.parse_args()


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.root.setLevel(logging.INFO)


class RLAugmentClassifierTrainer:
    def __init__(self):
        self._prepare_data_iterator()
        self._init_aug_model()
        self._init_classifier()

    def _prepare_data_iterator(self):
        tx.utils.maybe_create_dir(args.output_dir)

        # Loads data
        num_train_data = config_data.num_train_data
        self.num_train_steps = \
            int(num_train_data / config_data.train_batch_size *
                config_data.max_train_epoch)

        train_dataset = tx.data.RecordData(
            hparams=config_data.train_hparam, device=device)
        val_dataset = tx.data.RecordData(
            hparams=config_data.eval_hparam, device=device)
        test_dataset = tx.data.RecordData(
            hparams=config_data.test_hparam, device=device)
        self.iterator = tx.data.DataIterator(
            {"train": train_dataset,
             "dev": val_dataset,
             "test": test_dataset}
        )

        self.val_data_iterator = tx.data.DataIterator({"dev": val_dataset})
        self.val_data_iterator.switch_to_dataset("dev")

    def _init_aug_model(self):
        # pylint: disable=protected-access
        # Builds data augmentation BERT
        aug_model = BertForMaskedLM.from_pretrained(
            args.augmentation_model_name)
        aug_model.to(device)
        aug_tokenizer = tx.data.BERTTokenizer(
            pretrained_model_name=args.augmentation_model_name)
        input_mask_ids = aug_tokenizer._map_token_to_id('[MASK]')
        # Builds augmentation optimizer
        aug_lr = 4e-5
        param_optimizer = list(aug_model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if
                        not any(nd in n for nd in no_decay)],
             'weight_decay': 0.01},
            {'params': [p for n, p in param_optimizer if
                        any(nd in n for nd in no_decay)],
             'weight_decay': 0.0}]
        aug_optim = tx.core.BertAdam(
            optimizer_grouped_parameters, betas=(0.9, 0.999),
            eps=1e-6, lr=aug_lr)
        # Builds data augmentation wrapper
        self.aug_wrapper = MetaAugmentationWrapper(
            aug_model, aug_optim, input_mask_ids, device, args.num_aug)

    def _init_classifier(self):
        # Builds BERT for classification task.
        config_downstream = {
            k: v for k, v in config_classifier.__dict__.items()
            if not k.startswith('__') and k != "hyperparams"}

        self.classifier = tx.modules.BERTClassifier(
            pretrained_model_name=args.pretrained_model_name,
            hparams=config_downstream)
        self.classifier.to(device)

        # Builds learning rate decay scheduler
        classifier_lr = 4e-5
        vars_with_decay = []
        vars_without_decay = []
        for name, param in self.classifier.named_parameters():
            if 'layer_norm' in name or name.endswith('bias'):
                vars_without_decay.append(param)
            else:
                vars_with_decay.append(param)
        opt_params = [{
            'params': vars_with_decay,
            'weight_decay': 0.01,
        }, {
            'params': vars_without_decay,
            'weight_decay': 0.0,
        }]
        self.optim = tx.core.BertAdam(
            opt_params, betas=(0.9, 0.999), eps=1e-6, lr=classifier_lr)
        num_warmup_steps = int(self.num_train_steps
                               * config_data.warmup_proportion)
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optim, functools.partial(model_utils.get_lr_multiplier,
                                     total_steps=self.num_train_steps,
                                     warmup_steps=num_warmup_steps))

    def pre_train_classifier_epoch(self):
        r"""Pre-trains model on the training set
        for better weight initialization.
        """
        self.iterator.switch_to_dataset("train")
        self.classifier.train()

        for _ in range(args.classifier_pretrain_epoch):
            for batch in self.iterator:
                self.optim.zero_grad()

                input_ids = batch["input_ids"]
                segment_ids = batch["segment_ids"]
                labels = batch["label_ids"]
                input_length = (1 - (input_ids == 0).int()).sum(dim=1)

                logits, _ = self.classifier(
                    input_ids, input_length, segment_ids)
                loss = self._compute_loss(logits, labels)

                loss.backward()
                self.optim.step()
                self.scheduler.step()

    def train_epoch(self):
        r"""Trains on the training set, and evaluates on the validation set
        periodically.
        """
        self.iterator.switch_to_dataset("train")
        self.classifier.train()
        self.optim.zero_grad()

        for batch in self.iterator:
            input_ids = batch["input_ids"]
            input_mask = batch["input_mask"]
            segment_ids = batch["segment_ids"]
            labels = batch["label_ids"]

            # Train augmentation model params phi.
            self.aug_wrapper.reset_model()
            # Iterate over training instances.
            num_instances = len(input_ids)
            for i in range(num_instances):
                features = (input_ids[i], input_mask[i],
                            segment_ids[i], labels[i])

                # Augmented instance with params phi exposed
                aug_probs, input_mask_aug, segment_ids_aug, label_ids_aug = \
                    self.aug_wrapper.augment_instance(features)

                # Compute classifier loss.
                self.classifier.zero_grad()
                input_length_aug = ((input_mask_aug == 1).int()).sum(dim=1)
                logits, _ = self.classifier(
                    aug_probs, input_length_aug, segment_ids_aug)
                loss = self._compute_loss(logits, label_ids_aug)
                # Update classifier params on meta_model.
                meta_model = TexarBertMetaModule(self.classifier)
                meta_model = self.aug_wrapper.update_meta_model(
                    meta_model, loss, self.classifier, self.optim)

                # Compute grads of aug_model on validation data.
                for val_batch in self.val_data_iterator:   # one batch
                    val_input_ids = val_batch["input_ids"]
                    val_segment_ids = val_batch["segment_ids"]
                    val_labels = val_batch["label_ids"]
                    val_input_length = \
                        (1 - (val_input_ids == 0).int()).sum(dim=1)
                    val_logits, _ = meta_model(val_input_ids,
                                               val_input_length,
                                               val_segment_ids)
                    val_loss = self._compute_loss(val_logits, val_labels)
                    val_loss = val_loss / num_instances / args.num_aug \
                               / len(self.val_data_iterator)
                    val_loss.backward()

            # Update aug_model param phi.
            self.aug_wrapper.update_phi()

            # Train classifier with augmented batch
            input_probs, input_masks, segment_ids, label_ids = \
                self.aug_wrapper.augment_batch((input_ids, input_mask,
                                          segment_ids, labels))

            input_length = ((input_masks == 1).int()).sum(dim=1)
            self.optim.zero_grad()
            logits, _ = self.classifier(input_probs, input_length, segment_ids)
            loss = self._compute_loss(logits, label_ids)
            loss.backward()
            self.optim.step()
            self.scheduler.step()
            self._display_logging(loss)

    @torch.no_grad()
    def eval_epoch(self):
        """Evaluates on the dev set.
        """
        self.iterator.switch_to_dataset("dev")
        self.classifier.eval()

        nsamples = 0
        avg_rec = tx.utils.AverageRecorder()
        for batch in self.iterator:
            input_ids = batch["input_ids"]
            segment_ids = batch["segment_ids"]
            labels = batch["label_ids"]
            input_length = (1 - (input_ids == 0).int()).sum(dim=1)

            logits, preds = \
                self.classifier(input_ids, input_length, segment_ids)
            loss = self._compute_loss(logits, labels)
            accu = tx.evals.accuracy(labels, preds)

            batch_size = input_ids.size()[0]
            avg_rec.add([accu, loss], batch_size)
            nsamples += batch_size
        logging.info("eval accu: %.4f; loss: %.4f; nsamples: %d",
                     avg_rec.avg(0), avg_rec.avg(1), nsamples)

    @torch.no_grad()
    def test_epoch(self, test_file):
        """Does predictions on the test set.
        """
        self.iterator.switch_to_dataset("test")
        self.classifier.eval()

        _all_preds = []
        nsamples = 0
        avg_rec = tx.utils.AverageRecorder()
        for batch in self.iterator:
            input_ids = batch["input_ids"]
            segment_ids = batch["segment_ids"]
            labels = batch["label_ids"]
            input_length = (1 - (input_ids == 0).int()).sum(dim=1)

            logits, preds = \
                self.classifier(input_ids, input_length, segment_ids)
            loss = self._compute_loss(logits, labels)
            accu = tx.evals.accuracy(labels, preds)

            batch_size = input_ids.size()[0]
            avg_rec.add([accu, loss], batch_size)
            nsamples += batch_size

            _all_preds.extend(preds.tolist())

        logging.info("test accu: %.4f; loss: %.4f; nsamples: %d",
                     avg_rec.avg(0), avg_rec.avg(1), nsamples)

        output_file = os.path.join(args.output_dir, test_file)
        with open(output_file, "w+") as writer:
            writer.write("\n".join(str(p) for p in _all_preds))
        logging.info("test output written to %s", output_file)

    def _compute_loss(self, logits, labels):
        r"""Compute loss.
        """
        if self.classifier.is_binary:
            loss = F.binary_cross_entropy(
                logits.view(-1), labels.view(-1), reduction='mean')
        else:
            loss = F.cross_entropy(
                logits.view(-1, self.classifier.num_classes),
                labels.view(-1), reduction='mean')
        return loss

    def _display_logging(self, loss):
        step = self.scheduler.last_epoch
        dis_steps = config_data.display_steps
        if dis_steps > 0 and step % dis_steps == 0:
            logging.info("step: %d; loss: %f", step, loss)

        eval_steps = config_data.eval_steps
        if eval_steps > 0 and step % eval_steps == 0:
            self._eval_epoch()
            self.classifier.train()


def main():
    trainer = RLAugmentClassifierTrainer()
    trainer.pre_train_classifier_epoch()
    if args.do_train:
        for k in range(config_data.max_train_epoch):
            logging.info("training epoch %d", k)
            trainer.train_epoch()
    if args.do_eval:
        trainer.eval_epoch()
    if args.do_test:
        trainer.test_epoch("test_results.tsv")


if __name__ == "__main__":
    main()
