import torch
import nltk
from nltk.translate.bleu_score import SmoothingFunction
from tqdm import tqdm

def calculate_perplexity(model, tokens, prompt_len, bsz=1, marker=False):
    """
    Calculate perplexity of given tokens using provided model, ignoring padding tokens. 
    Args:
        model: Llama model
        tokens (List[List[int]] or torch.Tensor): Input tokens (n_prompt * n_draft, seqlen)
        prompt_len (int): Prefix length
        bsz (int): Batch size
        marker (bool): Whether to show progress bar
    Returns:
        Perplexity across all generations (n_prompt * n_drafts)
    """
    it = range(0, len(tokens), bsz)
    if marker:
        it = tqdm(it)
    start = 0
    ppl = torch.zeros(len(tokens))
    for start in it:
        end = start + bsz
        data = tokens[start : end]
        if not isinstance(data, list):
            data = data.tolist()
        # Remove any padding tokens (-1) in generations
        for d_idx in range(len(data)):
            cur = data[d_idx]
            if -1 in cur:
                data[d_idx] = cur[:cur.index(-1)]
        # Calculate cross entropy loss on tokens
        ce_loss = model.generate(data, max_gen_len=0, temperature=-1, top_p=-1, grade=True)
        # Cut off everything past `prompt_len`
        ce_loss = ce_loss[:, prompt_len-1:]  # Subtract 1 because the first token (start token) is removed
        # Calculate perplexity 
        lengths = (ce_loss != 0).sum(dim=-1)
        mean = ce_loss.sum(dim=-1) / lengths
        ppl[start : end] = torch.exp(-1 * mean)
    return ppl
    
def calculate_diversity(generations, k=4):
    """
    Calculate diversity of generations using SELF-BLEU.
    Args:
        generations (List[List[List[int]]]): Tokenized input
        k (int, Optional): Number of n-grams to use for bleu
    Returns:
        Average diversity across all generations (float)
    """
    nltk.download('punkt')  # Can be deleted once downloaded
    smooth = SmoothingFunction()
    bleus = []
    
    for drafts in generations:
        tokenized_drafts = []
        # Stringify tokens
        for d in drafts:
            if -1 in d:
                d = d[:d.index(-1)]
            tokenized_drafts.append([str(n) for n in d])
        # Calculate SELF-BLEU
        minlength = min([len(g) for g in tokenized_drafts])
        minlength = min(minlength, k)
        weights = tuple((1. / minlength for _ in range(minlength)))
        for i in range(len(drafts)):
            # Create source and reference (all other drafts)
            src = tokenized_drafts[i]
            ref = tokenized_drafts[:i] + tokenized_drafts[i+1:]
            tmp = nltk.translate.bleu_score.sentence_bleu(references=ref, 
                                                          hypothesis=src, 
                                                          weights=weights,
                                                          smoothing_function=smooth.method1)
            bleus.append(tmp)
    bleus = torch.Tensor(bleus)
    return torch.mean(bleus)


def calculate_ngram_repetition(sequences):
    """
    Calculate uniqueness scores of `sequences`.
    Args:
        sequences (List[List[int]]): Generated sequences
    Returns:
        (unigram_uniqueness, bigram_uniqueness, trigram_uniqueness)
    """
    u_total = 0
    b_total = 0
    t_total = 0
    # Iterate through all sequences indiscriminately
    for gen in sequences:
        if -1 in gen:
            gen = gen[:gen.index(-1)]
        unigrams, bigrams, trigrams = [], [], []
        o = [str(i) for i in gen]
        # Create lists of n-grams for the generation
        for i in range(len(o)):
            unigrams.append(o[i])
        for i in range(len(o) - 1):
            bigrams.append(o[i] + '_' + o[i + 1])
        for i in range(len(o) - 2):
            trigrams.append(o[i] + '_' + o[i + 1] + '_' + o[i + 2])
        # Calculate uniqueness of the generation
        u, b, t = len(set(unigrams)) / len(unigrams), len(set(bigrams)) / len(bigrams), len(set(trigrams)) / len(trigrams)
        u_total += u
        b_total += b
        t_total += t
    return u_total / len(sequences), b_total / len(sequences), t_total / len(sequences)