import chromadb
import os
import gradio as gr
import json
from huggingface_hub import InferenceClient

path = '/Users/thiloid/Desktop/Avah/chroma'
if not os.path.exists(path): 
    path = "/home/user/app/chroma"

print(path)

client = chromadb.PersistentClient(path=path)
print(client.heartbeat()) 
print(client.get_version())  
print(client.list_collections()) 

from chromadb.utils import embedding_functions
default_ef = embedding_functions.DefaultEmbeddingFunction()
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="T-Systems-onsite/cross-en-de-roberta-sentence-transformer")

collection = client.get_collection(name="chromatsc", embedding_function=sentence_transformer_ef)

client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")



def format_prompt(message, history):
    prompt = ""
    if history:
        user_prompt, bot_response = history[-1]
        prompt += f"[INST] {user_prompt} [/INST] {bot_response}</s> "
    prompt += f"[INST] {message} [/INST]"    
    return prompt


def response(prompt, history, selected_option, temperature=0.9, max_new_tokens=700, top_p=0.95, repetition_penalty=1.0):
    temperature = float(temperature)
    if temperature < 1e-2: temperature = 1e-2
    top_p = float(top_p)
    generate_kwargs = dict(
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
        do_sample=True,
        seed=42,
    )
    if selected_option == "Keine These ausgewählt":
        selected_option = "Die Frage bezieht sich auf keine These speziell und ist somit allgemeiner zu behandeln."
        search_prompt = format_prompt(prompt, history)
    else:
        search_prompt = selected_option
        selected_option = "Die Nutzer Frage bezieht sich auf folgende Wahl-O-Mat These: " + selected_option
    addon = ""
    results=collection.query(
          query_texts=[search_prompt],
          n_results=8
    )
    
    dists = ["<br><small>(relevance: " + str(round((1 - d) * 100) / 100) + ";" for d in results['distances'][0]]
    results = results['documents'][0]
    combination = zip(results, dists)
    combination = [' '.join(triplets) for triplets in combination]
    print(results)
    if len(results) > 1:
        addon = "Bitte berücksichtige bei deiner Antwort ausschießlich folgende Auszüge aus unserer Wahlhilfe Datenbank, sofern sie für die Antwort erforderlich sind. Beantworte die Frage knapp und präzise. Ignoriere unpassende Datenbank-Auszüge OHNE sie zu kommentieren, zu erwähnen oder aufzulisten:\n" + "\n".join(results)
    system = "Du heißt Avah und bist ein deutschsprachiges KI-basiertes Wahlhilfe Beratungssystem zum Wahl-O-Maten für die Bundestagswahl 2021, an welchem insgesamt 38 Parteien teilgenommen haben.  Du beinhaltest aber nur Informationen zu den 6 Hauptparteien. Du nimmst selber keine eigene politische Postion ein!" + addon + "\n\nUser-Anliegen:"
    
    formatted_prompt = format_prompt(system + "\n" + selected_option + "\n" + prompt, history)
    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""
    for response in stream:
        output += response.token.text
    return output


