import re
import uuid

import pandas as pd
import streamlit as st
import re
import matplotlib.pyplot as plt

import subprocess
import sys
import io

from utils.default_values import get_system_prompt, get_guidelines_dict
from utils.epfl_meditron_utils import get_llm_response, gptq_model_options
from utils.openai_utils import get_available_engines, get_search_query_type_options

from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from sklearn.metrics import classification_report

DATA_FOLDER = "data/"

POC_VERSION = "0.1.0"
MAX_QUESTIONS = 10
AVAILABLE_LANGUAGES = ["DE", "EN", "FR"]

st.set_page_config(page_title='Medgate Whisper PoC', page_icon='public/medgate.png')

# Azure apparently truncates message if longer than 200, see
MAX_SYSTEM_MESSAGE_TOKENS = 200


def format_question(q):
    res = q

    # Remove numerical prefixes, if any, e.g. '1. [...]'
    if re.match(r'^[0-9].\s', q):
        res = res[3:]

    # Replace doc reference by doc name
    if len(st.session_state["citations"]) > 0:
        for source_ref in re.findall(r'\[doc[0-9]+\]', res):
            citation_number = int(re.findall(r'[0-9]+', source_ref)[0])
            citation_index = citation_number - 1 if citation_number > 0 else 0
            citation = st.session_state["citations"][citation_index]
            source_title = citation["title"]
            res = res.replace(source_ref, '[' + source_title + ']')

    return res.strip()


def get_text_from_row(text):
    res = str(text)
    if res == "nan":
        return ""
    return res
def get_questions_from_df(df, lang, test_scenario_name):
    questions = []
    for i, row in df.iterrows():
        questions.append({
            "question": row[lang + ": Fragen"],
            "answer": get_text_from_row(row[test_scenario_name]),
            "question_id": uuid.uuid4()
        })
    return questions


def get_questions(df, lead_symptom, lang, test_scenario_name):
    print(str(st.session_state["lead_symptom"]) + " -> " + lead_symptom)
    print(str(st.session_state["scenario_name"]) + " -> " + test_scenario_name)
    if st.session_state["lead_symptom"] != lead_symptom or st.session_state["scenario_name"] != test_scenario_name:
        st.session_state["lead_symptom"] = lead_symptom
        st.session_state["scenario_name"] = test_scenario_name
        symptom_col_name = st.session_state["language"] + ": Symptome"
        df_questions = df[(df[symptom_col_name] == lead_symptom)]
        st.session_state["questions"] = get_questions_from_df(df_questions, lang, test_scenario_name)

    return st.session_state["questions"]


def display_streamlit_sidebar():
    st.sidebar.title("Local LLM PoC " + str(POC_VERSION))

    st.sidebar.write('**Parameters**')
    form = st.sidebar.form("config_form", clear_on_submit=True)

    model_name_or_path = form.selectbox("Select model", gptq_model_options())
    
    temperature = form.slider(label="Temperature", min_value=0.0, max_value=1.0, step=0.01, value=st.session_state["temperature"])
    do_sample = form.checkbox('do_sample', value=st.session_state["do_sample"])
    top_p = form.slider(label="top_p", min_value=0.0, max_value=1.0, step=0.01, value=st.session_state["top_p"])
    top_k = form.slider(label="top_k", min_value=1, max_value=1000, step=1, value=st.session_state["top_k"])
    max_new_tokens = form.slider(label="max_new_tokens", min_value=32, max_value=512, step=1, value=st.session_state["max_new_tokens"])
    repetition_penalty = form.slider(label="repetition_penalty", min_value=0.0, max_value=5.0, step=0.01, value=st.session_state["repetition_penalty"])

    submitted = form.form_submit_button("Start session")
    if submitted:
        print('Parameters updated...')
        st.session_state['session_started'] = True

        st.session_state["session_events"] = []
        st.session_state["model_name_or_path"] = model_name_or_path
        st.session_state["temperature"] = temperature
        st.session_state["do_sample"] = do_sample
        st.session_state["top_p"] = top_p
        st.session_state["top_k"] = top_k
        st.session_state["max_new_tokens"] = max_new_tokens
        st.session_state["repetition_penalty"] = repetition_penalty

        st.rerun()
      

def init_session_state():
    print('init_session_state()')
    st.session_state['session_started'] = False
    st.session_state["session_events"] = []
    st.session_state["model_name_or_path"] = "TheBloke/meditron-7B-GPTQ"
    st.session_state["temperature"] = 0.01
    st.session_state["do_sample"] = True
    st.session_state["top_p"] = 0.95
    st.session_state["top_k"] = 40
    st.session_state["max_new_tokens"] = 512
    st.session_state["repetition_penalty"] = 1.1
    st.session_state["system_prompt"] = "You are a medical expert that provides answers for a medically trained audience"
    st.session_state["prompt"] = ""

