import sys
import json
import requests
import json
import re
from tqdm import tqdm, trange
import pandas as pd

def save_answer(answers, path):
    file_name = path + '.json'
    with open(file_name, "a", encoding='utf8') as fout:
        for i in range(len(answers)):
            fout.write(json.dumps({
                "answer": answers[i],
            }, ensure_ascii=False) + "\n")


def postprocess(answers, label):
    model_answer = []
    correct = 0.0
    for i in range(len(answers)):
        ans = answers[i].replace('(', '').replace(')', '').replace(' ', '').replace(':', '').replace('\n', '')
        poss_ans = ans[0]
        if poss_ans.lower() == label[i].lower():
            correct += 1
        model_answer.append(poss_ans)

    acc = correct / len(answers)
    return acc, model_answer

def mmlu_preprocess(data_path):
    prompt_list = []
    label_list = []
    pre_text = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. \n### Human: Got any creative ideas for a 10 year old’s birthday?\n### Assistant: Of course! Here are some creative ideas for a 10-year-old's birthday party:\n1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.\n2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.\n3. Outdoor Movie Night: Set up a backyard movie night with a projector and a large screen or white sheet. Create a cozy seating area with blankets and pillows, and serve popcorn and snacks while the kids enjoy a favorite movie under the stars.\n4. DIY Crafts Party: Arrange a craft party where kids can unleash their creativity. Provide a variety of craft supplies like beads, paints, and fabrics, and let them create their own unique masterpieces to take home as party favors.\n5. Sports Olympics: Host a mini Olympics event with various sports and games. Set up different stations for activities like sack races, relay races, basketball shooting, and obstacle courses. Give out medals or certificates to the participants.\n6. Cooking Party: Have a cooking-themed party where the kids can prepare their own mini pizzas, cupcakes, or cookies. Provide toppings, frosting, and decorating supplies, and let them get hands-on in the kitchen.\n7. Superhero Training Camp: Create a superhero-themed party where the kids can engage in fun training activities. Set up an obstacle course, have them design their own superhero capes or masks, and organize superhero-themed games and challenges.\n8. Outdoor Adventure: Plan an outdoor adventure party at a local park or nature reserve. Arrange activities like hiking, nature scavenger hunts, or a picnic with games. Encourage exploration and appreciation for the outdoors.\nRemember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!\n ### Human:"
    mid_text = "Here is a question and some options. Please read the passage and then choose the right anwser from the given options for the question.\n"
    # pos_text = "Now chose the right answer and return the ansewr in the following format: \n correct option:(the ID of your chooesed options). \n only return the answer, no other explanation is needed. \n ### Assistant:"
    pos_text = "The correct answer is: "

    data_df = pd.read_csv(data_path, names=['question', 'A', 'B', 'C', 'D', 'answer'])
    for idx, row in data_df.iterrows():
        question = row['question']
        options = f"{row['A']}\nB. {row['B']}\nC. {row['C']}\nD. {row['D']}"
        label = row['answer']
        prmt = pre_text + mid_text + '\n Question: \n' + question + '\n Options: \n' + options + '\n' + pos_text
        prompt_list.append(prmt)
        label_list.append(label)

    return prompt_list, label_list


