"""
Based on https://github.com/huggingface/transformers
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import logging
import argparse

import torch
import numpy as np
import torch.nn.functional as F

from tqdm import trange
from pathlib import Path
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from transformers import XLNetLMHeadModel, XLNetTokenizer
from transformers import TransfoXLLMHeadModel, TransfoXLTokenizer
from transformers import OpenAIGPTLMHeadModel, OpenAIGPTTokenizer
from transformers import GPT2Config, OpenAIGPTConfig, XLNetConfig, TransfoXLConfig


logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt = '%m/%d/%Y %H:%M:%S',
                    level = logging.INFO)

logger = logging.getLogger(__name__)

MAX_LENGTH = int(10000)  # Hardcoded max length to avoid infinite loop
ALL_MODELS = sum((tuple(conf.pretrained_config_archive_map.keys())
                  for conf in (GPT2Config, OpenAIGPTConfig, XLNetConfig, TransfoXLConfig)), ())

MODEL_CLASSES = {
    'gpt2': (GPT2LMHeadModel, GPT2Tokenizer),
    'gpt2-medium': (GPT2LMHeadModel, GPT2Tokenizer),
    'gpt2-large': (GPT2LMHeadModel, GPT2Tokenizer),
    'gpt2-xl': (GPT2LMHeadModel, GPT2Tokenizer),
    'openai-gpt': (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
    'xlnet': (XLNetLMHeadModel, XLNetTokenizer),
    'transfo-xl': (TransfoXLLMHeadModel, TransfoXLTokenizer)
}

# Padding text to help Transformer-XL and XLNet with short prompts as proposed by Aman Rusia
# in https://github.com/rusiaaman/XLNet-gen#methodology
# and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
PADDING_TEXT = """ In 1991, the remains of Russian Tsar Nicholas II and his family
(except for Alexei and Maria) are discovered.
The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the
remainder of the story. 1883 Western Siberia,
a young Grigori Rasputin is asked by his father and a group of men to perform magic.
Rasputin has a vision and denounces one of the men as a horse thief. Although his
father initially slaps him for making such an accusation, Rasputin watches as the
man is chased outside and beaten. Twenty years later, Rasputin sees a vision of
the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous,
with people, even a bishop, begging for his blessing. <eod> </s> <eos>"""


def set_seed(args):
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
        Args:
            logits: logits distribution shape (batch size x vocabulary size)
            top_k > 0: keep only top k tokens with highest probability (top-k filtering).
            top_p > 0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
                Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
        From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
    """
    top_k = min(top_k, logits.size(-1))  # Safety check
    if top_k > 0:
        # Remove all tokens with a probability less than the last token of the top-k
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens with cumulative probability above the threshold
        sorted_indices_to_remove = cumulative_probs > top_p
        # Shift the indices to the right to keep also the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        # scatter sorted tensors to original indexing
        indices_to_remove = sorted_indices_to_remove.scatter(dim=1, index=sorted_indices, src=sorted_indices_to_remove)
        logits[indices_to_remove] = filter_value
    return logits


