import torch
# import torch.nn as nn
# import torch.nn.functional as F
import torch.optim as optim
# from torch.optim.lr_scheduler import ReduceLROnPlateau, ExponentialLR
from beam import BeamSearch
from tensorboardX import SummaryWriter
from utils import batch_bleu, bert_dual_sequence_mask
import os
import logging
from utils import mkdir
from dotmap import DotMap
import tqdm
import json
# import sacrebleu
import numpy as np
# from transformers import BertModel, RobertaModel
from itertools import chain
from bert_score import BERTScorer
# from itertools import islice

from data import bos_token

from pytorch_memlab import profile

scorer = BERTScorer(lang="en")
bert_score = scorer.score

# __transformers__ = [BertModel, RobertaModel]


def build_trainer(model, args, datamaker, phase="train"):
    if phase not in ["train"]:
        raise NotImplementedError(
            "PRETRAIN and TUNE modes to be implemented, only TRAIN mode is supported"
        )

    trainer = Trainer(
        model,
        patience=args.patience,
        # val_interval=100,
        serialization_dir=args.serialization_dir,
        # max_vals=50,
        device=args.device,
        clip_grad_norm_val=args.clip,
        initial_lr=args.initial_lr,
        lr_decay=None,
        min_lr=args.min_lr,
        lr_patience=args.lr_patience,
        keep_all_checkpoints=args.keep_all_checkpoints,
        val_data_limit=args.val_data_limit,
        max_epochs=args.max_epochs,
        training_data_fraction=args.training_data_fraction,
        beam_size=args.beam_size,
        min_length=args.min_length,
        max_length=args.max_length,
        n_best=1,
        ratio=None,
        datamaker=datamaker,
        lr_scheduling_metric=args.lr_scheduling_metric,
        metric_decreases=args.metric_decreases,
        load_model=args.load_model,
        load_optimizer=args.load_optimizer,
        kl_reach_point=args.kl_reach_point,
        warmup_steps=args.warmup_steps,
        validation_interval=args.validation_interval,
    )

    return trainer


logging.basicConfig(
    filename="app.log",
    filemode="w",
    format="%(name)s - %(levelname)s - %(message)s",
    level=logging.DEBUG,
)


