# %%
import os
from openai import AzureOpenAI, OpenAI
import sys
sys.path.append('../')
from Util.constants import llama_2_7b_chat_api_key, llama_2_7b_chat_api_endpoint, phi_api_key, phi_api_endpoint, llama_2_7b_model_name, phi_model_name
from Util.smallModelOutput import outputModel
os.environ['AZURE_OPENAI_API_KEY'] = "YOUR API KEY"
os.environ['AZURE_OPENAI_ENDPOINT'] = "YOUR API ENDPOINT"

# %%
# import os
# from openai import AzureOpenAI

# client = AzureOpenAI(
#   azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"), 
#   api_key=os.getenv("AZURE_OPENAI_API_KEY"),  
#   api_version="2023-03-15-preview"
# )

# response = client.chat.completions.create(
#     model="GPT4", # model = "deployment_name".
#     messages=[
#         {"role": "system", "content": """Strictly follow the below conditions.
#          1. Output format: \nReasoning Chain: \nFinal Answer: 
#          2. Reasoning Chain should be seperated by new line only.
#          3. Final answer should not be in the reasoning chain.
#          4. Do not include any additional information in the final answer (only the final answer).
#          """},
#         {"role": "user", "content": """A robe takes 2 bolts of blue fiber and half that much white fiber.  How many bolts in total does it take? Let us think step by step."""},
#     ]
# )

# print(response.choices[0].message.content)

# %%
import json
# dir = 'MATH_Dataset/levelWise/COT/Level_1_data'
dir = '../Dataset/JEEBENCH/defaultCOT/mcq_data'
model = 'llama-2-7b-chat'

with open(f'{dir}.json', 'r') as f:
    data = json.load(f)

# %%
import os
from openai import AzureOpenAI

def createCOTask(question):
    client = AzureOpenAI(
        azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"), 
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),  
        api_version="2023-03-15-preview"
    )
    response = client.chat.completions.create(
    model="GPT4", # model = "deployment_name".
    messages=[
            {"role": "system", "content": """Strictly follow the below conditions.
            1. Output format: \nReasoning Chain: \nFinal Answer: 
            2. Reasoning Chain should be seperated by new line only.
            3. Reasoning chain cannot have the final answer. (Replace the final answer in the reasoning chain with its calculation or ####)
            4. Do not include any additional information in the final answer (only the answer).
            """},
            {"role": "user", "content": f"""{question} Let us think step by step."""},
        ]
    )
    return response.choices[0].message.content

def outputMistrael(messages):
    try:
        client = OpenAI(api_key="YOUR API KEY",
                        base_url="YOUR API ENDPOINT")
        response = client.chat.completions.create(
        model='mistralai/Mixtral-8x7B-Instruct-v0.1',
        messages=messages
        )
        return response.choices[0].message.content
    except Exception as e:
        return str(e)

def getMessage(question):
    return [
            {"role": "system", "content": """Strictly follow the below conditions.
            1. Output format: \nReasoning Chain: \nFinal Answer: 
            2. Reasoning Chain should be seperated by new line only.
            3. Reasoning chain cannot have the final answer. (Replace the final answer in the reasoning chain with its calculation or ####)
            4. Do not include any additional information in the final answer (only the answer).
            """},
            {"role": "user", "content": f"""{question} Let us think step by step."""},
        ]
def getMessage_2Shot(question):
    return [
            {"role": "system", "content": """Strictly follow the below conditions.
            1. Output format: \nReasoning Chain: \nFinal Answer: 
            2. Reasoning Chain should be seperated by new line only.
            3. Reasoning chain cannot have the final answer. (Replace the final answer in the reasoning chain with its calculation or ####)
            4. Do not include any additional information in the final answer (only the answer).
            """},
            {"role": "user", "content": f"""Raymond and Samantha are cousins. Raymond was born 6 years before Samantha. Raymond had a son at the age of 23. If Samantha is now 31, how many years ago was Raymond's son born? Let us think step by step."""},
            {"role": "assistant", "content": f"""Reasoning Chain:\nRaymond is 6 years older than Samantha.\nSo, Raymond's current age = Samantha's age + 6\n= 31 + 6\n= 37 years\n\nRaymond had a son 23 years ago.\nSo, Raymond's son's age = 23 years\n\nNow, to find out how many years ago Raymond's son was born, we subtract Raymond's son's age from Raymond's current age.\n= 37 - 23\n= ####\n\nFinal Answer: 14"""},
            {"role": "user", "content": f"""Kylar went to the store to buy glasses for his new apartment. One glass costs $5, but every second glass costs only 60% of the price. Kylar wants to buy 16 glasses. How much does he need to pay for them? Let us think step by step."""},
            {"role": "assistant", "content": f"""Reasoning Chain: \nKylar wants to buy 16 glasses. \nHalf of these glasses (8) will be sold full price, and the other half with a discount. \nSo, Kylar pays $5*8 = $40 for the first 8 glasses. \nThe other 8 glasses are sold with a 60% discount, so they cost $5*0.6 = $3. \nTherefore, Kylar pays $3*8 = $24 for the second 8 glasses. \n\nFinal Answer: 64"""},
            {"role": "user", "content": f"""{question} Let us think step by step."""}
        ]
import os
from openai import AzureOpenAI

