import re
import random
import sys
import json
import time
import openai
import jsonlines
import random
import numpy as np
import backoff
from openai import OpenAI
client = OpenAI(api_key = "")


# Only output the above json object without any explanation.
@backoff.on_exception(backoff.expo, openai.RateLimitError)
def completions_with_backoff(**kwargs):
    return client.chat.completions.create(**kwargs)



def prompt_diverse4(sen_set):
    prompt = f"""
    Task: Classify the diversity of each sentence in the provided list by assigning it to one of five categories. These categories reflect how diverse each sentence is in comparison to others in the set:

    Categories:
    - "Very Similar": The sentence is very similar to or identical to other sentences in the list.
    - "Somewhat Similar": The sentence shows a moderate level of similarity to other sentences in the list.
    - "Neutral": The sentence neither shows significant similarity nor difference compared to others in the list.
    - "Somewhat Diverse": The sentence is somewhat different from other sentences in the list, either semantically or lexically.
    - "Highly Diverse": The sentence is highly different from other sentences in the list, both semantically and lexically.

    Sentence list:
    {sen_set}
    
    Instructions for Classification:
    - Create a list of classifications in JSON format.
    - Each entry should include the sentence and its corresponding diversity category.
    - Format for each entry: {{"sentence": "<sentence>", "category": "<category>"}}
    - No number before each entry.

    Example of Classification List Format:
    
    [
        {{"sentence": "Example sentence 1", "category": "Somewhat Diverse"}},
        {{"sentence": "Example sentence 2", "category": "Very Similar"}},
        ...
    ]

    Please note that no explanations for the scores are required. Just focus on providing the score for each sentence.
    """

    return prompt

def map_category_to_score(category):
    mapping = {
        "Very Similar": 1,
        "Somewhat Similar": 2,
        "Neutral": 3,
        "Somewhat Diverse": 4,
        "Highly Diverse": 5
    }
    return mapping.get(category, "Unknown Category")



# diversity
def get_score(sentences:list):
    sen_set = "\n".join(sentences)
    prompt = prompt_diverse4(sen_set)
    role = "You are a highly experienced NLP evaluator. Your expertise lies in assessing sentences generated by language models for their diversity. Your current task is to evaluate a set of sentences produced by a language model."
    chat_record = [{"role": "system", "content": role
                    },{"role": "user", "content": prompt}]
    output = gpt_output(chat_record)
    return output

def get_top_sentences(sentence_set, n, k):
    max_retries = 5  # Max number of retries if an error occurs
    for attempt in range(max_retries):
        try:
            scores_dict = {}
            random.shuffle(sentence_set)
            for _ in range(k):
                data = get_score(sentence_set)

                if '[' in data and ']' in data:
                    addressed = re.findall(r"\[.*\]", data, re.DOTALL)[0]

                else:
                    first_brace = data.find('{')
                    last_brace = data.rfind('}')
                    addressed = '[' + data[first_brace:last_brace+1] + ']'
                
                temp_data = eval(addressed)

                for item in temp_data:
                    sentence = item['sentence']
                    temp_cate = item['category']
                    score = map_category_to_score(temp_cate)
                    if sentence not in scores_dict:
                        scores_dict[sentence] = score
                    else:
                        scores_dict[sentence] += score

            sorted_data = sorted(scores_dict.items(), key=lambda x: x[1], reverse=True)
            print(sorted_data)
            sorted_sentence = sorted_data[:n]
            return [sentence for sentence, score in sorted_sentence]
            #return [sentence for sentence, score in sorted_data], [score for sentence, score in sorted_data]
        except (IndexError, json.JSONDecodeError):
            print(f"An error occurred on attempt {attempt + 1}. Retrying...")
            print(data)
            print("#############################################")
            if attempt == max_retries - 1:
                raise  
    raise


def get_avg_score(sentence_set, key):
    data = get_score(sentence_set)
    if isinstance(data, str):
        data = json.loads(data)
    total_score = sum(sentence['scores'][key] for sentence in data)
    # Calculate the average
    average_score = total_score / len(data)
    
    return average_score
    
def load_machine_data(filename):
    with open(filename, 'r', encoding='utf-8') as file:
        data = [json.loads(line) for line in file]
        for sample in data:
            sample['nomalized_scores'] = [sc/3 for sc in sample['scores']]
            scores = sample['nomalized_scores']
            sample['variance'] = np.var(scores)
            sample['mean'] = sum(scores) / len(scores)
    return data
    
def load_human_data(filename):
    with open(filename, 'r', encoding='utf-8') as file:
        data = [json.loads(line) for line in file]
    
    human_dict = {sample['src'].strip(): sample['human_rate'] for sample in data}

    return human_dict 


def extract_top_samples(data, key, reverse=False):
    sorted_samples = sorted(data, key=lambda x: x[key], reverse=reverse)
    
    return [(sample['src'],sample['sentences'],sample['variance'],sample['mean']) for sample in sorted_samples[:10]]
    
def gpt_output(chat_record:list) -> str:
    completion = completions_with_backoff(model="gpt-3.5-turbo", messages=chat_record, n=1, temperature=0.2)
    content = completion.choices[0].message.content
    return content


if __name__ == "__main__":
    # print("#############################################")
    # lst1 = [
    #     "A dog catches a frisbee thrown by its owner.",
    #     "A dog jumps in the air to catch a frisbee thrown by its owner.",
    #     "A dog catches a frisbee thrown by its owner."
    # ]
    # scores = get_score(lst1)
    # print(scores)
    
    lst1 = [
        "A dog leaps to catch a thrown frisbee.",
        "The dog catches the frisbee when the boy throws it.",
        "A man throws away his dog's favourite frisbee expecting him to catch it in the air."
    ]
    scores = get_score(lst1)
    print(scores)

    
     