def get_genders():
    return ['Male', 'Female']


def display_session_overview():
    st.subheader('History of LLM queries')
    st.write(st.session_state["llm_messages"])
    st.subheader("Session costs overview")

    df_session_overview = pd.DataFrame.from_dict(st.session_state["session_events"])
    st.write(df_session_overview)

    if "prompt_tokens" in df_session_overview:
        prompt_tokens = df_session_overview["prompt_tokens"].sum()
        st.write("Prompt tokens: " + str(prompt_tokens))

        prompt_cost = df_session_overview["prompt_cost_chf"].sum()
        st.write("Prompt CHF: " + str(prompt_cost))

        completion_tokens = df_session_overview["completion_tokens"].sum()
        st.write("Completion tokens: " + str(completion_tokens))

        completion_cost = df_session_overview["completion_cost_chf"].sum()
        st.write("Completion CHF: " + str(completion_cost))

        completion_cost = df_session_overview["total_cost_chf"].sum()
        st.write("Total costs CHF: " + str(completion_cost))

        total_time = df_session_overview["response_time"].sum()
        st.write("Total compute time (ms): " + str(total_time))


def plot_report(title, expected, predicted, display_labels):
    st.markdown('#### ' + title)
    conf_matrix = confusion_matrix(expected, predicted, labels=display_labels)
    conf_matrix_plot = ConfusionMatrixDisplay(confusion_matrix=conf_matrix, display_labels=display_labels)
    conf_matrix_plot.plot()
    st.pyplot(plt.gcf())

    report = classification_report(expected, predicted, output_dict=True)
    df_report = pd.DataFrame(report).transpose()
    st.write(df_report)

    df_rp = df_report
    df_rp = df_rp.drop('support', axis=1)
    df_rp = df_rp.drop(['accuracy', 'macro avg', 'weighted avg'])

    try:
        ax = df_rp.plot(kind="bar", legend=True)
        for container in ax.containers:
            ax.bar_label(container, fontsize=7)
        plt.xticks(rotation=45)
        plt.legend(loc=(1.04, 0))
        st.pyplot(plt.gcf())
    except Exception as e:
        # Out of bounds
        pass


def get_prompt_format(model_name):
    if model_name == "TheBloke/Llama-2-13B-chat-GPTQ":
        return '''[INST] <<SYS>>
                You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
                <</SYS>>
                {prompt}[/INST]

                '''
    if model_name == "TheBloke/Llama-2-7B-Chat-GPTQ":
        return "[INST] <<SYS>>You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.＼n<</SYS>>＼n{prompt}[/INST]"
    
    if model_name == "TheBloke/meditron-7B-GPTQ" or model_name == "TheBloke/meditron-70B-GPTQ":
        return '''<|im_start|>system
                {system_message}<|im_end|>
                <|im_start|>user
                {prompt}<|im_end|>
                <|im_start|>assistant'''
    
    return ""

def format_prompt(template, system_message, prompt):
    if template == "":
        return f"{system_message} {prompt}"
    return template.format(system_message=system_message, prompt=prompt)

def display_llm_output():
    st.header("LLM")

    form = st.form('llm')

    prompt_format_str = get_prompt_format(st.session_state["model_name_or_path"])
    prompt_format = form.text_area('Prompt format', value=prompt_format_str)
    system_prompt = ""#form.text_area('System prompt', value=st.session_state["system_prompt"])
    prompt = form.text_area('Prompt', value=st.session_state["prompt"])

    submitted = form.form_submit_button('Submit')

    if submitted:
        st.session_state["system_prompt"] = system_prompt
        st.session_state["prompt"] = prompt
        formatted_prompt = format_prompt(prompt_format, system_prompt, prompt)
        print(f"Formatted prompt: {format_prompt}")
        llm_response = get_llm_response(
            st.session_state["model_name"],
            st.session_state["temperature"],
            st.session_state["do_sample"],
            st.session_state["top_p"],
            st.session_state["top_k"],
            st.session_state["max_new_tokens"],
            st.session_state["repetition_penalty"],
            formatted_prompt)
        st.write(llm_response)
        st.write('Done displaying LLM response')

def main():
    print('Running Local LLM PoC Streamlit app...')
    session_inactive_info = st.empty()
    if "session_started" not in st.session_state or not st.session_state["session_started"]:
        init_session_state()
        display_streamlit_sidebar()
    else:
        display_streamlit_sidebar()
        session_inactive_info.empty()
        display_llm_output()
        display_session_overview()


if __name__ == '__main__':
    main()