def createCOTaskMathBench(endpoint, apikey,question, optionList):
    client = AzureOpenAI(
        azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT"), 
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),  
        api_version="2023-03-15-preview"
    )
    response = client.chat.completions.create(
    model="GPT4", # model = "deployment_name".
    messages=[
            {"role": "system", "content": """Strictly follow the below conditions.
            1. You will be provided with a question and four answer options. Select the correct answer option (A, B, C or D). 
            2. Output format: \nReasoning Chain: \nFinal Answer Option: 
            3. Reasoning Chain should be seperated by new line only.
            4. Reasoning chain cannot have the final answer option. (Replace the final answer in the reasoning chain with its calculation or ####)
            5. Only output A, B, C or D in final answer option.
            """},
            {"role": "user", "content": f"""Suppose  $15\\%$ of  $x$ equals  $20\\%$ of  $y.$ What percentage of  $x$ is  $y?$ Let us think step by step. \nOption A: 75 \nOption B: 5 \nOption C: 35 \nOption D: 300 \nWhich option is correct?"""},
            {"role": "assistant", "content": f"""Reasoning Chain: \nGiven that 15% of x equals 20% of y. We can write this as 0.15x = 0.2y.\nWe need to find y as a percentage of x. This means we need to rearrange the equation to make y the subject, or in other words, find y = ? in terms of x.\nDividing both sides of the equation by 0.2 we get y = 0.15x / 0.2 which simplifies to y = 0.75x. \nTranslated into a percentage y is 75% of x.\n\nFinal Answer Option: A"""},
            {"role": "user", "content": f"""{question} Let us think step by step. \nOption A: {optionList[0]} \nOption B: {optionList[1]} \nOption C: {optionList[2]} \nOption D: {optionList[3]} \nWhich option is correct?"""},
        ]
    )
    return response.choices[0].message.content

def jeeBenchMessage(question, options):
    return [
            {"role": "system", "content": """Strictly follow the below conditions.
            1. You will be provided with a question and four answer options. Select the correct answer option (A, B, C or D). 
            2. Output format: \nReasoning Chain: \nFinal Answer Option: 
            3. Reasoning Chain should be seperated by new line only.
            4. Reasoning chain cannot have the final answer option. (Replace the final answer in the reasoning chain with its calculation or ####)
            5. Only output A, B, C or D in final answer option.
            """},
            {"role": "user", "content": f"""Consider an experiment of tossing a coin repeatedly until the outcomes of two consecutive tosses are same. If the probability of a random toss resulting in head is $\\frac{1}{3}$, then the probability that the experiment stops with head is\n\n\n Let us think step by step. (A) $\\frac{1}{3}$\n\n\n\n(B) $\\frac{5}{21}$\n\n\n\n(C) $\\frac{4}{21}$\n\n\n\n(D) $\\frac{2}{7}$ \nWhich option is correct?"""},
            {"role": "assistant", "content": f"""Reasoning Chain: \nGiven that 15% of x equals 20% of y. We can write this as 0.15x = 0.2y.\nWe need to find y as a percentage of x. This means we need to rearrange the equation to make y the subject, or in other words, find y = ? in terms of x.\nDividing both sides of the equation by 0.2 we get y = 0.15x / 0.2 which simplifies to y = 0.75x. \nTranslated into a percentage y is 75% of x.\n\nFinal Answer Option: A"""},
            {"role": "user", "content": f"""{question} Let us think step by step. {options} \nWhich option is correct?"""},
        ]
def mathbenchMessage(question, optionList):
    return [
            {"role": "system", "content": """Strictly follow the below conditions.
            1. You will be provided with a question and four answer options. Select the correct answer option (A, B, C or D). 
            2. Output format: \nReasoning Chain: \nFinal Answer Option: 
            3. Reasoning Chain should be seperated by new line only.
            4. Reasoning chain cannot have the final answer option. (Replace the final answer in the reasoning chain with its calculation or ####)
            5. Only output A, B, C or D in final answer option.
            """},
            {"role": "user", "content": f"""Suppose  $15\\%$ of  $x$ equals  $20\\%$ of  $y.$ What percentage of  $x$ is  $y?$ Let us think step by step. \nOption A: 75 \nOption B: 5 \nOption C: 35 \nOption D: 300 \nWhich option is correct?"""},
            {"role": "assistant", "content": f"""Reasoning Chain: \nGiven that 15% of x equals 20% of y. We can write this as 0.15x = 0.2y.\nWe need to find y as a percentage of x. This means we need to rearrange the equation to make y the subject, or in other words, find y = ? in terms of x.\nDividing both sides of the equation by 0.2 we get y = 0.15x / 0.2 which simplifies to y = 0.75x. \nTranslated into a percentage y is 75% of x.\n\nFinal Answer Option: A"""},
            {"role": "user", "content": f"""{question} Let us think step by step. \nOption A: {optionList[0]} \nOption B: {optionList[1]} \nOption C: {optionList[2]} \nOption D: {optionList[3]} \nWhich option is correct?"""},
        ]
    
# %%
import copy
with open(f'{dir}_{model}.json', 'r') as f:
    COTData = json.load(f)
from tqdm import tqdm
for d in tqdm(data[len(COTData):400]):
    temp = d
    # temp['question'] = d['problem']
    # del temp['problem']
    # temp['question'] = d['question']
    # temp['reasoning_chain'] = d['answer']
    # temp['answer'] = d['answer'].split('####')[-1]
    if(model == 'llama-2-7b-chat'):
        temp['COT'] = outputModel(llama_2_7b_chat_api_key, llama_2_7b_chat_api_endpoint,llama_2_7b_model_name, jeeBenchMessage(d['question'], d['options']))
    elif(model == 'phi'):
        temp['COT'] = outputModel(phi_api_key, phi_api_endpoint,phi_model_name, jeeBenchMessage(d['question'], d['options']))
        # temp['COT'] = createCOTask(d['question'])
    if(model == 'mistral'):
        temp['COT'] = outputMistrael(jeeBenchMessage(d['question'], d['options']))
    COTData.append(temp)
    with open(f'{dir}_{model}.json', 'w') as f:
        json.dump(COTData, f, indent=4)
    
    

# %%



