import io
import json
import re

import gradio as gr
import pandas as pd
import plotly
import plotly.express as px
from pandas.api.types import is_numeric_dtype

from pipeline.config import LLMBoardConfig, QueriesConfig
from app_constants import README, JS, TIME_PERIODS_EXPLANATION_DF

queries_config = QueriesConfig()

output_types_df = pd.DataFrame(
    {"Output Type": queries_config.query_template.keys(), "Added text": queries_config.query_template.values()}
)

summary_df: pd.DataFrame = pd.read_csv("data/summary.csv")
time_of_day_comparison_df = pd.read_csv("data/time_of_day_comparison.csv")
general_plots = pd.read_csv("data/general_plots.csv")
model_costs_df = pd.read_csv("data/model_costs.csv")
time_of_day_plots = pd.read_csv("data/time_of_day_plots.csv")
summary_metrics_plots = pd.read_csv("data/summary_metrics_plots.csv")
output_plots = pd.read_csv("data/output_plots.csv")
combined_plots = pd.read_csv("data/combined_plots.csv")

searched_query = ""
collapse_languages = False
collapse_output_method = False


def filter_dataframes(input: str):
    global searched_query
    input = input.lower()
    searched_query = input
    return get_updated_dataframes()


def collapse_languages_toggle():
    global collapse_languages
    if collapse_languages:
        collapse_languages = False
        button_text = "Collapse languages"
    else:
        collapse_languages = True
        button_text = "Un-collapse languages"
    return get_updated_dataframes()[0], button_text


def collapse_output_method_toggle():
    global collapse_output_method
    if collapse_output_method:
        collapse_output_method = False
        button_text = "Collapse output method"
    else:
        collapse_output_method = True
        button_text = "Un-collapse output method"
    return get_updated_dataframes()[0], button_text

def filter_dataframe_by_models(df, searched_model_names):
    if not searched_model_names:
        return df
    filter_series = df.model == "" # False values
    for n in searched_model_names:
        filter_series = filter_series | df.model.str.lower().str.contains(n)
    return df[filter_series]

def get_updated_dataframes():
    global collapse_languages, collapse_output_method, searched_query, summary_df, time_of_day_comparison_df, model_costs_df

    summary_df_columns = summary_df.columns.to_list()
    group_columns = LLMBoardConfig().group_columns.copy()
    if collapse_languages:
        summary_df_columns.remove("language")
        group_columns.remove("language")
    if collapse_output_method:
        summary_df_columns.remove("template_name")
        group_columns.remove("template_name")

    summary_df_processed = summary_df[summary_df_columns].groupby(by=group_columns).mean().reset_index()

    searched_model_names = searched_query.split("|")
    searched_model_names = [n.lower().strip() for n in searched_model_names]
    searched_model_names = [n for n in searched_model_names if n]

    def for_dataframe(df):
        return dataframe_style(filter_dataframe_by_models(df, searched_model_names))

    return (
        for_dataframe(summary_df_processed),
        for_dataframe(time_of_day_comparison_df),
        for_dataframe(model_costs_df),
    )


def dataframe_style(df: pd.DataFrame):
    df = df.copy()
    column_formats = {}
    new_column_names = []

    for column in df.columns:
        if is_numeric_dtype(df[column]):
            if column == "execution_time":
                column_formats[column] = "{:.4f}"
            else:
                column_formats[column] = "{:.2f}"
        new_column_name = snake_case_to_title(column)
        if "time" in column and column != "time_of_day":
            new_column_name += " (Seconds)"
        elif "chunk" in column:
            new_column_name += " (Characters)"
        new_column_names.append(new_column_name)
    
    df.columns = new_column_names
    df = df.style.format(column_formats, na_rep="")
    return df

def snake_case_to_title(text):
    # Convert snake_case to title-case
    words = re.split(r"_", text)
    title_words = [word.capitalize() for word in words]
    return " ".join(title_words)

plots = []

