#!/usr/bin/env python3 -u

from collections import namedtuple

import math
import torch
from torch.nn.utils.rnn import pad_sequence

from fairseq import options, tasks, utils
from eet.fairseq.transformer import EETTransformerDecoder


Batch = namedtuple('Batch', 'ids src_tokens src_lengths')

def make_batches(lines, task, max_positions, encode_fn):

    tokens = [task.source_dictionary.encode_line(encode_fn(line),
                                                 add_if_not_exist=False,
                                                 append_eos=False,
                                                 reverse_order=True).long()
              for line in lines]
    lengths = [t.numel() for t in tokens]
    tokens = pad_sequence(tokens, batch_first=True,
                          padding_value=1).flip(dims=(1,))

    return Batch(ids=torch.arange(len(tokens)),
                 src_tokens=tokens,
                 src_lengths=torch.tensor(lengths))


def encode_fn(x_str):
    x_str = x_str.replace(" ", "")
    x_str = x_str.split("</s>")
    x_str = " </s> ".join([" ".join(list(x)) for x in x_str])
    x_str = "</s> " + x_str
    return x_str


def decode_fn(x):
    x = x.replace(" ", "")
    return x


def eos_token_filter(sent):
    if "</s>" in sent:
        return True
    return False


def post_precess(line):
    line = "</s>".join(line.split("</s>")[:-1])
    return line


class Inference(object):

    def __init__(self, model_path, data_path, eet_batch_size):

        parser = options.get_generation_parser(
            default_task="language_modeling")
        args = options.parse_args_and_arch(parser)
        args.data = data_path
        args.path = model_path
        self.args = args

        # generate parameter
        args.beam = 1  # don't change
        args.min_len = 5
        args.max_len_b = 200
        args.lenpen = 1.0
        args.sampling = True
        args.sampling_topp = 0.8
        # args.sampling_topk = 20
        args.temperature = 0.8
        args.no_repeat_ngram_size = 1
        args.fp16 = True

        # Setup task, e.g., translation
        task = tasks.setup_task(args)
        self.task = task
        # Set dictionaries
        self.src_dict = task.source_dictionary
        self.tgt_dict = task.target_dictionary

        use_cuda = torch.cuda.is_available() and not args.cpu
        self.use_cuda = use_cuda

        model_path = args.path
        checkpoint = torch.load(model_path.replace("best.pt", "best_part_1.pt"))
        checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_2.pt")))
        checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_3.pt")))
        torch.save(checkpoint, model_path)

        state = torch.load(args.path, map_location=torch.device("cpu"))
        cfg_args = eval(str(state["cfg"]))["model"]
        del cfg_args["_name"]
        keys_list = []
        values_list = []
        for key, value in cfg_args.items():
            keys_list.append(key)
            values_list.append(value)
        Model_args = namedtuple("Model_args", keys_list)
        model_args = Model_args._make(values_list)
        del state

        eet_seq_len = 1024  # max sequence length, (input length + generation length) shouldn't be larger than this 
        eet_batch_size = eet_batch_size
        data_type = torch.float16
        eet_config = {"data_type": data_type,
                      "max_batch": eet_batch_size,
                      "full_seq_len": eet_seq_len}
        print(model_args)

        eet_model = EETTransformerDecoder.from_fairseq_pretrained(model_id_or_path=args.path,
                                                                  dictionary=self.src_dict, args=model_args,
                                                                  config=eet_config,
                                                                  no_encoder_attn=True)
        self.models = [eet_model]
        # Initialize generator
        self.generator = task.build_generator(self.models, args)

        # Load alignment dictionary for unknown word replacement
        # (None if no unknown word replacement, empty if no path to align dictionary)
        self.align_dict = utils.load_align_dict(args.replace_unk)

        self.max_positions = 1024  # the model config
        self.eos_index = self.tgt_dict.eos()
        self.pad_index = self.tgt_dict.pad()

    def __call__(self, inputs, append_right_eos=True):

        results = []
        start_id = 0

        batch = make_batches(inputs, self.task, self.max_positions, encode_fn)
        inputs_str = inputs

        src_tokens = batch.src_tokens
        src_lengths = batch.src_lengths
        # a new paragraph always
        if src_tokens[0][-1].item() != self.eos_index and append_right_eos:
            src_tokens = torch.cat([src_tokens, src_tokens.new_ones(
                src_tokens.size(0), 1) * self.eos_index], dim=1)
            src_lengths += 1
        if self.use_cuda:
            src_tokens = src_tokens.cuda()
            src_lengths = src_lengths.cuda()
        sample = {
            'net_input': {
                'src_tokens': src_tokens,
                'src_lengths': src_lengths,
            },
        }

        translations = self.task.inference_step(
            self.generator, self.models, sample)

        for i, (id, hypos) in enumerate(zip(batch.ids.tolist(), translations)):
            results.append((start_id + id, src_tokens[i], hypos))

        # sort output to match input order
        final_results = []
        for id, src_tokens, hypos in sorted(results, key=lambda x: x[0]):
            # Process top predictions
            tmp_res = []
            for hypo in hypos[:min(len(hypos), self.args.nbest)]:
                hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
                    hypo_tokens=hypo['tokens'].int().cpu()[
                        len(src_tokens) - 1:],
                    src_str=None,
                    alignment=hypo['alignment'],
                    align_dict=self.align_dict,
                    tgt_dict=self.tgt_dict)

                detok_hypo_str = decode_fn(hypo_str)
                if eos_token_filter(detok_hypo_str):
                    detok_hypo_str = post_precess(detok_hypo_str)
                    score = hypo['score'] / math.log(2)  # convert to base 2
                    tmp_res.append([detok_hypo_str, score])
            final_results.append(tmp_res)
        return final_results