def agi_preprocess(data_path):
    # path = "/yuxiaoxi/dataset/AGI/sat-en.jsonl"
    prompt_list = []
    label_list = []
    pre_text = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. \n### Human: Got any creative ideas for a 10 year old’s birthday?\n### Assistant: Of course! Here are some creative ideas for a 10-year-old's birthday party:\n1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.\n2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.\n3. Outdoor Movie Night: Set up a backyard movie night with a projector and a large screen or white sheet. Create a cozy seating area with blankets and pillows, and serve popcorn and snacks while the kids enjoy a favorite movie under the stars.\n4. DIY Crafts Party: Arrange a craft party where kids can unleash their creativity. Provide a variety of craft supplies like beads, paints, and fabrics, and let them create their own unique masterpieces to take home as party favors.\n5. Sports Olympics: Host a mini Olympics event with various sports and games. Set up different stations for activities like sack races, relay races, basketball shooting, and obstacle courses. Give out medals or certificates to the participants.\n6. Cooking Party: Have a cooking-themed party where the kids can prepare their own mini pizzas, cupcakes, or cookies. Provide toppings, frosting, and decorating supplies, and let them get hands-on in the kitchen.\n7. Superhero Training Camp: Create a superhero-themed party where the kids can engage in fun training activities. Set up an obstacle course, have them design their own superhero capes or masks, and organize superhero-themed games and challenges.\n8. Outdoor Adventure: Plan an outdoor adventure party at a local park or nature reserve. Arrange activities like hiking, nature scavenger hunts, or a picnic with games. Encourage exploration and appreciation for the outdoors.\nRemember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!\n ### Human:"
    mid_text = "Here are a passage, a question and some options. Please read the passage and then choose the right anwser from the given options for the question.\n"
    # pos_text = "Now chose the right answer and return the ansewr in the following format: \n correct option:(the ID of your chooesed options). \n only return the answer, no other explanation is needed. \n ### Assistant:"
    pos_text = "The correct answer is: "
    with open(data_path, 'r') as f:
        for line in f.readlines():
            passage = json.loads(line)['passage']
            question = json.loads(line)['question']
            options = ' '.join(json.loads(line)['options'])
            label = json.loads(line)['label']
            if not passage:
                passage = ''
            elif passage == 'null':
                passage = ''
            prmt = pre_text + mid_text + '\n Passage: \n' + passage + '\n Question: \n' + question + '\n Options: \n' + options + '\n' + pos_text
            prompt_list.append(prmt)
            label_list.append(label)
    return prompt_list, label_list


def ceval_preprocess(data_path):
    prompt_list = []
    label_list = []
    pre_text = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. \n### Human: Got any creative ideas for a 10 year old’s birthday?\n### Assistant: Of course! Here are some creative ideas for a 10-year-old's birthday party:\n1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.\n2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.\n3. Outdoor Movie Night: Set up a backyard movie night with a projector and a large screen or white sheet. Create a cozy seating area with blankets and pillows, and serve popcorn and snacks while the kids enjoy a favorite movie under the stars.\n4. DIY Crafts Party: Arrange a craft party where kids can unleash their creativity. Provide a variety of craft supplies like beads, paints, and fabrics, and let them create their own unique masterpieces to take home as party favors.\n5. Sports Olympics: Host a mini Olympics event with various sports and games. Set up different stations for activities like sack races, relay races, basketball shooting, and obstacle courses. Give out medals or certificates to the participants.\n6. Cooking Party: Have a cooking-themed party where the kids can prepare their own mini pizzas, cupcakes, or cookies. Provide toppings, frosting, and decorating supplies, and let them get hands-on in the kitchen.\n7. Superhero Training Camp: Create a superhero-themed party where the kids can engage in fun training activities. Set up an obstacle course, have them design their own superhero capes or masks, and organize superhero-themed games and challenges.\n8. Outdoor Adventure: Plan an outdoor adventure party at a local park or nature reserve. Arrange activities like hiking, nature scavenger hunts, or a picnic with games. Encourage exploration and appreciation for the outdoors.\nRemember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!\n ### Human:"
    mid_text = "Here is a question and some options. Please read the passage and then choose the right anwser from the given options for the question.\n"
    # pos_text = "Now chose the right answer and return the ansewr in the following format: \n correct option:(the ID of your chooesed options). \n only return the answer, no other explanation is needed. \n ### Assistant:"
    pos_text = "The correct answer is: "

    with open(data_path, 'r') as f:
        for line in f.readlines():

            passage = json.loads(line)['inputs_pretokenized']
            question = passage.split('\n')[0]
            options = passage.split('\n')[1:]
            options = '\n'.join(options)
            label = json.loads(line)['targets_pretokenized'][0]

            prmt = pre_text + mid_text + '\n Question: \n' + question + '\n Options: \n' + options + '\n' + pos_text
            prompt_list.append(prmt)
            label_list.append(label)

    return prompt_list, label_list