class Trainer(object):
    def __init__(
        self,
        model,
        patience=4,
        # val_interval=100,
        serialization_dir=None,
        device="cuda",
        clip_grad_norm_val=None,
        initial_lr=None,
        lr_decay=None,
        min_lr=None,
        lr_patience=None,
        keep_all_checkpoints=False,
        val_data_limit=None,
        max_epochs=-1,
        training_data_fraction=0,
        beam_size=1,
        min_length=3,
        max_length=512,
        n_best=1,
        ratio=None,
        datamaker=None,
        lr_scheduling_metric=None,
        metric_decreases=None,
        kl_reach_point=None,
        load_model=None,
        load_optimizer=None,
        warmup_steps=None,
        validation_interval=None,
    ):
        """
        The training coordinator. Unusually complicated to handle MTL with tasks of
        diverse sizes.
        Parameters
        ----------
        model : ``Model``, required.
            An PyTorch model to be optimized. Can  be optimized if
            their ``forward`` method returns a dictionary with a "loss" key, containing a
            scalar tensor representing the loss function to be optimized.
        patience , optional (default=2)
            Number of validations to be patient before early stopping.
        val_metric , optional (default="loss")
            Validation metric to measure for whether to stop training using patience
            and whether to serialize an ``is_best`` model after each validation. The metric name
            must be prepended with either "+" or "-", which specifies whether the metric
            is an increasing or decreasing function.
        serialization_dir , optional (default=None)
            Path to directory for saving and loading model files. Models will not be saved if
            this parameter is not passed.
        cuda_device , optional (default = -1)
            An integer specifying the CUDA device to use. If -1, the CPU is used.
            Multi-gpu training is not currently supported, but will be once the
            Pytorch DataParallel API stabilises.
        grad_norm : float, optional, (default = None).
            If provided, gradient norms will be rescaled to have a maximum of this value.
        grad_clipping : ``float``, optional (default = ``None``).
            If provided, gradients will be clipped `during the backward pass` to have an (absolute)
            maximum of this value.  If you are getting ``NaNs`` in your gradients during training
            that are not solved by using ``grad_norm``, you may need this.
        keep_all_checkpoints : If set, keep checkpoints from every validation. Otherwise, keep only
            best and (if different) most recent.
        val_data_limit: During training, use only the first N examples from the validation set.
            Set to -1 to use all.
        training_data_fraction: If set to a float between 0 and 1, load only the specified
            percentage of examples. Hashing is used to ensure that the same examples are loaded
            each epoch.
        """
        self._model = model
        print(model)
        self._load_model = load_model
        if self._load_model:
            self._model.load_state_dict(torch.load(self._load_model))

        self._patience = patience
        self._serialization_dir = serialization_dir
        self._device = device
        self._clip_grad_norm_val = clip_grad_norm_val
        self._lr_decay = lr_decay
        self._min_lr = min_lr
        self._lr_patience = lr_patience
        self._keep_all_checkpoints = keep_all_checkpoints
        self._max_epochs = max_epochs
        self._training_data_fraction = training_data_fraction
        self._initial_lr = initial_lr
        self._lr_scheduling_metric = lr_scheduling_metric
        self._metric_decreases = metric_decreases
        self._val_data_limit = val_data_limit
        self._kl_reach_point = kl_reach_point
        self._warmup_steps = warmup_steps

        self._metric_infos = {}
        # TO BE REMOVED
        self._test_metric_infos = {}
        self._patience_exceeded = False

        no_decay = ["bias", "LayerNorm.weight"]
        if ("transformers" in str(type(self._model.encoder))) or (
            "dataparallel" in str(type(self._model.encoder)).lower()
        ):
            self._trainable_params = filter(
                lambda p: p.requires_grad, self._model.parameters()
            )

            self._optimizer = optim.AdamW(
                [
                    {
                        "params": list(
                            chain(
                                *[
                                    list(
                                        (
                                            filter(
                                                lambda p: p.requires_grad,
                                                module.parameters(),
                                            )
                                        )
                                    )
                                    for module in self._model.children()
                                    if (
                                        ("transformers" in str(type(module)).lower())
                                        or ("dataparallel" in str(type(module)).lower())
                                    )
                                ]
                            )
                        ),
                        "lr": 5e-5,
                        "weight_decay": 0.0,
                    },
                    {
                        "params": list(
                            chain(
                                *[
                                    list(
                                        (
                                            filter(
                                                lambda p: p.requires_grad,
                                                module.parameters(),
                                            )
                                        )
                                    )
                                    for module in self._model.children()
                                    if (
                                        ("transformers" not in str(type(module)))
                                        and (
                                            "dataparallel"
                                            not in str(type(module)).lower()
                                        )
                                    )
                                ]
                            )
                        ),
                        "weight_decay": 0.0,
                    },
                ],
                lr=self._initial_lr,
                eps=1e-8,
            )
        else:
            self._trainable_params = filter(
                lambda p: p.requires_grad, self._model.parameters()
            )
            self._optimizer = optim.Adam(self._trainable_params, lr=self._initial_lr)
        self._epoch_steps = 0
        self._train_counter = 0
        self._validation_counter = 0
        self._validation_steps = 0
        self._initial_lr = [
            param_group["lr"] for param_group in self._optimizer.param_groups
        ]
        self._load_optimizer = load_optimizer
        self._optimizer_loaded = False
        if self._load_optimizer:
            self._optimizer.load_state_dict(torch.load(load_optimizer))
            self._optimizer_loaded = True

        self._bad_epochs = 0
        if self._kl_reach_point:
            self._kl_anneal_function = lambda x: 1 / (
                1 + np.exp(-((x / (self._kl_reach_point / 10)) - 5))
            )
            if self._optimizer_loaded:
                self._train_counter = self._kl_reach_point

        if beam_size == 1:
            log.warining(
                "WARNING: Beam size is 1, note that this is equivalent to greedy search"
            )
        self._beam_size = beam_size
        self._n_best = n_best
        self._min_length = min_length
        self._max_length = max_length
        self._ratio = ratio

        self._datamaker = datamaker

        self._tgt_pad_idx = self._datamaker.vocab.definition.stoi["<pad>"]
        self._tgt_bos_idx = self._datamaker.vocab.definition.stoi["<sos>"]
        self._tgt_eos_idx = self._datamaker.vocab.definition.stoi["<eos>"]
        self._tgt_unk_idx = self._datamaker.vocab.definition.stoi["<unk>"]
        self._exclusion_idxs = {self._tgt_unk_idx, self._tgt_pad_idx, self._tgt_bos_idx}

        self._validation_interval = validation_interval
        self._TB_dir = None
        if self._serialization_dir is not None:
            self._TB_dir = mkdir(os.path.join(self._serialization_dir, "tensorboard"))
            self._TB_train_log = SummaryWriter(
                mkdir(os.path.join(self._TB_dir, "train"))
            )
            self._TB_validation_log = SummaryWriter(
                mkdir(os.path.join(self._TB_dir, "val"))
            )

            self._validation_log_dir = mkdir(
                os.path.join(self._serialization_dir, "valid")
            )
            mkdir(os.path.join(self._serialization_dir, "model"))
            mkdir(os.path.join(self._serialization_dir, "optimizer"))
            self._train_log_dir = mkdir(os.path.join(self._serialization_dir, "train"))
            with open(self._serialization_dir + "/optimizer_printout", "w") as f:
                f.write(repr(self._optimizer))

    def _check_metric_history(
        self, metric_history, current_score, should_decrease=False
    ):
        """
        Given a the history of the performance on a metric
        and the current score, check if current score is
        best so far and if out of patience.
        """
        assert current_score in metric_history

        patience = self._patience + 1
        best_fn = min if should_decrease else max
        best_score = best_fn(metric_history)
        if best_score == current_score:
            best_so_far = metric_history.index(best_score) == len(metric_history) - 1
        else:
            best_so_far = False

        if should_decrease:
            index_of_last_improvement = metric_history.index(min(metric_history))
            out_of_patience = index_of_last_improvement <= len(metric_history) - (
                patience + 1
            )
        else:
            index_of_last_improvement = metric_history.index(max(metric_history))
            out_of_patience = index_of_last_improvement <= len(metric_history) - (
                patience + 1
            )

        return best_so_far, out_of_patience

    def _update_metric_history(
        self, val_pass, metric, current_value, metric_infos, metric_decreases,
    ):
        """
        This function updates metric history with the best validation score so far.
        Parameters
        ---------
        val_pass: int.
        all_val_metrics: dict with performance on current validation pass.
        metric: str, name of metric
        task_name: str, name of task
        metric_infos: dict storing information about the various metrics
        metric_decreases: bool, marker to show if we should increase or
        decrease validation metric.
        should_save: bool, for checkpointing
        new_best: bool, indicator of whether the previous best preformance score was exceeded
        Returns
        ________
        metric_infos: dict storing information about the various metrics
        this_val_metric: dict, metric information for this validation pass, used for optimization
            scheduler
        should_save: bool
        new_best: bool
        """

        metric_exists = metric_infos.get(metric)
        if metric_exists is None:
            metric_infos[metric] = {}
        metric_history = metric_infos[metric].get("hist")
        if metric_history is None:
            metric_infos[metric]["hist"] = []
            metric_history = metric_infos[metric]["hist"]
        metric_history.append(current_value)
        is_best_so_far, out_of_patience = self._check_metric_history(
            metric_history, current_value, metric_decreases
        )
        if is_best_so_far:
            logging.info("Best result seen so far for %s.", metric)
            metric_infos[metric]["best"] = (val_pass, current_value)
            should_save = True
        if out_of_patience:
            metric_infos[metric]["stopped"] = True
        else:
            metric_infos[metric]["stopped"] = False
        return is_best_so_far, out_of_patience

    # @profile
    def _train(self, batch_size):

        assert isinstance(self._model, torch.nn.Module), (
            "Before calling train, you must supply a PyTorch model using the"
            " `Trainer._set_model` method"
        )
        self._epoch_steps += 1
        if self._epoch_steps > self._max_epochs:
            logging.info(
                f"Max Epoch Steps {self._max_epochs} reached. Training Stopped."
            )
            return
        if self._patience_exceeded:
            logging.info(
                f"Patience has already been exceeded for every metric. In other words,"
                f" I've become IMPATIENT. Training Stopped."
            )
            return

        train_iterator = self._datamaker.get_iterator(
            "train", batch_size, device=self._device
        )

        # for testing
        # _test_train_iterator = islice(train_iterator, 0, 1)

        validate_interval = None
        if self._validation_interval is not None:
            if len(train_iterator) > self._validation_interval:
                validation_iters_per_epoch = round(
                    len(train_iterator) / self._validation_interval
                )
                validate_interval = len(train_iterator) // validation_iters_per_epoch
                max_val_interval = (
                    self._epoch_steps * len(train_iterator)
                ) + validate_interval * (validation_iters_per_epoch - 1)

        generations = []
        targets = []
        sources = []
        words = []
        log = {
            "bleu": [],
            "perplexity": [],
            "kld": [],
            "bert-score-p": [],
            "bert-score-r": [],
            "bert-score-f1": [],
        }
        for i, batch in enumerate(
            tqdm.tqdm(train_iterator, desc=f"Training (Epoch {self._epoch_steps}): ")
        ):
            try:
                self._train_counter += 1
                self._model.zero_grad()
                self._model.train()

                example, example_lens = batch.example
                definition, definition_lens = batch.definition
                word, word_lens = batch.word
                if self._model.variational or self._model.defbert:
                    definition_ae, definition_ae_lens = batch.definition_ae
                else:
                    definition_ae, definition_ae_lens = None, None

                sentence_mask = bert_dual_sequence_mask(
                    # example, self._datamaker.vocab.example.encode("</s>")[1:-1], device=self._device
                    # example, self._datamaker.vocab.example.encode(self._datamaker.vocab.example._sep_token)[1:-1], device=self._device
                    example, self._datamaker.vocab.example.encode("[SEP]")[1:-1], device=self._device
                    # example, self._tgt_bos_idx, device=self._device
                )
                current_batch_size = word.shape[0]

                model_out = self._forward(
                    "train",
                    input=example,
                    seq_lens=example_lens,
                    span_token_ids=word,
                    target=definition,
                    target_lens=definition_lens,
                    definition=definition_ae,
                    definition_lens=definition_ae_lens,
                    sentence_mask=sentence_mask,
                )
                if torch.isnan(model_out.perplexity):
                    print(
                        "Loss is NaN. If this happens to often you must debug. YOU MUST"
                    )
                    continue
                torch.cuda.empty_cache()

                torch.nn.utils.clip_grad_norm_(
                    self._trainable_params, self._clip_grad_norm_val
                )

                generations.extend(
                    self._datamaker.decode(
                        model_out.predictions, "definition", batch=True
                    )
                )
                targets.extend(
                    self._datamaker.decode(definition, "definition", batch=True)
                )
                sources.extend(self._datamaker.decode(example, "example", batch=True))
                words.extend(self._datamaker.decode(word, "word", batch=True))

                self._TB_train_log.add_scalar(
                    "loss", model_out.loss.mean().item(), self._train_counter
                )
                current_bleu = batch_bleu(
                    targets[-current_batch_size:],
                    generations[-current_batch_size:],
                    reduction="average",
                )
                try:
                    P, R, F1 = bert_score(
                        generations[-current_batch_size:],
                        targets[-current_batch_size:],
                    )
                except:
                    P, R, F1 = (torch.Tensor([0]), torch.Tensor([0]), torch.Tensor([0]))
                log["bert-score-p"].append(P.mean().item())
                log["bert-score-r"].append(R.mean().item())
                log["bert-score-f1"].append(F1.mean().item())

                self._TB_train_log.add_scalar(
                    "batch_BLEU", current_bleu, self._train_counter
                )

                log["bleu"].append(current_bleu)
                log["perplexity"].append(model_out.perplexity.item())

                loss = model_out.loss
                if model_out.kl is not None:
                    self._TB_train_log.add_scalar(
                        "KL_Divergence", model_out.kl.mean().item(), self._train_counter
                    )
                    self._TB_train_log.add_scalar(
                        "KL_Weight",
                        self._kl_anneal_function(self._train_counter),
                        self._train_counter,
                    )
                    loss = (
                        loss
                        + self._kl_anneal_function(self._train_counter)
                        * model_out.fake_kl
                    ).mean()
                    log["kld"].append(model_out.kl.mean().item())

                loss.backward()
                self._optimizer.step()
                del loss

                for i, param_group in enumerate(self._optimizer.param_groups):
                    self._TB_train_log.add_scalar(
                        f"Learning_rate_{i}", param_group["lr"], self._train_counter,
                    )
                if self._warmup_steps:
                    if not self._optimizer_loaded:
                        if self._train_counter <= self._warmup_steps:
                            learning_rate = [
                                lr * (self._train_counter / self._warmup_steps)
                                for lr in self._initial_lr
                            ]
                            for i, param_group in enumerate(
                                self._optimizer.param_groups
                            ):
                                param_group["lr"] = learning_rate[i]

            except RuntimeError as e:
                # catch out of memory exceptions during fwd/bck (skip batch)
                if "out of memory" in str(e):
                    logging.warning(
                        "| WARNING: ran out of memory, skipping batch. "
                        "if this happens frequently, decrease batch_size or "
                        "truncate the inputs to the model."
                    )
                    torch.cuda.empty_cache()

                    if "loss" in locals():
                        del loss
                    
                    continue
                else:
                    raise e
            if self._train_counter % 200 == 0:
                write_str = (
                    f"| epoch {self._epoch_steps} |"
                    f" {self._train_counter%len(train_iterator)}/{len(train_iterator)}"
                    f" | ppl {sum(log['perplexity'])/len(log['perplexity']):.2f} | bleu"
                    f" {100*sum(log['bleu'])/len(log['bleu']):.2f} | bert-score,"
                    f" p:{sum(log['bert-score-p'])/len(log['bert-score-p']):.2f}"
                    f" r:{sum(log['bert-score-r'])/len(log['bert-score-r']):.2f}"
                    f" f1:{sum(log['bert-score-f1'])/len(log['bert-score-f1']):.2f} |"
                )
                if self._model.variational:
                    write_str += f" KLD {sum(log['kld'])/len(log['kld']):.2f} |"
                for i, param_group in enumerate(self._optimizer.param_groups):
                    write_str += f" lr #{i} {param_group['lr']} |"
                tqdm.tqdm.write(write_str)
                for key in log:
                    log[key] = []
            if validate_interval is not None:
                if validate_interval != 0:
                    if (
                        self._train_counter % validate_interval == 0
                        and self._train_counter <= max_val_interval
                    ):
                        self._validate(64)

        bleu = batch_bleu(targets, generations, reduction="average")
        self._TB_train_log.add_scalar("BLEU", bleu, self._epoch_steps)

        with open(
            os.path.join(self._train_log_dir, f"iter_{self._epoch_steps}.json"), "w",
        ) as f:
            f.write(
                "\n".join(
                    [
                        json.dumps(
                            {
                                "src": sources[i],
                                "tgt": targets[i],
                                "gen": generations[i],
                                "word": words[i],
                            }
                        )
                        for i in range(len(generations))
                    ]
                )
            )
        return DotMap({"src": sources, "tgt": targets, "gen": generations})

    def _validate(self, batch_size):

        assert isinstance(self._model, torch.nn.Module), (
            "Before calling _validate, you must supply a PyTorch model using the"
            " `Trainer._set_model` method"
        )
        valid_iterator = self._datamaker.get_iterator(
            "valid", batch_size, device=self._device
        )

        # for testing
        # _test_valid_iterator = islice(valid_iterator, 1)

        generations = []
        targets = []
        sources = []
        words = []

        self._validation_steps += 1
        ppl = 0
        kld = 0
        for i, batch in enumerate(
            tqdm.tqdm(
                valid_iterator, desc=f"Validating (Epoch {self._validation_steps}): "
            )
        ):
            try:
                self._validation_counter += 1
                self._model.zero_grad()
                self._model.eval()

                example, example_lens = batch.example
                definition, definition_lens = batch.definition
                word, word_lens = batch.word
                if self._model.variational:
                    definition_ae, definition_ae_lens = batch.definition_ae
                else:
                    definition_ae, definition_ae_lens = None, None

                sentence_mask = bert_dual_sequence_mask(
                    # example, self._datamaker.vocab.example.encode("</s>")[1:-1], device=self._device
                    # example, self._datamaker.vocab.example.encode(self._datamaker.vocab.example._sep_token)[1:-1], device=self._device
                    # example, self._tgt_bos_idx, device=self._device
                    example, self._datamaker.vocab.example.encode("[SEP]")[1:-1], device=self._device
                )
                current_batch_size = word.shape[0]

                decode_strategy = BeamSearch(
                    self._beam_size,
                    current_batch_size,
                    pad=self._tgt_pad_idx,
                    bos=self._tgt_bos_idx,
                    eos=self._tgt_eos_idx,
                    n_best=1 if self._n_best is None else self._n_best,
                    global_scorer=self._model.global_scorer,
                    min_length=self._min_length,
                    max_length=self._max_length,
                    return_attention=False,
                    block_ngram_repeat=3,
                    exclusion_tokens=self._exclusion_idxs,
                    stepwise_penalty=None,
                    ratio=self._ratio if self._ratio is not None else 0,
                )
                with torch.no_grad():
                    model_out = self._forward(
                        "valid",
                        input=example,
                        seq_lens=example_lens,
                        span_token_ids=word,
                        target=definition,
                        target_lens=definition_lens,
                        decode_strategy=decode_strategy,
                        definition=definition_ae,
                        definition_lens=definition_ae_lens,
                        sentence_mask=sentence_mask,
                        device=self._device,
                    )
                torch.cuda.empty_cache()

                generations.extend(
                    [
                        self._datamaker.decode(gen[0], "definition", batch=False)
                        for gen in model_out.predictions
                    ]
                )
                targets.extend(
                    self._datamaker.decode(definition, "definition", batch=True)
                )
                sources.extend(self._datamaker.decode(example, "example", batch=True))
                words.extend(self._datamaker.decode(word, "word", batch=True))

                if torch.isnan(model_out.perplexity):
                    tqdm.tqdm.write(
                        "NaN Fouuuuuuuuuund!!!!!!!!!!!!!!! If this happens too often,"
                        " check WTF is going on"
                    )
                    continue
                ppl += model_out.perplexity.item()

                self._TB_validation_log.add_scalar(
                    "batch_perplexity",
                    model_out.perplexity.item(),
                    self._validation_counter,
                )

                current_bleu = batch_bleu(
                    targets[-current_batch_size:],
                    generations[-current_batch_size:],
                    reduction="average",
                )
                self._TB_validation_log.add_scalar(
                    "batch_BLEU", current_bleu, self._validation_counter
                )

                if model_out.kl is not None:
                    kld += model_out.kl.item()
                    self._TB_validation_log.add_scalar(
                        "kl", model_out.kl.item(), self._validation_counter
                    )
                if self._val_data_limit:
                    if i * batch_size > self._val_data_limit:
                        break
            except RuntimeError as e:
                # catch out of memory exceptions during fwd/bck (skip batch)
                if "out of memory" in str(e):
                    logging.warning(
                        "| WARNING: ran out of memory, skipping batch. "
                        "if this happens frequently, decrease batch_size or "
                        "truncate the inputs to the model."
                    )
                    continue
                else:
                    raise e

        torch.cuda.empty_cache()

        bleu = batch_bleu(targets, generations, reduction="average")
        self._TB_validation_log.add_scalar("BLEU", bleu, self._validation_steps)
        try:
            P, R, F1 = bert_score(generations, targets)
        except:
            P, R, F1 = (torch.Tensor([0]), torch.Tensor([0]), torch.Tensor([0]))

        self._TB_validation_log.add_scalar(
            "bert-score-p", P.mean().item(), self._validation_counter
        )
        self._TB_validation_log.add_scalar(
            "bert-score-r", R.mean().item(), self._validation_counter
        )
        self._TB_validation_log.add_scalar(
            "bert-score-f1", F1.mean().item(), self._validation_counter
        )
        ppl = ppl / len(valid_iterator)
        # kld = kld / len(valid_iterator)
        self._TB_validation_log.add_scalar("kl", kld, self._validation_counter)

        # Had to do this for memory issues

        self._TB_validation_log.add_scalar("Perplexity", ppl, self._validation_steps)

        metric_dict = {"bleu": bleu, "perplexity": ppl, "kl": kld}

        bleu_best, bleu_patience = self._update_metric_history(
            self._validation_steps,
            "bleu",
            bleu,
            self._metric_infos,
            metric_decreases=False,
        )

        ppl_best, ppl_patience = self._update_metric_history(
            self._validation_steps,
            "perplexity",
            ppl,
            self._metric_infos,
            metric_decreases=True,
        )

        bert_score_p_best, bert_score_p_patience = self._update_metric_history(
            self._validation_steps,
            "bert_score_p",
            P.mean().item(),
            self._metric_infos,
            metric_decreases=False,
        )

        bert_score_r_best, bert_score_r_patience = self._update_metric_history(
            self._validation_steps,
            "bert_score_r",
            R.mean().item(),
            self._metric_infos,
            metric_decreases=False,
        )
        bert_score_f1_best, bert_score_f1_patience = self._update_metric_history(
            self._validation_steps,
            "bert_score_f1",
            F1.mean().item(),
            self._metric_infos,
            metric_decreases=False,
        )
        # kld_best, kld_patience = self._update_metric_history(
        #    self._epoch_steps,
        #    "KL Divergence",
        #    kld,
        #    self._metric_infos,
        #    metric_decreases=True,
        # )

        if self._keep_all_checkpoints:
            torch.save(
                self._model.state_dict(),
                os.path.join(
                    self._serialization_dir,
                    "model",
                    f"iter_{self._validation_steps}.pth",
                ),
            )
            torch.save(
                self._optimizer.state_dict(),
                os.path.join(
                    self._serialization_dir,
                    "optimizer",
                    f"iter_{self._validation_steps}.pth",
                ),
            )
        if bleu_best:
            torch.save(
                self._model.state_dict(),
                os.path.join(self._serialization_dir, "model", f"bleu_best.pth",),
            )
            torch.save(
                self._optimizer.state_dict(),
                os.path.join(self._serialization_dir, "optimizer", f"bleu_best.pth",),
            )
            self._bad_epochs = 0

        self._write_metric_info()

        if ppl_best:
            torch.save(
                self._model.state_dict(),
                os.path.join(self._serialization_dir, "model", f"ppl_best.pth",),
            )
            torch.save(
                self._optimizer.state_dict(),
                os.path.join(self._serialization_dir, "optimizer", f"ppl_best.pth",),
            )
            self._bad_epochs = 0
        if bert_score_f1_best:
            torch.save(
                self._model.state_dict(),
                os.path.join(self._serialization_dir, "model", f"bert_score_best.pth",),
            )
            torch.save(
                self._optimizer.state_dict(),
                os.path.join(
                    self._serialization_dir, "optimizer", f"bert_score_best.pth",
                ),
            )
            self._bad_epochs = 0
        if not bleu_best and not ppl_best and not bert_score_f1_best:
            self._bad_epochs += 1
        if bleu_patience and ppl_patience and bert_score_f1_patience:
            logging.info(
                "Ran out of patience for both BLEU, BERTScore and perplexity. Stopping"
                " training"
            )

            self._patience_exceeded = True
        with open(
            os.path.join(self._validation_log_dir, f"iter_{self._epoch_steps}.json"),
            "w",
        ) as f:
            f.write(
                "\n".join(
                    [
                        json.dumps(
                            {
                                "src": sources[i],
                                "tgt": targets[i],
                                "gen": generations[i],
                                "word": words[i],
                            }
                        )
                        for i in range(len(generations))
                    ]
                )
            )

        if self._bad_epochs == 3:
            tqdm.tqdm.write(
                "4 Bad Epochs in a row, reverting optimizer and model to the last"
                " bert_score_best.pth"
            )
            self._model.load_state_dict(
                torch.load(
                    os.path.join(
                        self._serialization_dir, "model", f"bert_score_best.pth"
                    )
                ),
            )
            self._optimizer.load_state_dict(
                torch.load(
                    os.path.join(
                        self._serialization_dir, "optimizer", f"bert_score_best.pth"
                    )
                )
            )
            for param_group in self._optimizer.param_groups:
                param_group["lr"] *= 0.5

        return DotMap({"src": sources, "tgt": targets, "gen": generations})

    def _test(self, batch_size):

        self.load_model()
        assert isinstance(self._model, torch.nn.Module), (
            "Before calling _validate, you must supply a PyTorch model using the"
            " `Trainer._set_model` method"
        )
        test_iterator = self._datamaker.get_iterator(
            "test", batch_size, device=self._device
        )

        generations = []
        targets = []
        sources = []
        words = []

        ppl = 0
        kld = 0
        for i, batch in enumerate(
            tqdm.tqdm(test_iterator, desc=f"Testing (Epoch {self._validation_steps}): ")
        ):
            try:
                self._model.zero_grad()
                self._model.eval()

                example, example_lens = batch.example
                definition, definition_lens = batch.definition
                word, word_lens = batch.word
                if self._model.variational or self._model.defbert:
                    definition_ae, definition_ae_lens = batch.definition_ae
                else:
                    definition_ae, definition_ae_lens = None, None

                sentence_mask = bert_dual_sequence_mask(
                    # example, self._datamaker.vocab.example.encode("</s>")[1:-1], device=self._device
                    # example, self._datamaker.vocab.example.encode(self._datamaker.vocab.example._sep_token)[1:-1], device=self._device
                    # example, self._tgt_bos_idx, device=self._device
                    example, self._datamaker.vocab.example.encode("[SEP]")[1:-1], device=self._device
                )
                current_batch_size = word.shape[0]

                decode_strategy = BeamSearch(
                    self._beam_size,
                    current_batch_size,
                    pad=self._tgt_pad_idx,
                    bos=self._tgt_bos_idx,
                    eos=self._tgt_eos_idx,
                    n_best=1 if self._n_best is None else self._n_best,
                    global_scorer=self._model.global_scorer,
                    min_length=self._min_length,
                    max_length=self._max_length,
                    return_attention=False,
                    block_ngram_repeat=3,
                    exclusion_tokens=self._exclusion_idxs,
                    stepwise_penalty=None,
                    ratio=self._ratio if self._ratio is not None else 0,
                )
                with torch.no_grad():
                    model_out = self._forward(
                        "test",
                        input=example,
                        seq_lens=example_lens,
                        span_token_ids=word,
                        target=definition,
                        target_lens=definition_lens,
                        decode_strategy=decode_strategy,
                        definition=definition_ae,
                        definition_lens=definition_ae_lens,
                        sentence_mask=sentence_mask,
                    )
                torch.cuda.empty_cache()

                generations.extend(
                    [
                        self._datamaker.decode(gen[0], "definition", batch=False)
                        for gen in model_out.predictions
                    ]
                )
                targets.extend(
                    self._datamaker.decode(definition, "definition", batch=True)
                )
                sources.extend(self._datamaker.decode(example, "example", batch=True))
                words.extend(self._datamaker.decode(word, "word", batch=True))

                ppl += model_out.perplexity.item()
                if model_out.kl is not None:
                    kld += model_out.kl.item()
                    self._TB_validation_log.add_scalar(
                        "kl", model_out.kl.item(), self._validation_counter
                    )

                current_bleu = batch_bleu(
                    targets[-current_batch_size:],
                    generations[-current_batch_size:],
                    reduction="average",
                )
            except RuntimeError as e:
                # catch out of memory exceptions during fwd/bck (skip batch)
                if "out of memory" in str(e):
                    logging.warning(
                        "| WARNING: ran out of memory, skipping batch. "
                        "if this happens frequently, decrease batch_size or "
                        "truncate the inputs to the model."
                    )
                    torch.cuda.empty_cache()
                    continue
                else:
                    raise e

        torch.cuda.empty_cache()

        bleu = batch_bleu(targets, generations, reduction="average")

        ppl = ppl / len(test_iterator)
        if self._model.variational:
            kld = kld / len(test_iterator)
            kld_best, kld_patience = self._update_metric_history(
                self._validation_steps,
                "KL Divergence",
                kld,
                self._test_metric_infos,
                metric_decreases=False,
            )

        metric_dict = {"bleu": bleu, "perplexity": ppl, "kl": kld}

        try:
            P, R, F1 = bert_score(generations, targets)
        except:
            P, R, F1 = (torch.Tensor([0]), torch.Tensor([0]), torch.Tensor([0]))

        bleu_best, bleu_patience = self._update_metric_history(
            self._validation_steps,
            "bleu",
            bleu,
            self._test_metric_infos,
            metric_decreases=False,
        )

        ppl_best, ppl_patience = self._update_metric_history(
            self._validation_steps,
            "perplexity",
            ppl,
            self._test_metric_infos,
            metric_decreases=True,
        )

        bert_score_p_best, bert_score_p_patience = self._update_metric_history(
            self._validation_steps,
            "bert_score_p",
            P.mean().item(),
            self._test_metric_infos,
            metric_decreases=False,
        )

        bert_score_r_best, bert_score_r_patience = self._update_metric_history(
            self._validation_steps,
            "bert_score_r",
            R.mean().item(),
            self._test_metric_infos,
            metric_decreases=False,
        )
        bert_score_f1_best, bert_score_f1_patience = self._update_metric_history(
            self._validation_steps,
            "bert_score_f1",
            F1.mean().item(),
            self._test_metric_infos,
            metric_decreases=False,
        )
        # kld_best, kld_patience = self._update_metric_history(
        #     self._epoch_steps,
        #     "KL Divergence",
        #     kld,
        #     self._test_metric_infos,
        #     metric_decreases=False,
        # )

        self._test_write_metric_info()

        with open(
            os.path.join(
                self._validation_log_dir, f"test_iter_{self._epoch_steps}.json"
            ),
            "w",
        ) as f:
            f.write(
                "\n".join(
                    [
                        json.dumps(
                            {
                                "src": sources[i],
                                "tgt": targets[i],
                                "gen": generations[i],
                                "word": words[i],
                            }
                        )
                        for i in range(len(generations))
                    ]
                )
            )

        return DotMap({"src": sources, "tgt": targets, "gen": generations})

    def _reset_steps(self):
        self._epoch_steps = 0
        self._validation_counter = 0
        self._train_counter = 0

    def _forward(self, phase: str = "train",  **batch):

        if phase not in ["train", "valid", "test"]:
            raise NotImplementedError(f"{phase} must be in ['train','test','valid']")
        if phase == "train":
            return self._model(device=self._device, **batch)
        elif phase in ["valid", "test"]:
            return self._model._validate(**batch)
        else:
            raise NotImplementedError

    def _set_model(self, model):
        self._model = model
        if self._load_model:
            self._model.load_state_dict(torch.load(self._load_model))

    def load_model(self):
        if self._load_model:
            self._model.load_state_dict(torch.load(self._load_model))
        else:
            self._model.load_state_dict(
                torch.load(self._serialization_dir + "/model/bert_score_best.pth")
            )

    def _write_metric_info(self):
        with open(os.path.join(self._serialization_dir, "metric.json"), "w") as fp:
            json.dump(self._metric_infos, fp)
        # TO BE REMOVED

    def _test_write_metric_info(self):
        with open(os.path.join(self._serialization_dir, "test_metric.json"), "w") as fp:
            json.dump(self._test_metric_infos, fp)
