#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from parlai.utils.distributed import is_distributed
from parlai.core.torch_ranker_agent import TorchRankerAgent
from parlai.utils.misc import padded_3d
from parlai.zoo.bert.build import download

from .bert_dictionary import BertDictionaryAgent
from .helpers import (get_bert_optimizer, BertWrapper, BertModel,
                      add_common_args, surround, MODEL_PATH,
                      get_bert_model)

import os
import torch
import pdb
from tqdm import tqdm


class BiEncoderRankerMlpAgent(TorchRankerAgent):
    """ TorchRankerAgent implementation of the biencoder.
        It is a standalone Agent. It might be called by the Both Encoder.
    """

    @staticmethod
    def add_cmdline_args(parser):
        add_common_args(parser)
        parser.add_argument('--mlp-num-hidden-layers', type=int, default=1,
                            help='Additional mlp hidden layers')
        parser.add_argument('--mlp-num-hidden-units', type=int, default=1536,
                            help='Additional mlp hidden units')
        parser.add_argument('--mlp-dropout', type=float, default=0.1,
                            help='Additional mlp dropout')
        parser.add_argument('--first-pass', type=bool, default=False,
                            help='Make a first pass where we only fine tuned'
                                 ' the upper layers.')
        parser.add_argument('--second-pass-lr', type=float, default=5e-5,
                            help='When no improvement can be obtained on the '
                                 'first pass, start the second pass.')
        parser.add_argument('--mlp-concat-mode', type=str, default="concat",
                            choices=["concat", "concat_hadamar", "hadamar"],
                            help='Additional mlp dropout')
        parser.set_defaults(
            encode_candidate_vecs=True
        )

    def __init__(self, opt, shared=None):
        # download pretrained models
        download(opt['datapath'])
        self.pretrained_path = os.path.join(opt['datapath'], 'models',
                                            'bert_models', MODEL_PATH)
        opt['pretrained_path'] = self.pretrained_path

        self.clip = -1

        super().__init__(opt, shared)
        # it's easier for now to use DataParallel when
        self.data_parallel = opt.get('data_parallel') and self.use_cuda
        if self.data_parallel:
            self.model = torch.nn.DataParallel(self.model)
        if is_distributed():
            raise ValueError('Cannot combine --data-parallel and distributed mode')
        self.set_special_idx()
        # default one does not average
        self.rank_loss = torch.nn.CrossEntropyLoss(reduce=True, size_average=True)
        self.is_first_pass = opt["first_pass"]
        self.previous_accuracy = -1e32

    def build_model(self):
        self.model = BiEncoderModule(self.opt)

    def set_special_idx(self):
        self.NULL_IDX = self.dict.pad_idx
        self.START_IDX = self.dict.start_idx
        self.END_IDX = self.dict.end_idx

    def receive_metrics(self, metrics_dict):
        if self.is_first_pass:
            # TODO checkout what is the metric that is used here.
            if metrics_dict['accuracy'] < self.previous_accuracy:
                print("End of the first pass")
                for group in self.optimizer.param_groups:
                    group['lr'] = self.opt['second_pass_lr']
                self.is_first_pass = False
            else:
                self.previous_accuracy = metrics_dict['accuracy']
                super().receive_metrics(metrics_dict)
        else:
            super().receive_metrics(metrics_dict)

    @staticmethod
    def dictionary_class():
        return BertDictionaryAgent

    def init_optim(self, params, optim_states=None, saved_optim_type=None):
        if self.opt['classic_parlai_optimizer']:
            super().init_optim(params, optim_states, saved_optim_type)
        else:
            self.optimizer = get_bert_optimizer([self.model],
                                                self.opt['type_optimization'],
                                                self.opt['learningrate'],
                                                fp16=self.opt.get('fp16'))

    def set_vocab_candidates(self, shared):
        """Load the tokens from the vocab as candidates

        self.vocab_candidates will contain a [num_cands] list of strings
        self.vocab_candidate_vecs will contain a [num_cands, 1] LongTensor
        """
        self.opt['encode_candidate_vecs'] = True
        if shared:
            self.vocab_candidates = shared['vocab_candidates']
            self.vocab_candidate_vecs = shared['vocab_candidate_vecs']
            self.vocab_candidate_encs = shared['vocab_candidate_encs']
        else:
            if 'vocab' in (self.opt['candidates'], self.opt['eval_candidates']):
                cands = []
                vecs = []
                for ind in range(1, len(self.dict)):
                    txt = self.dict[ind]
                    cands.append(txt)
                    vecs.append(
                        self._vectorize_text(txt, add_start=True, add_end=True,
                                             truncate=self.label_truncate)
                    )
                self.vocab_candidates = cands
                self.vocab_candidate_vecs = padded_3d([vecs]).squeeze(0)
                print("[ Loaded fixed candidate set (n = {}) from vocabulary ]"
                      "".format(len(self.vocab_candidates)))
                enc_path = self.opt.get('model_file') + '.vocab.encs'
                if os.path.isfile(enc_path):
                    self.vocab_candidate_encs = self.load_candidates(
                        enc_path, cand_type='vocab encodings')
                else:
                    cand_encs = []
                    vec_batches = [
                        self.vocab_candidate_vecs[i:i + 512] for i in
                        range(0, len(self.vocab_candidate_vecs), 512)
                    ]
                    print("[ Vectorizing vocab candidates ({} batch(es) of up "
                          "to 512) ]".format(len(vec_batches)))
                    for vec_batch in tqdm(vec_batches):
                        cand_encs.append(self.encode_candidates(vec_batch))
                    self.vocab_candidate_encs = torch.cat(cand_encs, 0)
                    self.save_candidates(
                        self.vocab_candidate_encs, enc_path,
                        cand_type='vocab encodings')
                if self.use_cuda:
                    self.vocab_candidate_vecs = self.vocab_candidate_vecs.cuda()
                    self.vocab_candidate_encs = self.vocab_candidate_encs.cuda()
            else:
                self.vocab_candidates = None
                self.vocab_candidate_vecs = None
                self.vocab_candidate_encs = None

    def vectorize_fixed_candidates(self, cands_batch):
        """Override from TorchRankerAgent.
        """
        return [self._vectorize_text(cand, add_start=True, add_end=True,
                truncate=self.label_truncate) for cand in cands_batch]

    def encode_candidates(self, padded_cands):
        token_idx_cands, segment_idx_cands, mask_cands = to_bert_input(
            padded_cands, self.NULL_IDX)
        _, embedding_cands = self.model(
            None, None, None, token_idx_cands, segment_idx_cands, mask_cands)

        return embedding_cands.cpu().detach()

    def _set_text_vec(self, *args, **kwargs):
        obs = super()._set_text_vec(*args, **kwargs)
        # concatenate the [CLS] and [SEP] tokens
        if obs is not None and 'text_vec' in obs:
            obs['text_vec'] = surround(obs['text_vec'], self.START_IDX,
                                       self.END_IDX)
        return obs

    def score_candidates(self, batch, cand_vecs, cand_encs=None):
        # Encode contexts first
        token_idx_ctxt, segment_idx_ctxt, mask_ctxt = to_bert_input(
            batch.text_vec, self.NULL_IDX)
        embedding_ctxt, _, _ = self.model(
            token_idx_ctxt, segment_idx_ctxt, mask_ctxt,
            None, None, None)
        if self.is_first_pass:
            embedding_ctxt = embedding_ctxt.detach()

        if cand_encs is not None:
            if len(cand_encs.size()) == 2:
                cand_encs = cand_encs.unsqueeze(0)
            _, _, scores = self.model(None, None, None, None, None, None,
                                      embedding_ctxt, cand_encs)
            return scores

        if len(cand_vecs.size()) == 2 and cand_vecs.dtype == torch.long:
            # train time. We compare with all elements of the batch
            token_idx_cands, segment_idx_cands, mask_cands = to_bert_input(
                cand_vecs, self.NULL_IDX)
            _, embedding_cands, _ = self.model(
                None, None, None,
                token_idx_cands, segment_idx_cands, mask_cands)
            bsz = embedding_ctxt.size(0)
            cand_size = cand_vecs.size(0)
            embedding_cands = embedding_cands.unsqueeze(0).expand(bsz, cand_size, -1)
            if self.is_first_pass:
                embedding_cands = embedding_cands.detach()
            _, _, scores = self.model(None, None, None, None, None, None,
                                      embedding_ctxt, embedding_cands)
            return scores

        # predict time with multiple candidates
        if len(cand_vecs.size()) == 3:
            csize = cand_vecs.size()  # batchsize x ncands x sentlength
            cands_idx_reshaped = cand_vecs.view(csize[0] * csize[1], csize[2])
            token_idx_cands, segment_idx_cands, mask_cands = to_bert_input(
                cands_idx_reshaped, self.NULL_IDX)
            _, embedding_cands, _ = self.model(
                None, None, None,
                token_idx_cands, segment_idx_cands, mask_cands)
            embedding_cands = embedding_cands.view(csize[0], csize[1], -1)
            if self.is_first_pass:
                embedding_cands = embedding_cands.detach()
            _, _, scores = self.model(None, None, None, None, None, None,
                                      embedding_ctxt, embedding_cands)
            return scores

        # otherwise: cand_vecs should be 2D float vector ncands x embed_size
        _, _, scores = self.model(None, None, None, None, None, None,
                                  embedding_ctxt, cand_vecs)
        return scores

    def share(self):
        """Share model parameters."""
        shared = super().share()
        shared['vocab_candidate_encs'] = self.vocab_candidate_encs
        return shared


