import os
import argparse
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import numpy as np
import warnings
warnings.filterwarnings('ignore')


def choose_from_top_k_top_n(probs, k=50, p=0.8):
    ind = np.argpartition(probs, -k)[-k:]
    top_prob = probs[ind]
    top_prob = {i: top_prob[idx] for idx, i in enumerate(ind)}
    sorted_top_prob = {
        k: v
        for k, v in sorted(
            top_prob.items(), key=lambda item: item[1], reverse=True)
    }

    t = 0
    f = []
    pr = []
    for k, v in sorted_top_prob.items():
        t += v
        f.append(k)
        pr.append(v)
        if t >= p:
            break
    top_prob = pr / np.sum(pr)
    token_id = np.random.choice(f, 1, p=top_prob)

    return int(token_id)


def generate(tokenizer, model, sentences, label):
    with torch.no_grad():
        for idx in range(sentences):
            finished = False
            cur_ids = torch.tensor(
                tokenizer.encode(label)).unsqueeze(0).to('cpu')
            for i in range(100):
                outputs = model(cur_ids, labels=cur_ids)
                loss, logits = outputs[:2]

                softmax_logits = torch.softmax(logits[0, -1], dim=0)

                if i < 5:
                    n = 10
                else:
                    n = 5

                next_token_id = choose_from_top_k_top_n(
                    softmax_logits.to('cpu').numpy())  #top-k-top-n sampling
                DEVICE = 'cpu'
                if torch.cuda.is_available():
                    DEVICE = 'cuda'
                cur_ids = torch.cat([
                    cur_ids,
                    torch.ones((1, 1)).long().to(DEVICE) * next_token_id
                ],
                                    dim=1)

                if next_token_id in tokenizer.encode('<|endoftext|>'):
                    finished = True
                    break

            if finished:
                output_list = list(cur_ids.squeeze().to('cpu').numpy())
                output_text = tokenizer.decode(output_list)
                return output_text
            else:
                output_list = list(cur_ids.squeeze().to('cpu').numpy())
                output_text = tokenizer.decode(output_list)
                return output_text


def load_models(model_name):
    """
	Summary:
		Loading the trained model
	"""
    print('Loading Trained GPT-2 Model')
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
    model = GPT2LMHeadModel.from_pretrained('gpt2-medium')
    model_path = model_name
    model.load_state_dict(torch.load(model_path))
    return tokenizer, model


def generate_text(sentences=5, model_name='text_augmenter_model.pt', label=''):
    """Generates augmented text using the fine-tuned gpt-2 model on the given dataset for a particular label to start with.

    Args:
        sentences (int): Number of sentences in outputs
        model_name (RhoModel): Name of the model file'
        label ([str]): Label for which to produce text

    Returns:
        [str]: Augmented text
    """

    SENTENCES = sentences
    MODEL_NAME = model_name
    LABEL = label

    TOKENIZER, MODEL = load_models(MODEL_NAME)

    output_text = generate(TOKENIZER, MODEL, SENTENCES, LABEL)
    return output_text
