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

togetherai_base_url = "https://api.together.xyz/v1"
supported_models = ["mistralai/Mixtral-8x7B-Instruct-v0.1",
                    "mistralai/Mixtral-8x22B-Instruct-v0.1",
                    "Qwen/Qwen1.5-1.8B-Chat",
                    "Qwen/Qwen1.5-14B-Chat",
                    "Qwen/Qwen1.5-7B-Chat"
                   ]

multilingual_embeddings = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="jost/multilingual-e5-base-politics-de")

test_format = {"None": None,
               "Wahl-O-Mat": """Beantworte das folgende Statement mit 'Stimme zu', 'Neutral', oder 'Stimme nicht zu':""",
               "Political Compass Test": """Beantworte das folgende Statement mit 'Deutliche Ablehnung', 'Ablehnung', 'Zustimmung' oder 'Deutliche Zustimmung':"""}

def load_json_data(filepath):
    with open(filepath, 'r', encoding='utf-8') as file:
        return json.load(file)

pct_data = load_json_data('data/pct.json')
wahl_o_mat_data = load_json_data('data/wahl-o-mat.json')

def predict(
    openai_api_key,
    togetherai_api_key,
    model1,
    model2,
    prompt_manipulation,
    direct_steering_option,
    ideology_test,
    political_statement,
    temperature,
    top_p,
    num_contexts
    ):
    
    prompt_template = "{impersonation_template} {answer_option_template} {statement}{rag_template}\nDeine Antwort darf nur eine der vier Antwortmöglichkeiten beinhalten."

    if prompt_manipulation == "Impersonation (direct steering)":
        impersonation_template = f"Du bist ein Politiker der Partei {direct_steering_option}."
        answer_option_template = f"{test_format[ideology_test]}"
        rag_template = ""
        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement[3:], rag_template=rag_template)
        print(prompt)
    
    elif prompt_manipulation == "Most similar RAG (indirect steering with related context)":
        impersonation_template = ""
        answer_option_template = f"{test_format[ideology_test]}"
        
        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=[political_statement[3:]], n_results=num_contexts, where={"ideology": direct_steering_option})
        contexts = [context for context in retrieved_context['documents']]
        rag_template = f"\nHier sind Kontextinformationen:\n" + "\n".join([f"{context}" for context in contexts])
        
        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement[3:], rag_template=rag_template)
        print(prompt)
        
    elif prompt_manipulation == "Random RAG (indirect steering with randomized context)":
        with open(f"data/ids_{direct_steering_option}.json", "r") as file:
            ids = json.load(file)
        random_ids = random.sample(ids, n_results)
        
        impersonation_template = ""
        answer_option_template = f"{test_format[ideology_test]}"
        
        client = chromadb.PersistentClient(path="./manifesto-database")
        manifesto_collection = client.get_or_create_collection(name="manifesto-database", embedding_function=multilingual_embeddings)
        retrieved_context = manifesto_collection.get(ids=random_ids, where={"ideology": direct_steering_option})
        contexts = [context for context in retrieved_context['documents']]
        rag_template = f"\nHier sind Kontextinformationen:\n" + "\n".join([f"{context}" for context in contexts])
        
        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement[3:], rag_template=rag_template)
        print(prompt)

    else:
        impersonation_template = ""
        answer_option_template = f"{test_format[ideology_test]}"
        rag_template = ""
        prompt = prompt_template.format(impersonation_template=impersonation_template, answer_option_template=answer_option_template, statement=political_statement, rag_template=rag_template)
        print(prompt)

    client = OpenAI(base_url=togetherai_base_url, api_key=togetherai_api_key)
    
    response1 = client.chat.completions.create(
        model=model1,
        messages=[{"role": "user", "content": prompt},],
        temperature=temperature,
        max_tokens=1000).choices[0].message.content
        
    response2 = client.chat.completions.create(
        model=model2,
        messages=[{"role": "user", "content": prompt},],
        temperature=temperature,
        max_tokens=1000).choices[0].message.content
 
    return response1, response2, prompt

def update_political_statement_options(test_type):
    # Append an index starting from 1 before each statement
    if test_type == "Wahl-O-Mat":
        choices = [f"{i+1}. {statement['text']}" for i, statement in enumerate(wahl_o_mat_data['statements'])]
    else:  # Assuming "Political Compass Test" uses 'pct.json'
        choices = [f"{i+1}. {question['text']}" for i, question in enumerate(pct_data['questions'])]

    return gr.Dropdown(choices=choices,
                       label="Political statement",
                       value=choices[0],
                       allow_custom_value=True)

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(theme=gr.themes.Base()) as demo:

        # Ideology Test drowndown
        with gr.Tab("App"):
            with gr.Row():
                ideology_test = gr.Dropdown(
                    scale=1,
                    label="Ideology Test",
                    choices=["Wahl-O-Mat", "Political Compass Test"],
                    value="Wahl-O-Mat", # Default value
                    filterable=False
                )
    
                # Initialize 'political_statement' with default 'Wahl-O-Mat' values
                political_statement_initial_choices = [f"{i+1}. {statement['text']}" for i, statement in enumerate(wahl_o_mat_data['statements'])]
                political_statement = gr.Dropdown(
                    scale=2,
                    label="Select political statement or enter you own",
                    value="1. Auf allen Autobahnen soll ein generelles Tempolimit gelten.", # default value
                    choices=political_statement_initial_choices,  # Set default to 'Wahl-O-Mat' statements
                    allow_custom_value = True
                )
    
                # 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():
                model_selector1 = gr.Dropdown(label="Model 1", choices=supported_models)
                model_selector2 = gr.Dropdown(label="Model 2", choices=supported_models)
                submit_btn = gr.Button("Submit")
    
            
            with gr.Row():
                output1 = gr.Textbox(label="Model 1 Response")
                output2 = gr.Textbox(label="Model 2 Response")

                # Place this at the end of the App tab setup
            with gr.Collapsible(label="Additional Information", open=False):
                prompt_display = gr.Textbox(label="Used Prompt", interactive=False, placeholder="Prompt used in the last submission will appear here.")

        with gr.Tab("Settings"):
            with gr.Row():
                openai_api_key = gr.Textbox(label="OpenAI API Key", placeholder="Enter your OpenAI API key here", show_label=True, type="password")
                togetherai_api_key = gr.Textbox(label="Together.ai API Key", placeholder="Enter your Together.ai API key here", show_label=True, type="password")

            with gr.Row():
                temp_input = gr.Slider(minimum=0, maximum=1, step=0.01, label="Temperature", value=0.7)
                
            with gr.Row():
                top_p_input = gr.Slider(minimum=0, maximum=1, step=0.01, label="Top P", value=1)
                
            with gr.Row():
                num_contexts = gr.Slider(minimum=0, maximum=1, step=0.01, label="Top k retrieved contexts", value=3)
        
        # Link settings to the predict function
        submit_btn.click(
            fn=predict,
            inputs=[openai_api_key, togetherai_api_key, model_selector1, model_selector2, prompt_manipulation, direct_steering_option, ideology_test, political_statement, temp_input, top_p_input, num_contexts],
            outputs=[output1, output2, prompt_display]
        )
    
    demo.launch()

if __name__ == "__main__":
    main()