import openai
import time

def set_api_key(tag):
    # GPT4-turbo
    if tag == 0:
        raise NotImplementedError
    elif tag == 1:
        client = openai.AzureOpenAI(
        api_key = "f6052a23f2154134bd7ed1e37c896814",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest1.openai.azure.com/"
        )
    elif tag == 2:
        client = openai.AzureOpenAI(
        api_key = "c526ab27d0e94df6906cc2874c4a1ab9",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest3.openai.azure.com/"
        )
    elif tag == 3:
        client = openai.AzureOpenAI(
        api_key = "ee2cbb8fda114a509317d2a2c4a703bc",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest7.openai.azure.com/"
        )
    elif tag == 4:
        client = openai.AzureOpenAI(
        api_key = "e0f5b4e9daae45419068360d9ced2bcf",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest8.openai.azure.com/"
        )
    elif tag == 5:
        client = openai.AzureOpenAI(
        api_key = "aa1dc0d5e551408ca9b5c36b1e830431",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest9.openai.azure.com/"
        )
    elif tag == 6:
        client = openai.AzureOpenAI(
        api_key = "36e3e1c4d4ca403d8842efa13cfdd918",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest10.openai.azure.com/"
        )
    elif tag == 7:
        client = openai.AzureOpenAI(
        api_key = "ea52f340bc1d433083d473263aef3732",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest5.openai.azure.com/"
        )
    elif tag == 8:
        client = openai.AzureOpenAI(
        api_key = "7c08884cc0a14252ac25ffc68591db6f",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest6.openai.azure.com/"
        )
    elif tag == 9:
        client = openai.AzureOpenAI(
        api_key = "e73a07643f0a45b28f43c219b156925e",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest2.openai.azure.com/"
        )
    # GPT4
    elif tag == 10:
        client = openai.AzureOpenAI(
        api_key = "704ea8e23c5749cf8fde870db7048468",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://ai4mtest4.openai.azure.com/"
        )
    elif tag == 11:
        client = openai.AzureOpenAI(
        api_key = "a3ab304fdc3849caa63f74eddd6b88f7",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://promptengineer.openai.azure.com/"
        )
    elif tag == 12:
        client = openai.AzureOpenAI(
        api_key = "a3bb33e9d0fe420e8ead829884bd44a6",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://testcloseai.openai.azure.com/"
        )
    elif tag == 13:
        client = openai.AzureOpenAI(
        api_key = "8f1874d5cd1b4d0c9e32f7b177690fee",  
        api_version = "2023-07-01-preview",
        azure_endpoint = "https://agatha-gpt-playground.openai.azure.com/"
        )
    else:
        raise NotImplementedError
    return client

def get_ids(min_index=1, max_index=30):
    for i in range(min_index, max_index+1):
        try:
            _ = set_api_key(i)
        except NotImplementedError:
            break
    return list(range(min_index, i))


# to send GPT4 a request and get a response back
# autoformalize the problem
def gpt4_response(client, prompt, examples, max_tokens=800, top_p=0.95):  
    content = ""
    while True:
        try:
            response = client.chat.completions.create(
                model="gpt-4",
                messages = examples+[{"role":"user","content":f"{prompt}"}],  # Load the examples at prob_message
                temperature=0.7,
                # max token is configured to be 6000 for a complete response
                max_tokens=max_tokens,
                top_p=top_p,
                frequency_penalty=0,
                presence_penalty=0,
                stop=None)
            content = response.choices[0].message.content
        except openai.BadRequestError as e:
            content = ""
        except openai.RateLimitError as e:
            time.sleep(5)
        break
    if content is None:
        content = ""
    return content

