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

README = """
Projects compares different large language models and their providers for real time applications and mass data processing. 
While other benchmarks compare LLMs on different human intelligence tasks this benchmark focus on features related to business and engineering aspects such as response times, pricing and data streaming capabilities. 

To preform evaluation we chose a task of newspaper articles summarization from [GEM/xlsum](https://huggingface.co/datasets/GEM/xlsum) dataset as it represents a very standard type of task where model has to understand unstructured natural language text, process it and output text in a specified format.
For this version we chose English, Polish and Japanese languages, with Japanese representing languages using logographic alphabets. This enable us also validate the effectiveness of the LLM for different language groups.

Each of the models was asked to summarize the text using the following prompt:

```
{}
```

Where {{language}} stands for original language of the text as we wanted to avoid the model translating the text to English during summarization.

LLM was asked to return the output in three formats: markdown, json and function call. Note that currently function calls are only supported by Open AI API.
To do that we added following text to the query:

{}

All of the call were made from the same machine with the same internet connection with usage of the LiteLLM library which may adds some time overhead compared to pure curl calls. Call were made from Poland, UTC +1.

Please take a look at the following project and let us know if you have any questions or suggestions.
"""

time_periods_explanation_df = pd.DataFrame(
    {
        "time_of_day": [
            "early morning",
            "morning",
            "afternoon",
            "late afternoon",
            "evening",
            "late evening",
            "midnight",
            "night",
        ],
        "hour_range": ["6-8", "9-11", "12-14", "15-17", "18-20", "21-23", "0-2", "3-5"],
    }
)

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")
output_plots = pd.read_csv("data/output_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 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 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 dataframes()[0], button_text

def filter_dataframe(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 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(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)


filter_textbox = gr.Textbox(label="Model name parts *", scale=2)
filter_button = gr.Button("Filter", scale=1)
collapse_languages_button = gr.Button("Collapse languages")
collapse_output_method_button = gr.Button("Collapse output method")
last_textbox = 0
plots = []
single_model_plots = []


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(df, searched_model_names)
            plot = px.bar(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


def display_plot(plot_df_row):
    row = dict(plot_df_row)
    plot = plotly.io.from_json(row["plot_json"])
    plot.update_layout(autosize=True)
    plots.append((gr.Plot(plot, label=row["header"], scale=1), plot, row))
    if "description" in row and pd.notna(row["description"]):
        gr.Markdown(str(row["description"]))


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

    with gr.Row():
        filter_textbox.render()
        filter_button.render()
    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_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_plot, axis=1)
    with gr.Tab("Output characteristics"):
        with gr.Row():
            collapse_languages_button.render()
            collapse_output_method_button.render()
        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_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_plot, axis=1)
    with gr.Tab("Context length and parameters count"):
        general_plots[general_plots.plot_name != "execution_costs"].apply(display_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.
"""
        )
    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()