import os
def getDataPath(dataset, type, COT_type):
    count = 100
    if(dataset == 'gsm8k'):
        if(COT_type == 'smallerModel'):
            if(type == 'llama-2-7b-chat_incorrect_reasoning'):
                return '../Dataset/gsm8k/dataFiles/smallModel/COT_Incorrect_llama-2-7b-chat.json', count
            elif(type == 'mistral_incorrect_reasoning'):
                return '../Dataset/gsm8k/dataFiles/smallModel/COT_Incorrect_mistral.json', count
            elif(type == 'phi_incorrect_reasoning'):
                return '../Dataset/gsm8k/dataFiles/smallModel/COT_Incorrect_phi.json', count
            else:
                return None, None
        elif(COT_type == 'default'):
            return '../Dataset/gsm8k/dataFiles/defaultCOT/nonCOT_InCorrectReasoning.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/gsm8k/dataFiles/COT/COT_InCorrectReasoning.json',count
    count = 50
    if(dataset == 'MATH-Level-1'):
        if(COT_type == 'smallerModel'):
            if(type == 'llama-2-7b-chat_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_1_data_raw_COTData_Incorrect_llama-2-7b-chat.json', count
            elif(type == 'mistral_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_1_data_raw_COTData_Incorrect_mistral.json', count
            elif(type == 'phi_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_1_data_raw_COTData_Incorrect_phi.json', count
            else:
                return None, None
        elif(COT_type == 'default'):
            return '../Dataset/MATH/dataFiles/defaultCOT/Level_1_nonCOT.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/MATH/dataFiles/COT/Level_1_data_COT_inCorrectReasoning.json', count
    if(dataset == 'MATH-Level-2'):
        if(COT_type == 'smallerModel'):
            if(type == 'llama-2-7b-chat_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_2_data_raw_COTData_Incorrect_llama-2-7b-chat.json', count
            elif(type == 'mistral_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_2_data_raw_COTData_Incorrect_mistral.json', count
            elif(type == 'phi_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_2_data_raw_COTData_Incorrect_phi.json', count
            else:
                return None, None
        elif(COT_type == 'default'):
            return '../Dataset/MATH/dataFiles/defaultCOT/Level_2_nonCOT.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/MATH/dataFiles/COT/Level_2_data_COT_inCorrectReasoning.json', count
    if(dataset == 'MATH-Level-3'):
        if(COT_type == 'smallerModel'):
            if(type == 'llama-2-7b-chat_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_3_data_raw_COTData_Incorrect_llama-2-7b-chat.json', count
            elif(type == 'mistral_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_3_data_raw_COTData_Incorrect_mistral.json', count
            elif(type == 'phi_incorrect_reasoning'):
                return '../Dataset/MATH/dataFiles/smallerModel/IncorrectCOT/Level_3_data_raw_COTData_Incorrect_phi.json', count
            else:
                return None, None
        elif(COT_type == 'default'):
            return '../Dataset/MATH/dataFiles/defaultCOT/Level_3_nonCOT.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/MATH/dataFiles/COT/Level_3_data_COT_inCorrectReasoning.json', count
    count = 100
    if(dataset == 'MATHBENCH'):
        if(COT_type == 'smallerModel'):
            if(type == 'llama-2-7b-chat_incorrect_reasoning'):
                return '../Dataset/MATHBENCH/dataFiles/smallerModel/math_bench_COTData_Incorrect_llama-2-7b-chat.json', count
            elif(type == 'mistral_incorrect_reasoning'):
                return '../Dataset/MATHBENCH/dataFiles/smallerModel/math_bench_COTData_Incorrect_mistral.json', count
            elif(type == 'phi_incorrect_reasoning'):
                return '../Dataset/MATHBENCH/dataFiles/smallerModel/math_bench_COTData_Incorrect_phi.json', count
            else:
                return None, None
        elif(COT_type == 'COT'):
            return '../Dataset/MATHBENCH/dataFiles/COT/COT_inCorrectReasoning.json', count
    count = 100
    if(dataset == 'JEEBENCH'):
        if(COT_type == 'smallerModel'):
            if(type == 'llama-2-7b-chat_incorrect_reasoning'):
                return '../Dataset/JEEBENCH/smallerModel/mcq_data_Incorrect_llama-2-7b-chat.json', count
            elif(type == 'mistral_incorrect_reasoning'):
                return '../Dataset/JEEBENCH/smallerModel/mcq_data_Incorrect_mistral.json', count
            elif(type == 'phi_incorrect_reasoning'):
                return '../Dataset/JEEBENCH/smallerModel/mcq_data_Incorrect_phi.json', count
            else:
                return None, None
        elif(COT_type == 'default'):
            return '../Dataset/JEEBENCH/defaultCOT/COT_IncorrectReasoning.json', count
        
def getDataPathMemorization(dataset, type, COT_type):
    count = 100
    if(dataset == 'gsm8k'):
        if(COT_type == 'smallerModel'):
            return '../Dataset/gsm8k/dataFiles/smallModel/COT_Correct_Combined.json', count
        elif(COT_type == 'default'):
            return '../Dataset/gsm8k/dataFiles/defaultCOT/nonCOT_InCorrectReasoning.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/gsm8k/dataFiles/COT/COT_InCorrectReasoning.json',count
    count = 50
    if(dataset == 'MATH-Level-1'):
        if(COT_type == 'smallerModel'):
            return '../Dataset/MATH/dataFiles/smallerModel/CorrectCOT/Combined/Level_1_data_raw_COTData_Correct_Combined.json', count
        elif(COT_type == 'default'):
            return '../Dataset/MATH/dataFiles/defaultCOT/Level_1_nonCOT.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/MATH/dataFiles/COT/Level_1_data_COT_inCorrectReasoning.json', count
    if(dataset == 'MATH-Level-2'):
        if(COT_type == 'smallerModel'):
            return '../Dataset/MATH/dataFiles/smallerModel/CorrectCOT/Combined/Level_2_data_raw_COTData_Correct_Combined.json', count
        elif(COT_type == 'default'):
            return '../Dataset/MATH/dataFiles/defaultCOT/Level_2_nonCOT.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/MATH/dataFiles/COT/Level_2_data_COT_inCorrectReasoning.json', count
    if(dataset == 'MATH-Level-3'):
        if(COT_type == 'smallerModel'):
            return '../Dataset/MATH/dataFiles/smallerModel/CorrectCOT/Combined/Level_3_data_raw_COTData_Correct_Combined.json', count
        elif(COT_type == 'default'):
            return '../Dataset/MATH/dataFiles/defaultCOT/Level_3_nonCOT.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/MATH/dataFiles/COT/Level_3_data_COT_inCorrectReasoning.json', count
    count = 100
    if(dataset == 'MATHBENCH'):
        if(COT_type == 'smallerModel'):
            return '../Dataset/MATHBENCH/dataFiles/smallerModel/math_bench_COTData_correct_combined.json', count
        elif(COT_type == 'COT'):
            return '../Dataset/MATHBENCH/dataFiles/COT/COT_inCorrectReasoning.json', count
    count = 100
    if(dataset == 'JEEBENCH'):
        if(COT_type == 'smallerModel'):
            return '../Dataset/JEEBENCH/smallerModel/mcq_data_correct_combined.json', count
        elif(COT_type == 'default'):
            return '../Dataset/JEEBENCH/defaultCOT/COT_IncorrectReasoning.json', count
            
        

def getSavePath(dataset, type, model, COT_type, parsed=False, onlyFinalAnswer=False):
    if(dataset == 'gsm8k'):
        if(parsed == False):
            os.makedirs(f'../Dataset/gsm8k/results/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/gsm8k/results/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/gsm8k/results/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
        if(parsed == True):
            os.makedirs(f'../Dataset/gsm8k/results/parsed/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/gsm8k/results/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/gsm8k/results/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
    if(dataset == 'MATH-Level-1'):
        if(parsed == False):
            os.makedirs(f'../Dataset/MATH/results/level-1/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer):
                return f'../Dataset/MATH/results/level-1/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/MATH/results/level-1/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
        if(parsed == True):
            os.makedirs(f'../Dataset/MATH/results/level-1/parsed/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/MATH/results/level-1/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/MATH/results/level-1/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
    if(dataset == 'MATH-Level-2'):
        if(parsed == False):
            os.makedirs(f'../Dataset/MATH/results/level-2/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/MATH/results/level-2/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:  
                return f'../Dataset/MATH/results/level-2/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
        if(parsed == True):
            os.makedirs(f'../Dataset/MATH/results/level-2/parsed/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/MATH/results/level-2/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/MATH/results/level-2/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
    if(dataset == 'MATH-Level-3'):
        if(parsed == False):
            os.makedirs(f'../Dataset/MATH/results/level-3/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/MATH/results/level-3/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/MATH/results/level-3/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
        if(parsed == True):
            os.makedirs(f'../Dataset/MATH/results/level-3/parsed/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/MATH/results/level-3/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/MATH/results/level-3/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
    if(dataset == 'MATHBENCH'):
        if(parsed == False):
            os.makedirs(f'../Dataset/MATHBENCH/results/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/MATHBENCH/results/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/MATHBENCH/results/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
        if(parsed == True):
            os.makedirs(f'../Dataset/MATHBENCH/results/parsed/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/MATHBENCH/results/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/MATHBENCH/results/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
    if(dataset == 'JEEBENCH'):
        if(parsed == False):
            os.makedirs(f'../Dataset/JEEBENCH/results/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/JEEBENCH/results/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/JEEBENCH/results/{model}/{COT_type}/{type}_{COT_type}_{model}.json'
        if(parsed == True):
            os.makedirs(f'../Dataset/JEEBENCH/results/parsed/{model}/{COT_type}', exist_ok=True)
            if(onlyFinalAnswer == True):
                return f'../Dataset/JEEBENCH/results/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}_onlyFinalAnswer.json'
            else:
                return f'../Dataset/JEEBENCH/results/parsed/{model}/{COT_type}/{type}_{COT_type}_{model}.json'

        
        

    