class BiEncoderModule(torch.nn.Module):
    """ Groups context_encoder and cand_encoder together.
    """

    def __init__(self, opt):
        super(BiEncoderModule, self).__init__()
        self.opt = opt
        self.context_encoder = BertWrapper(
            get_bert_model(opt),
            opt['out_dim'],
            add_transformer_layer=opt['add_transformer_layer'],
            layer_pulled=opt['pull_from_layer'],
            aggregation=opt['bert_aggregation'],
            add_linear_layer=False,
        )
        self.cand_encoder = BertWrapper(
            get_bert_model(opt),
            opt['out_dim'],
            add_transformer_layer=opt['add_transformer_layer'],
            layer_pulled=opt['pull_from_layer'],
            aggregation=opt['bert_aggregation'],
            add_linear_layer=False,
        )
        n_layers = opt['mlp_num_hidden_layers']
        hidden_dim = opt['mlp_num_hidden_units']
        dropout = opt['mlp_dropout']
        input_dim = 2 * opt['out_dim']
        if opt["mlp_concat_mode"] == "concat_hadamar":
            input_dim = 3 * opt['out_dim']
        if opt["mlp_concat_mode"] == "hadamar":
            input_dim = opt['out_dim']
        layers = [torch.nn.Linear(input_dim, hidden_dim),
                  torch.nn.ReLU(),
                  torch.nn.Dropout(dropout)]
        for _ in range(n_layers - 1):
            layers += [torch.nn.Linear(hidden_dim, hidden_dim),
                       torch.nn.ReLU(),
                       torch.nn.Dropout(dropout)]
        layers +=[torch.nn.Linear(hidden_dim, 1)]
        self.additional_mlp = torch.nn.Sequential(*layers)


    def get_score(self, emb_ctxt, emb_cands):
        """ None of input are None.
            Returns the score
            emb_ctxt is (batchsize x emb_size)
            emb_cands is (batchsize x num_candidates x emb_size)
        """
        bsz = emb_ctxt.size(0)
        num_cands = emb_cands.size(1)
        emb_ctxt = emb_ctxt.unsqueeze(1).expand(bsz, num_cands, -1)

        # now emb_ctxt and emb_cands are both (batchsize x num_candidates
        # x emb_size) and we can sum them or other things
        emb_concat = None
        if self.opt["mlp_concat_mode"] == "concat":
            emb_concat = torch.cat([emb_ctxt, emb_cands], 2)
        elif self.opt["mlp_concat_mode"] == "concat_hadamar":
            hadamar = emb_ctxt * emb_cands
            emb_concat = torch.cat([emb_ctxt, emb_cands, hadamar], 2)
        elif self.opt["mlp_concat_mode"] == "hadamar":
            emb_concat = emb_ctxt * emb_cands
        scores = self.additional_mlp(emb_concat).squeeze(2)
        return scores

    def forward(self, token_idx_ctxt, segment_idx_ctxt, mask_ctxt,
                token_idx_cands, segment_idx_cands, mask_cands,
                embedding_ctxt=None, embedding_cands=None):
        # In order to bypass the "only one model philosophy" of parlai
        # I had to make 2 forward pass
        if embedding_ctxt is not None and embedding_cands is not None:
            return None, None, self.get_score(embedding_ctxt, embedding_cands)

        # First, encode if it's not done already
        embedding_ctxt = embedding_ctxt
        if token_idx_ctxt is not None and embedding_ctxt is None:
            embedding_ctxt = self.context_encoder(
                token_idx_ctxt, segment_idx_ctxt, mask_ctxt)
        embedding_cands = embedding_cands
        if token_idx_cands is not None and embedding_cands is None:
            embedding_cands = self.cand_encoder(
                token_idx_cands, segment_idx_cands, mask_cands)

        return embedding_ctxt, embedding_cands, None


def to_bert_input(token_idx, null_idx):
    """ token_idx is a 2D tensor int.
        return token_idx, segment_idx and mask
    """
    segment_idx = token_idx * 0
    mask = (token_idx != null_idx)
    # nullify elements in case self.NULL_IDX was not 0
    token_idx = token_idx * mask.long()
    return token_idx, segment_idx, mask