def display_plot(plot_df_row):
    row = dict(plot_df_row)
    plot = plotly.io.from_json(row["plot_json"])
    plot.update_layout(autosize=True)
    return (gr.Plot(plot, label=row["header"], scale=1), plot)

def display_filtered_plot(plot_df_row):
    row = dict(plot_df_row)
    plot_element, plot = display_plot(plot_df_row)
    plots.append((plot_element, plot, row))
    if "description" in row and pd.notna(row["description"]):
        gr.Markdown(str(row["description"]))

def filter_plots(searched_query: str):
    searched_model_names = searched_query.split("|")
    searched_model_names = [n.lower().strip() for n in searched_model_names]
    searched_model_names = [n for n in searched_model_names if n]

    results = []
    for plot_display, plot, row in plots:
        visible = True
        if "df" in row and pd.notna(row["df"]):
            buffer = io.StringIO(row["df"])
            df = pd.read_csv(buffer)
            df = filter_dataframe_by_models(df, searched_model_names)
            plot_constructor = px.bar
            if "plot_type" in row and pd.notna(row["plot_type"]) and row["plot_type"]:
                if row["plot_type"] == "scatter":
                    plot_constructor = px.scatter
            plot = plot_constructor(df, **json.loads(row["arguments"]))
            plot.update_layout(autosize=True)
        elif "for model" in row["header"] and searched_model_names:
            plot_model = row["header"].split("for model")[1].lower()
            if not any(n in plot_model for n in searched_model_names):
                visible = False

        results.append(gr.Plot(plot, visible=visible))

    return results