def sample_sequence(model, length, context, num_samples=1, temperature=1, top_k=0, top_p=0.0, repetition_penalty=1.0,
                    is_xlnet=False, device='cpu', next_word_prob=False):
    if next_word_prob:
        length = 1
        temperature = 0
    context = torch.tensor(context, dtype=torch.long, device=device)
    context = context.unsqueeze(0).repeat(num_samples, 1)
    generated = context
    with torch.no_grad():
        for _ in trange(length):

            inputs = {'input_ids': generated}
            if is_xlnet: 
                # XLNet is a direct (predict same token, not next token) and bi-directional model by default
                # => need one additional dummy token in the input (will be masked),
                # attention mask and target mapping (see model docstring)
                input_ids = torch.cat((generated, torch.zeros((1, 1), dtype=torch.long, device=device)), dim=1)
                perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float, device=device)
                perm_mask[:, :, -1] = 1.0  # Previous tokens don't see last token
                target_mapping = torch.zeros((1, 1, input_ids.shape[1]), dtype=torch.float, device=device)
                target_mapping[0, 0, -1] = 1.0  # predict last token
                inputs = {'input_ids': input_ids, 'perm_mask': perm_mask, 'target_mapping': target_mapping}

            outputs = model(**inputs)
            next_token_logits = outputs[0][:, -1, :] / (temperature if temperature > 0 else 1.)

            # get pure next-token distribution, no reweighting
            if next_word_prob:
                next_token_logits_pure = outputs[0][:, -1, :]
                sorted_logits, sorted_indices = torch.sort(next_token_logits_pure, descending=True)
                return sorted_logits, sorted_indices

            # repetition penalty from CTRL (https://arxiv.org/abs/1909.05858)
            for i in range(num_samples):
                for _ in set(generated[i].tolist()):
                    next_token_logits[i, _] /= repetition_penalty
                
            filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)
            if temperature == 0:  # greedy sampling:
                next_token = torch.argmax(filtered_logits, dim=-1).unsqueeze(-1)
            else:
                next_token = torch.multinomial(F.softmax(filtered_logits, dim=-1), num_samples=1)
            generated = torch.cat((generated, next_token), dim=1)
    return generated


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_type", default=None, type=str, required=True,
                        help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()))
    parser.add_argument("--model_name_or_path", default=None, type=str, required=True,
                        help="Path to pre-trained model or shortcut name selected in the list: " +
                             ", ".join(ALL_MODELS))
    parser.add_argument("--prompt", type=str, default="")
    parser.add_argument("--padding_text", type=str, default="")
    parser.add_argument("--length", type=int, default=20)
    parser.add_argument("--num_samples", type=int, default=1)
    parser.add_argument("--temperature", type=float, default=1.0,
                        help="temperature of 0 implies greedy sampling")
    parser.add_argument("--repetition_penalty", type=float, default=1.0,
                        help="primarily useful for CTRL model; in that case, use 1.2")
    parser.add_argument("--top_k", type=int, default=0)
    parser.add_argument("--top_p", type=float, default=0.9)
    parser.add_argument("--no_cuda", action='store_true',
                        help="Avoid using CUDA when available")
    parser.add_argument('--seed', type=int, default=42,
                        help="random seed for initialization")
    parser.add_argument('--stop_token', type=str, default=None,
                        help="Token at which text generation is stopped")
    parser.add_argument('--bert', action='store_true', help="Use BERT model for next_word_prob experiment.")
    parser.add_argument('--next_word_prob', help="Next-word probability experiment.", action='store_true')
    parser.add_argument('--names_file', help="Names tsv for next word probability experiment.", type=str,
                        default="./NewsNames/all_names.tsv")
    parser.add_argument('--context_type', choices=["news", "history", "personal", "none"], default="none",
                        help="Implicit domain of context/prefix.", type=str)
    args = parser.parse_args()

    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
    args.n_gpu = torch.cuda.device_count()

    set_seed(args)

    args.model_type = args.model_type.lower()
    model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
    tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
    model = model_class.from_pretrained(args.model_name_or_path)
    model.to(args.device)
    model.eval()

    if args.length < 0 and model.config.max_position_embeddings > 0:
        args.length = model.config.max_position_embeddings
    elif 0 < model.config.max_position_embeddings < args.length:
        args.length = model.config.max_position_embeddings  # No generation bigger than model size 
    elif args.length < 0:
        args.length = MAX_LENGTH  # avoid infinite loop

    logger.info(args)

    # Set the context prefix based on specified domain
    if args.context_type == "news":
        context_string = "A new report from CNN says that "
    elif args.context_type == "history":
        context_string = "A newly published biography of "
    elif args.context_type == "personal":
        context_string = "I want to introduce you to my best friend, "
    else:
        assert args.context_type == "none"
        context_string = ""

    if args.next_word_prob:
        contents = Path(args.names_file).read_text().rstrip()
        lines = contents.split('\n')
        rows = [L.split('\t') for L in lines]
        AE = args.names_file.endswith("ae_millennium.tsv")
        # Name    count   Prefix  Suffix  Simple
        # FullName  First   Last

    for i, row in enumerate(rows):
        if i == 0:
            setting = args.model_name_or_path + "." + args.context_type
            new_fields = ["Argmax." + setting, "GreedySeq." + setting, "Prob." + setting, "Ent." + setting]
            print("\t".join(new_fields))
            continue

        first_name = row[2]
        assert len(row) == 6
        raw_text = context_string + first_name

        if args.model_type in ["transfo-xl", "xlnet"]:
            # Models with memory likes to have a long prompt for short inputs.
            raw_text = (args.padding_text if args.padding_text else PADDING_TEXT) + raw_text

        context_tokens = tokenizer.encode(raw_text, add_special_tokens=False)

        if args.bert:
            raise NotImplementedError
        else:
            out = sample_sequence(
                model=model,
                context=context_tokens,
                num_samples=args.num_samples,
                length=args.length,
                temperature=args.temperature,
                top_k=args.top_k,
                top_p=args.top_p,
                repetition_penalty=args.repetition_penalty,
                is_xlnet=bool(args.model_type == "xlnet"),
                device=args.device,
                next_word_prob=args.next_word_prob
            )
        if args.next_word_prob:
            # Name    count   Prefix  Suffix  Simple
            sorted_logits, sorted_indices = out
            sorted_probs = F.softmax(sorted_logits, dim=-1)
            full_seq = sample_sequence(
                model=model,
                context=context_tokens,
                num_samples=args.num_samples,
                length=5,
                temperature=0,
                top_k=args.top_k,
                top_p=args.top_p,
                repetition_penalty=args.repetition_penalty,
                is_xlnet=bool(args.model_type == "xlnet"),
                device=args.device,
                next_word_prob=False
            )
            full_seq = full_seq[:, len(context_tokens):].tolist()
            for o in full_seq:
                text = tokenizer.decode(o, clean_up_tokenization_spaces=True)
                text = text[: text.find(args.stop_token) if args.stop_token else None]

            # remove whitespace that messes with formatting, don't need anyway
            decoded_text = text.replace('\t', '').replace('\n', '').replace('\r', '')

            ent = round(torch.distributions.Categorical(probs=sorted_probs).entropy().tolist()[0], 3)
            top_word = tokenizer.decode(sorted_indices.tolist()[0][0])
            top_prob = round(sorted_probs.tolist()[0][0], 3)
            print("\t".join([top_word, decoded_text, str(top_prob), str(ent)]))
        else:
            out = out[:, len(context_tokens):].tolist()
            for o in out:
                text = tokenizer.decode(o, clean_up_tokenization_spaces=True)
                text = text[: text.find(args.stop_token) if args.stop_token else None]
                print(text.replace('\t', ''))

        if args.prompt:
            break

    if args.next_word_prob:
        return 0
    return text


if __name__ == '__main__':
    main()