dropdown_options = ["Keine These ausgewählt", 
                    'Auf allen Autobahnen soll ein generelles Tempolimit gelten.',
 'Deutschland soll seine Verteidigungsausgaben erhöhen.',
 'Bei Bundestagswahlen sollen auch Jugendliche ab 16 Jahren wählen dürfen.',
 'Die Förderung von Windenergie soll beendet werden.',
 'Die Möglichkeiten der Vermieterinnen und Vermieter, Wohnungsmieten zu erhöhen, sollen gesetzlich stärker begrenzt werden.',
 'Impfstoffe gegen Covid-19 sollen weiterhin durch Patente geschützt sein.',
 'Der für das Jahr 2038 geplante Ausstieg aus der Kohleverstromung soll vorgezogen werden.',
 'Alle Erwerbstätigen sollen in der gesetzlichen Rentenversicherung versichert sein müssen.',
 'Das Recht anerkannter Flüchtlinge auf Familiennachzug soll abgeschafft werden.',
 'Auf den Umsatz, der in Deutschland mit digitalen Dienstleistungen erzielt wird, soll eine nationale Steuer erhoben werden.',
 'Die traditionelle Familie aus Vater, Mutter und Kindern soll stärker als andere Lebensgemeinschaften gefördert werden.',
 'Spenden von Unternehmen an Parteien sollen weiterhin erlaubt sein.',
 'Studentinnen und Studenten sollen BAföG unabhängig vom Einkommen ihrer Eltern erhalten.',
 'In Deutschland soll es generell möglich sein, neben der deutschen eine zweite Staatsbürgerschaft zu haben.',
 'Bundesbehörden sollen in ihren Veröffentlichungen unterschiedliche Geschlechtsidentitäten sprachlich berücksichtigen.',
 'Die Ostsee-Pipeline "Nord Stream 2", die Gas von Russland nach Deutschland transportiert, soll wie geplant in Betrieb gehen dürfen.',
 'Der Solidaritätszuschlag soll vollständig abgeschafft werden.',
 'Das Tragen eines Kopftuchs soll Beamtinnen im Dienst generell erlaubt sein.',
 'Die Zulassung von neuen Autos mit Verbrennungsmotor soll auch langfristig möglich sein.',
 'Der Bund soll mehr Zuständigkeiten in der Schulpolitik erhalten.',
 'Der Bund soll Projekte zur Bekämpfung des Antisemitismus stärker finanziell unterstützen.',
 'Chinesische Firmen sollen keine Aufträge für den Ausbau der Kommunikationsinfrastruktur in Deutschland erhalten dürfen.',
 'Der Staat soll weiterhin für Religionsgemeinschaften die Kirchensteuer einziehen.',
 'Der kontrollierte Verkauf von Cannabis soll generell erlaubt sein.',
 'Deutschland soll aus der Europäischen Union austreten.',
 'Die Landeslisten der Parteien für die Wahlen zum Deutschen Bundestag sollen abwechselnd mit Frauen und Männern besetzt werden müssen.',
 'Stationäre Behandlungen im Krankenhaus sollen weiterhin über eine Fallpauschale abgerechnet werden.',
 'Auf hohe Vermögen soll wieder eine Steuer erhoben werden.',
 'Bei der Videoüberwachung öffentlicher Plätze soll Gesichtserkennungssoftware eingesetzt werden dürfen.',
 'Auch Ehepaare ohne Kinder sollen weiterhin steuerlich begünstigt werden.',
 'Ökologische Landwirtschaft soll stärker gefördert werden als konventionelle Landwirtschaft.',
 'Islamische Verbände sollen als Religionsgemeinschaften staatlich anerkannt werden können.',
 'Der staatlich festgelegte Preis für den Ausstoß von CO2 beim Heizen und Autofahren soll stärker steigen als geplant.',
 'Die Schuldenbremse im Grundgesetz soll beibehalten werden.',
 'Asyl soll weiterhin nur politisch Verfolgten gewährt werden.',
 'Der gesetzliche Mindestlohn soll spätestens im Jahr 2022 auf mindestens 12 Euro erhöht werden.',
 'Der Flugverkehr soll höher besteuert werden.',
 'Unternehmen sollen selbst entscheiden, ob sie ihren Beschäftigten das Arbeiten im Homeoffice erlauben.']

with gr.Blocks(title="test") as demo:
    chatbot = gr.Chatbot(value=[[None, "Hallo mein Name ist AVAH (Artificial Voting Advice Helper), deine interaktive Hilfe zum Wahl-O-Maten für die Bundestagswahl 2021. Wenn du Fragen zu Thesen der Wahlhilfe hast, wähle gerne die besagt These aus. Ansonsten kannst du mir auch unabhängig von bestimmten Thesen Fragen stellen."]], render_markdown=True)
    dropdown = gr.Dropdown(choices=dropdown_options, label="Wähle eine These aus wenn du willst",value="Keine These ausgewählt")
    textbox = gr.Textbox(label="Deine Frage")

    def combined_response(prompt, history, dropdown_value):
        answer = response(prompt, history, dropdown_value)
        return history + [[prompt, answer]]

    submit_button = gr.Button("Submit")

    submit_button.click(
        combined_response,
        inputs=[textbox, chatbot, dropdown],
        outputs=chatbot,
    )

demo.launch()

print("Interface up and running!")