from __future__ import annotations
import sys, os, json, time
os.environ['CC'] = '/opt/rh/gcc-toolset-9/root/usr/bin/gcc'
os.environ['CXX'] = '/opt/rh/gcc-toolset-9/root/usr/bin/g++'
os.environ["PATH"] += os.pathsep + os.path.expanduser("~/.local/bin")
os.environ["PATH"] += os.pathsep + "/rhome/jyang311/shared/jyangfile/jyangfile/trainingmaterial/exllamav2"
os.environ["MAX_JOBS"] = '16'

from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer
from exllamav2.generator import ExLlamaV2DynamicGenerator, ExLlamaV2DynamicJob, ExLlamaV2Sampler
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
import os
import math
from tqdm import tqdm
import json

def calculate_perplexity(model, tokenizer, input_ids) -> torch.Tensor:
    # Tokenize the input text
    target_ids = input_ids[:, 1:]

    # Compute loss
    with torch.no_grad():
        logits = model.forward(input_ids).float()
        logprobs = F.log_softmax(logits, dim=-1).to('cuda')
        target_ids = target_ids.to('cuda')
        target_logprobs = logprobs[:, :-1].gather(dim=-1, index=target_ids.unsqueeze(-1)).squeeze(-1)
        nll = -target_logprobs

    # Perplexity is exp(loss)
    perplexity = nll.exp()
    return nll

def calculate_stdev(numbers):
    if len(numbers) < 2:
        raise ValueError("At least two numbers are required to calculate the standard deviation.")

    mean = sum(numbers) / len(numbers)
    variance = sum((x - mean) ** 2 for x in numbers) / (len(numbers) - 1)
    stdev = math.sqrt(variance)
    return stdev

import json
with open('clustered_embeded_full.json') as tmpinputCollector:
    col = json.load(tmpinputCollector)

modelDir = "../Mistral-Large-Instruct-2407-123B-exl2/"
max_response_len = 1024
cache_size = 100*1024  # Adjust as needed, 100k seems to be a safe size for L3-8B on a single 24 GB GPU
max_rows = 10000


config = ExLlamaV2Config(modelDir)
tokenizer = ExLlamaV2Tokenizer(config)
model = ExLlamaV2(config)
cache = ExLlamaV2Cache(model, max_seq_len = cache_size, lazy = True)
model.load_autosplit(cache)

def getPpl(input_text):
    input_ids = tokenizer.encode(input_text)

    # Calculate perplexity
    perplexity = calculate_perplexity(model, tokenizer, input_ids)

    vocab = tokenizer.get_id_to_piece_list()
    
    tokens = []
    ppls = []
    for idx in range(input_ids.shape[-1]):
        token = input_ids[0, idx].item()
        text = vocab[token]
        ppl = float("inf") if idx == 0 else perplexity[0, idx - 1].item()
        tokens.append(text)
        ppls.append(ppl)

    return tokens, ppls

with open('isolatedMiddle.json') as tmpinputCollector:
    res = json.load(tmpinputCollector)

i=0
for sample in tqdm(res):
    premise = sample['premise']
    conclusion = sample['conclusion']
    for middle in sample['middle4conclusionppl']:
        if sample['middle4conclusionppl'][middle] !='':
            continue
        #print(middle)
        finalInput = premise.strip()+' '+middle.strip()+' '+conclusion
        #print(finalInput)
        tokens, ppls = getPpl(finalInput)
        lenofConclusion = len(tokenizer.encode(conclusion)[0])
        #print(tokenizer.encode(conclusion)[0])
        concPpl = ppls[-lenofConclusion:]
        # print(sum(concPpl) / len(concPpl))
        sample['middle4conclusionppl'][middle] = concPpl
        # print(len(concPpl))
        i+=1
        if i%3000 != 0:
            continue
        with open("isolatedMiddle.json", "w") as outfile:
            json.dump(res, outfile)