with gr.Blocks(theme=gr.themes.Default(text_size="lg"), js=JS) as demo:
    gr.HTML("<h1>Performance LLM Board</h1>")

    with gr.Row():
        filter_textbox = gr.Textbox(label="Model name parts *", scale=2, elem_id="filter-textbox")
        filter_button = gr.Button("Filter", scale=1, elem_id="filter-button")
        with gr.Column(scale=1):
            open_ai_button = gr.Button("Compare Open AI models", elem_id="open-ai-button", scale=1)
            google_button = gr.Button("Compare Google Models", elem_id="google-button", scale=1)
            # gr.Button("Open Models", size="sm")
    gr.Markdown(
        '&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;\* You can use `|` operator to display multiple models at once, for example "gpt|mistral|zephyr"'
    )

    with gr.Tab("About this project"):
        gr.Markdown(
            README.format(
                queries_config.base_query_template.replace("```", "'''"), output_types_df.to_markdown(index=False)
            )
        )
    with gr.Tab("Performance by time of the day"):
        # display only first plot for all models
        time_of_day_plots[0:1].apply(display_filtered_plot, axis=1)
        time_periods_explanation_ui = gr.DataFrame(
            dataframe_style(TIME_PERIODS_EXPLANATION_DF), label="Times of day ranges"
        )
        time_of_day_comparison_ui = gr.DataFrame(dataframe_style(time_of_day_comparison_df), label="Time of day")
        gr.Markdown(
            """\
These measurements were made by testing the models using the same dataset as in the other comparisons every hour for 24 hours.
                    
Execution time refers to averaged time needed to execute one query.

Hours and times of day in the table and in the plot are based on Central European Time.
                    
Measurements were made during a normal work week.
"""
        )
        # display rest of the plots
        time_of_day_plots[1:].apply(display_filtered_plot, axis=1)
    with gr.Tab("Output characteristics"):
        with gr.Row():
            collapse_languages_button = gr.Button("Collapse languages")
            collapse_output_method_button = gr.Button("Collapse output method")
        summary_ui = gr.DataFrame(dataframe_style(summary_df), label="Output characteristics")
        gr.Markdown(
            """\
This table compares output characteristics of different models which include execution time, output size and chunking of the output. Some providers and models don't support output chunking, in this case chunk related fields are left empty.

Execution time refers to averaged time needed to execute one query.

To count words we split the output string by whitespace `\w` regex character.

Chunk sizes are measured in the characters count."""
        )
        output_plots.apply(display_filtered_plot, axis=1)
    with gr.Tab("Costs comparison"):
        models_costs_ui = gr.DataFrame(dataframe_style(model_costs_df), label="Costs comparison")
        gr.Markdown(
            """\
Provider pricing column contains pricing from the website of the provider.

Hugging Face Inference Endpoints are charged by hour so to compare different providers together, 
for models hosted this way we calculated "Cost Per Token" column using data collected during the experiment.

Note that pause and resume time cost was not included in the "Cost Per Token" column calculation.
            """
        )
        general_plots[general_plots.plot_name == "execution_costs"].apply(display_filtered_plot, axis=1)
    with gr.Tab("Summary metrics"):
        summary_metrics_plots.apply(display_filtered_plot, axis=1)
    with gr.Tab("Context length and parameters count"):
        general_plots[general_plots.plot_name != "execution_costs"].apply(display_filtered_plot, axis=1)
        gr.Markdown(
            """
LLM models context length and parameters count are based on release blogs and documentation of their respective developers.

A lot of models had to be omitted due to their developers not disclosing their parameters count.

Mainly OpenAI's GPT models and Google's Palm 2.
"""
        )
    with gr.Tab("Combined plots"):
        with gr.Row():
            choices = combined_plots.header
            choices = choices[choices.str.contains("for model")]
            choices = choices.str.split("for model").apply(lambda x: x[1])
            def handle_dropdown(dropdown, plot_element):
                def dropdown_change_handler(value):
                    for _, row in combined_plots.iterrows():
                        if value in row["header"]:
                            return display_plot(row)[0]
                dropdown.change(
                    fn=dropdown_change_handler,
                    inputs=[dropdown],
                    outputs=[plot_element],
                    api_name="dropdown_change_handler",
                )
            with gr.Column():
                dropdown = gr.Dropdown(choices.tolist(), label="First model for comparison", value=choices.iloc[0])
                plot_element, plot = display_plot(combined_plots.iloc[0])
                handle_dropdown(dropdown, plot_element)
            with gr.Column():
                dropdown = gr.Dropdown(choices.tolist(), label="Second model for comparison", value=choices.iloc[1])
                plot_element, plot = display_plot(combined_plots.iloc[1])
                handle_dropdown(dropdown, plot_element)
        gr.Markdown("""
Radial plots are used to compare the most important aspects of each model researched on this board using single images.

All values are normalized and scaled into 0.25 to 1 range, 0 is left for unknown values.
                    
To compare the parameters more thoroughly use the filtering box on top of this page and inspect individual tabs.
                    
In addition to side by side comparison all of the radial plots are displayed below.
""")
        combined_plots.apply(display_filtered_plot, axis=1)
    filter_button.click(
        fn=filter_dataframes,
        inputs=filter_textbox,
        outputs=[summary_ui, time_of_day_comparison_ui, models_costs_ui],
        api_name="filter_dataframes",
    )
    filter_textbox.submit(
        fn=filter_dataframes,
        inputs=filter_textbox,
        outputs=[summary_ui, time_of_day_comparison_ui, models_costs_ui],
        api_name="filter_dataframes",
    )
    filter_button.click(
        fn=filter_plots,
        inputs=filter_textbox,
        outputs=[v[0] for v in plots],
        api_name="filter_plots",
    )
    filter_textbox.submit(
        fn=filter_plots,
        inputs=filter_textbox,
        outputs=[v[0] for v in plots],
        api_name="filter_plots",
    )
    collapse_languages_button.click(
        fn=collapse_languages_toggle,
        outputs=[summary_ui, collapse_languages_button],
        api_name="collapse_languages_toggle",
    )
    collapse_output_method_button.click(
        fn=collapse_output_method_toggle,
        outputs=[summary_ui, collapse_output_method_button],
        api_name="collapse_output_method_toggle",
    )

demo.launch()