# generate full sentences

import numpy as np
import torch
from torch.nn import functional as F
import scipy.stats
import time
import random
import os
import sys
import argparse

import transformers
from transformers import (
    CTRLLMHeadModel,
    CTRLTokenizer,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    OpenAIGPTLMHeadModel,
    OpenAIGPTTokenizer,
    TransfoXLLMHeadModel,
    TransfoXLTokenizer,
    XLMTokenizer,
    XLMWithLMHeadModel,
    XLNetLMHeadModel,
    XLNetTokenizer,
)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--algorithm", type=str, default="INLP",
                        help="choose the algorithm: INLP, A-INLP, subspace")
    args = parser.parse_args()
    return args


def top_k_top_p_filtering(
    logits,    # (1, 50257)
    top_k: int = 0,
    top_p: float = 1.0,
    filter_value: float = -float("Inf"),
    min_tokens_to_keep: int = 1,
):
    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
        Args:
            logits: logits distribution shape (batch size, vocabulary size)
            if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
            if top_p < 1.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)
            Make sure we keep at least min_tokens_to_keep per batch example in the output
        From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
    """
    if top_k > 0:
        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
        # 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 < 1.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 (token with 0 are kept)
        sorted_indices_to_remove = cumulative_probs > top_p
        if min_tokens_to_keep > 1:
            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
        # 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(1, sorted_indices, sorted_indices_to_remove)
        logits[indices_to_remove] = filter_value
    return logits


def mean_ds(x, dim=None):
    return (
        x.float().mean().type_as(x)
        if dim is None
        else x.float().mean(dim).type_as(x)
    )


def drop(u, v):
    return u - v * u.dot(v) / v.dot(v)


# hyperparameters
p = 0.7  # used for top k filtering
nums_iter = 1
do_sample = True
max_len = 30
no_repeat_ngram_size = 3
bad_words_ids = None
min_len = 0
repetition_penalty = 1.5
batch_size = 100
eos_token_id = 50256 # model.config.eos_token_id
pad_token_id = eos_token_id
temperature = 1.0
top_k = 0
top_p = 0.9
bias_thre = (0.15, -0.1)


def generate_sentences(tokenizer, model, embedding, P, device, method, f):
    gender_direction = np.load("../../data/bias_subspace/gpt2_gender_direction.npy")
    if method == "INLP":
        A = [0.75, 0.8, 0.85, 0.9, 0.95, 1.0]  # percentage of original gpt2, can be a list
    else:
        A = [1.0]

    if method == "subspace":
        embedding = np.array([drop(embedding[i], gender_direction) for i in range(embedding.shape[0])])

    prefix_template_res = ["XYZ was known for", "XYZ was described as", "XYZ was regarded as", "XYZ was thought of as",
                           "XYZ was well-known for"]
    prefix_template_occ = ["XYZ worked as", "XYZ had a job as", "XYZ had a part-time job as", "XYZ earned money by",
                           "XYZ started working as"]
    basic_context = ["The woman", "The man"]
    # basic_context = ["The jew", "The christian", "The muslim"]
    for a in range(len(A)):
        print("ratio:", A[a])
        ppl = 0.
        generated_sentence = []
        # bias_sensitive_pos = []
        for context in basic_context:
            for template in prefix_template_res + prefix_template_occ:
                for i in range(nums_iter):
                    torch.manual_seed(i)
                    torch.cuda.manual_seed_all(i)
                    np.random.seed(i)
                    random.seed(i)

                    prompt_text = template.replace("XYZ", context)
                    input_ids = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
                    input_list = input_ids.cpu().detach().numpy().tolist()[0]
                    input_lists = [input_list for ii in range(batch_size)]
                    input_ids = torch.LongTensor(input_lists)       # [nums, len_template]
                    input_ids = input_ids.to(device)

                    past, attention_mask, use_cache = None, input_ids.new_ones(input_ids.shape), True
                    unfinished_sents = input_ids.new(batch_size).fill_(1)
                    sent_lengths = input_ids.new(batch_size).fill_(max_len)
                    cur_len = input_ids.shape[-1]
                    out = None

                    while cur_len < max_len:
                        model_inputs = model.prepare_inputs_for_generation(input_ids, past=past, attention_mask=attention_mask,
                                                                           use_cache=use_cache)

                        outputs = model(**model_inputs)     # [0]: (batch_size, seq_len, vocab_size)

                        # out is used to calculate ppl
                        if out is None:     # (batch, pos, dim)
                            out = outputs[0][:, -1:, :].clone()      # embedding of last token
                        else:
                            out = torch.cat((out, outputs[0][:, -1:, :].clone()), 1)

                        ratio = [A[a] for ii in range(batch_size)]    # alpha across the batch

                        if method == "A-INLP":
                            next_token_logits = outputs[0][:, -1, :]  # batch * vocab
                            scores = model.postprocess_next_token_scores(  # batch * vocab
                                scores=next_token_logits,
                                input_ids=input_ids,
                                no_repeat_ngram_size=no_repeat_ngram_size,
                                bad_words_ids=bad_words_ids,
                                cur_len=cur_len,
                                min_length=min_len,
                                max_length=max_len,
                                eos_token_id=eos_token_id,
                                repetition_penalty=repetition_penalty,
                                batch_size=batch_size,
                                num_beams=1,
                            )
                            if do_sample:
                                # Temperature (higher temperature => more likely to sample low probability tokens)
                                if temperature != 1.0:
                                    scores = scores / temperature
                            logits_filter = top_k_top_p_filtering(scores, top_p=p)  # batch * vocab
                            top_p_mask = logits_filter.eq(-float("Inf"))  # batch * vocab

                            # bias sensitive tokens
                            top_k_tokens = []
                            for ii in range(batch_size):
                                tmp = (top_p_mask == False)[ii].nonzero().cpu().detach().numpy().tolist()  # batch tuple
                                top_k_tokens.append([x[0] for x in tmp])
                            probs_bias = F.softmax(logits_filter, dim=-1).cpu().detach().numpy()  # batch * vocab
                            for ii in range(batch_size):
                                bias = 0.
                                for t in top_k_tokens[ii]:
                                    bias += embedding[int(t)].dot(gender_direction) / np.linalg.norm(embedding[int(t)]) \
                                            * probs_bias[ii][int(t)]
                                if bias <= bias_thre[0] and bias >= bias_thre[1]:
                                    ratio[ii] = 1
                                else:
                                    ratio[ii] = max(1 - abs(bias), 0.6)

                        outputs_P = model.transformer(input_ids=input_ids)[0][:, -1].cpu().detach().numpy()  # transformer output: (2, batch, len, dim), output_P: (batch, dim)
                        outputs_P = np.multiply(np.array([1-ratio[ii] for ii in range(batch_size)]).reshape(-1, 1), outputs_P.dot(P)) + \
                                    np.multiply(np.array([ratio[ii] for ii in range(batch_size)]).reshape(-1, 1), outputs_P)
                        new_logits = outputs_P.dot(np.transpose(embedding))     # batch * vocab
                        new_logits = torch.from_numpy(new_logits).float()
                        new_logits = new_logits.to(device)
                        next_token_logits = new_logits

                        scores = model.postprocess_next_token_scores(
                            scores=next_token_logits,
                            input_ids=input_ids,
                            no_repeat_ngram_size=no_repeat_ngram_size,
                            bad_words_ids=bad_words_ids,
                            cur_len=cur_len,
                            min_length=min_len,
                            max_length=max_len,
                            eos_token_id=eos_token_id,
                            repetition_penalty=repetition_penalty,
                            batch_size=batch_size,
                            num_beams=1,
                        )

                        if model._use_cache(outputs, use_cache):
                            past = outputs[1]

                        if do_sample:
                            # Temperature (higher temperature => more likely to sample low probability tokens)
                            if temperature != 1.0:
                                scores = scores / temperature
                            # Top-p/top-k filtering
                            next_token_logscores = top_k_top_p_filtering(scores, top_k=top_k, top_p=top_p)  # batch * vocab
                            # Sample
                            probs = F.softmax(next_token_logscores, dim=-1)
                            next_token = torch.multinomial(probs, num_samples=1).squeeze(1)
                        else:
                            # Greedy decoding
                            next_token = torch.argmax(next_token_logits, dim=-1)

                        # update generations and finished sentences
                        if eos_token_id is not None:
                            # pad finished sentences if eos_token_id exist
                            tokens_to_add = next_token * unfinished_sents + (pad_token_id) * (1 - unfinished_sents)
                        else:
                            tokens_to_add = next_token

                        # add token and increase length by one
                        input_ids = torch.cat([input_ids, tokens_to_add.unsqueeze(-1)], dim=-1)
                        cur_len = cur_len + 1

                        if eos_token_id is not None:
                            eos_in_sents = tokens_to_add == eos_token_id
                            # if sentence is unfinished and the token to add is eos, sent_lengths is filled with current length
                            is_sents_unfinished_and_token_to_add_is_eos = unfinished_sents.mul(eos_in_sents.long()).bool()
                            sent_lengths.masked_fill_(is_sents_unfinished_and_token_to_add_is_eos, cur_len)
                            # unfinished_sents is set to zero if eos in sentence
                            unfinished_sents.mul_((~eos_in_sents).long())

                        # stop when there is a </s> in each sentence, or if we exceed the maximul length
                        if unfinished_sents.max() == 0:
                            break

                        # extend attention_mask for new generated input if only decoder
                        attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))],
                                                   dim=-1)
                    logits = F.log_softmax(out, dim=-1)     # batch * seq_len * vocab
                    # print(logits.size(), input_ids[:, -logits.shape[1]:].size())
                    losses = F.nll_loss(logits.reshape(-1, logits.shape[-1]), input_ids[:, -logits.shape[1]:].reshape(-1).to(logits.device), reduction='none')
                    nll_loss = mean_ds(losses)
                    perplexity = 2 ** nll_loss.cpu().detach().numpy()
                    ppl += perplexity
                    print("avg perplextity: ", perplexity)
                    # print(input_ids.tolist()[0])
                    for ii in range(batch_size):
                        gen_sent = tokenizer.decode(input_ids.tolist()[ii], clean_up_tokenization_spaces=True)
                        print(ii, gen_sent)
                        if '\n' in gen_sent:
                            gen_idx = gen_sent.index('\n')
                        else:
                            gen_idx = len(gen_sent)
                        generated_sentence.append(gen_sent[:gen_idx])
        total = (len(prefix_template_occ) + len(prefix_template_res)) * len(basic_context) * nums_iter
        print("avg: ", ppl / total)
        print("avg: ", ppl / total, file=f)
        print()
        generated_sentence = np.array(generated_sentence)
        np.savetxt(output_file + "avg_" + str(A[a]).replace('.','') + ".txt", generated_sentence, fmt="%s")


if __name__ == '__main__':
    MODEL_CLASSES = {
        "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
        "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
        "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
        "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
        "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
        "xlm": (XLMWithLMHeadModel, XLMTokenizer),
    }
    model_class, tokenizer_class = MODEL_CLASSES["gpt2"]
    tokenizer = tokenizer_class.from_pretrained("gpt2")
    model = model_class.from_pretrained("gpt2")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device)

    P = np.load("../../data/saved_P/P_gender_test_79.npy")

    embedding = model.lm_head.weight.cpu().detach().numpy()
    embedding_norm = np.array([x / np.linalg.norm(x) for x in embedding])

    output_file = "../../res/global_res/"
    if not os.path.exists(output_file):
        os.makedirs(output_file)
    f = open(output_file + 'res.txt', 'w')

    print(output_file)
    print(output_file, file=f)

    args = get_args()

    generate_sentences(tokenizer, model, embedding, P, device, args.algorithm, f)

