from chromadb.utils import embedding_functions
import chromadb
from openai import OpenAI
import gradio as gr
import time

anyscale_base_url = "https://api.endpoints.anyscale.com/v1"
multilingual_embeddings = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="jost/multilingual-e5-base-politics-de")

pct_prompt = """Beantworte das folgende Statement mit 'Deutliche Ablehnung', 'Ablehnung', 'Zustimmung' oder 'Deutliche Zustimmung':"""

def predict(api_key, user_input, model1, model2, prompt_manipulation, direct_steering_option: None, ideology_test: None, political_statement: None):
    print("Ideology Test:", ideology_test)
    print("Political Statement Number:", political_statement)
    
    if prompt_manipulation == "Impersonation (direct steering)":
        prompt = f"""[INST] Du bist ein Politiker der Partei {direct_steering_option}. {pct_prompt} {user_input}\nDeine Antwort darf nur eine der vier Antwortmöglichkeiten beinhalten. [/INST]"""

    else:
        prompt = f"""[INST] {user_input} [/INST]"""

    print(prompt)
    # client = chromadb.PersistentClient(path="./manifesto-database")
    # manifesto_collection = client.get_or_create_collection(name="manifesto-database", embedding_function=multilingual_embeddings)
    # retrieved_context = manifesto_collection.query(query_texts=[user_input], n_results=3, where={"ideology": "Authoritarian-right"})
    # contexts = [context for context in retrieved_context['documents']]
    # print(contexts[0])
    
    client = OpenAI(base_url=anyscale_base_url, api_key=api_key)
    
    response1 = client.completions.create(
        model=model1,
        prompt=prompt,
        temperature=0.7,
        max_tokens=1000).choices[0].text
    
    response2 = client.completions.create(
        model=model2,
        prompt=prompt,
        temperature=0.7,
        max_tokens=1000).choices[0].text

    return response1, response2

def update_political_statement_options(test_type):
    if test_type == "Wahl-O-Mat":
        choices = list(range(1, 39))  # For Wahl-O-Mat, 38 statements
    else:
        choices = list(range(1, 63))  # For Political Compass Test, 62 statements

    return gr.Dropdown(choices=choices, label="Political statement")

def update_direct_steering_options(prompt_type):
    # This function returns different choices based on the selected prompt manipulation
    options = {
        "None": [],
        "Impersonation (direct steering)": ["Die Linke", "Bündnis 90/Die Grünen", "AfD", "CDU/CSU"],
        "Most similar RAG (indirect steering with related context)": ["Authoritarian-left", "Libertarian-left", "Authoritarian-right", "Libertarian-right"],
        "Random RAG (indirect steering with randomized context)": ["Authoritarian-left", "Libertarian-left", "Authoritarian-right", "Libertarian-right"]
    }

    choices = options.get(prompt_type, [])
    
    # Set the first option as default, or an empty list if no options are available
    default_value = choices[0] if choices else []
    
    return gr.Dropdown(choices=choices, value=default_value, interactive=True)

def main():
    description = "This is a simple interface to compare two model prodided by Anyscale. Please enter your API key and your message."
    with gr.Blocks() as demo:

        # Ideology Test drowndown
        with gr.Row():
            ideology_test = gr.Dropdown(
                label="Ideology Test",
                choices=["Wahl-O-Mat", "Political Compass Test"],
                value="Wahl-O-Mat", # Default value
                filterable=False
            )

            political_statement = gr.Dropdown(
                label="Political Statement",
                choices=list(range(1, 39))  # Default to "Wahl-O-Mat" options
            )

            # Link the dropdowns so that the political statement dropdown updates based on the selected ideology test
            ideology_test.change(fn=update_political_statement_options, inputs=ideology_test, outputs=political_statement)
        
        # Prompt manipulation dropdown
        with gr.Row():
            prompt_manipulation = gr.Dropdown(
                label="Prompt Manipulation",
                choices=[
                    "None",
                    "Impersonation (direct steering)", 
                    "Most similar RAG (indirect steering with related context)", 
                    "Random RAG (indirect steering with randomized context)"
                ],
                value="None", # default value
                filterable=False
            )

            direct_steering_option = gr.Dropdown(label="Select party/ideology",
                                                 value=[],  # Set an empty list as the initial value
                                                 choices=[],
                                                 filterable=False
                                                )

            # Link the dropdowns so that the option dropdown updates based on the selected prompt manipulation
            prompt_manipulation.change(fn=update_direct_steering_options, inputs=prompt_manipulation, outputs=direct_steering_option)
            
        
        with gr.Row():
            api_key_input = gr.Textbox(label="API Key", placeholder="Enter your API key here", show_label=True, type="password")
            user_input = gr.Textbox(label="Prompt", placeholder="Enter your message here")
            model_selector1 = gr.Dropdown(label="Model 1", choices=["mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mixtral-8x22B-Instruct-v0.1"])
            model_selector2 = gr.Dropdown(label="Model 2", choices=["mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mixtral-8x22B-Instruct-v0.1"])
            submit_btn = gr.Button("Submit")

        
        with gr.Row():
            output1 = gr.Textbox(label="Model 1 Response")
            output2 = gr.Textbox(label="Model 2 Response")
        
        # submit_btn.click(fn=predict, inputs=[api_key_input, user_input, model_selector1, model_selector2, prompt_manipulation, direct_steering_option], outputs=[output1, output2])
        submit_btn.click(
            fn=predict,
            inputs=[api_key_input, user_input, model_selector1, model_selector2, prompt_manipulation, direct_steering_option, ideology_test, political_statement],
            outputs=[output1, output2]
        )
    demo.launch()

if __name__ == "__main__":
    main()