import time  # to simulate a real time data, time loop
from os import listdir
from os.path import isfile, join
import numpy as np  # np mean, np random
import pandas as pd  # read csv, df manipulation
import plotly.express as px  # interactive charts
from plotly import graph_objs as go
import streamlit as st  # 🎈 data web app development
import plotly.figure_factory as ff
import numpy as np
from collections import Counter


print("Make sure to activate your VPN before running this script")

st.set_page_config(
    page_title="GroqFlow Progress Tracker",
    page_icon="🚀",
    layout="wide",
)


# Session State variables:
state = st.session_state
if "INFO_CLOSED" not in state:
    state.INFO_CLOSED = False

# dashboard title
st.title("GroqFlow Progress Tracker 🚀")

# Custom chart colors (https://plotly.com/python/discrete-color/)
colorway = ["#3366cc", "#FF7F0E"]


def add_filter(data_frame_list, name, label, options, num_cols=1):
    st.markdown(f"#### {name}")

    cols = st.columns(num_cols)
    instantiated_checkbox = []
    for idx in range(len(options)):
        with cols[idx % num_cols]:
            instantiated_checkbox.append(st.checkbox(options[idx], False))

    all_options = set(data_frame_list[-1][label])
    selected_options = [
        options[idx] for idx, checked in enumerate(instantiated_checkbox) if checked
    ]

    # The last checkbox will always correspond to "other"
    if instantiated_checkbox[-1]:
        selected_options = selected_options[:-1]
        other_options = [x for x in all_options if x not in options]
        selected_options = set(selected_options + other_options)

    if len(selected_options) > 0:
        for idx in range(len(data_frame_list)):
            data_frame_list[idx] = data_frame_list[idx][
                [
                    any([x == model_entry for x in selected_options])
                    for model_entry in data_frame_list[idx][label]
                ]
            ]
    return data_frame_list


with st.sidebar:

    st.markdown("# Filters")

    test_type = st.radio(
        "Test Type",
        ("Daily Tests (100 models)", "Monthly Tests (500+ models)"),
    )
    if test_type == "Daily Tests (100 models)":
        selected_test_type = "daily"
        report_folder = "reports/daily"
    else:
        selected_test_type = "monthly"
        report_folder = "reports/monthly"

    # Get ML Agility reports
    reports = sorted(
        [f for f in listdir(report_folder) if isfile(join(report_folder, f))]
    )

    selected_report = st.selectbox("Test date", reports, index=len(reports) - 1)
    selected_report_idx = reports.index(selected_report)
    prev_report = reports[max(0, selected_report_idx - 1)]
    mla_report = pd.read_csv(f"{report_folder}/{selected_report}")
    prev_mla_report = pd.read_csv(f"{report_folder}/{prev_report}")

    # Add chips filter
    num_chips_options = ["1", "2", "4", "8", "16", "32+"]
    mla_report = mla_report.astype({"chips_used": str})
    prev_mla_report = prev_mla_report.astype({"chips_used": str})
    mla_report, prev_mla_report = add_filter(
        [mla_report, prev_mla_report],
        "Number of GroqChips™",
        label="chips_used",
        options=num_chips_options,
        num_cols=3,
    )

    # Add author filter
    authors = [
        "google",
        "apple",
        "facebook",
        "openai",
        "microsoft",
        "huggingface",
        "CompVis",
        "others",
    ]
    mla_report, prev_mla_report = add_filter(
        [mla_report, prev_mla_report],
        "Authors",
        label="author",
        options=authors,
        num_cols=2,
    )

    # Add task filter
    tasks = [
        "Image Classification",
        "Translation",
        "Image Segmentation",
        "Fill-Mask",
        "Text-to-Image",
        "Token Classification",
        "Sentence Similarity",
        "Audio Classification",
        "Question Answering",
        "Summarization",
        "other",
    ]
    mla_report, prev_mla_report = add_filter(
        [mla_report, prev_mla_report], "Tasks", label="task", options=tasks
    )


def detailed_progress_list(df_new, df_old, filter=None):
    return
    """
    if filter is not None:
        df_new = df_new[(df_new[filter] == True)]
        df_old = df_old[(df_old[filter] == True)]

    progress = df_new[~(df_new["hash"].isin(df_old["hash"]))].reset_index(drop=True)
    regression = df_old[~(df_old["hash"].isin(df_new["hash"]))].reset_index(drop=True)

    for model_name in progress["model_name"]:
        st.markdown(
            f'<span style="color:green">↑ {model_name}</span>',
            unsafe_allow_html=True,
        )
    for model_name in regression["model_name"]:
        st.markdown(
            f'<span style="color:red">↓ {model_name}</span>',
            unsafe_allow_html=True,
        )
    """


# creating a single-element container
placeholder = st.empty()

with placeholder.container():

    st.markdown("## Summary Results")
    # create three columns
    kpi = st.columns(7)
    model_details = st.columns(7)

    # fill in those three columns with respective metrics or KPIs
    kpi[0].metric(
        label="All models",
        value=len(mla_report),
        delta=len(mla_report) - len(prev_mla_report),
    )
    if selected_test_type == "daily":
        with model_details[0]:
            detailed_progress_list(mla_report, prev_mla_report)

    kpi[1].metric(
        label="Convert to ONNX",
        value=np.sum(mla_report["base_onnx"]),
        delta=int(
            np.sum(mla_report["base_onnx"]) - np.sum(prev_mla_report["base_onnx"])
        ),
    )
    if selected_test_type == "daily":
        with model_details[1]:
            detailed_progress_list(mla_report, prev_mla_report, "base_onnx")

    kpi[2].metric(
        label="Optimize ONNX file",
        value=np.sum(mla_report["optimized_onnx"]),
        delta=int(
            np.sum(mla_report["optimized_onnx"])
            - np.sum(prev_mla_report["optimized_onnx"])
        ),
    )
    if selected_test_type == "daily":
        with model_details[2]:
            detailed_progress_list(mla_report, prev_mla_report, "optimized_onnx")

    kpi[3].metric(
        label="All ops supported",
        value=np.sum(mla_report["all_ops_supported"]),
        delta=int(
            np.sum(mla_report["all_ops_supported"])
            - np.sum(prev_mla_report["all_ops_supported"])
        ),
    )
    if selected_test_type == "daily":
        with model_details[3]:
            detailed_progress_list(mla_report, prev_mla_report, "all_ops_supported")

    kpi[4].metric(
        label="Converts to FP16",
        value=np.sum(mla_report["fp16_onnx"]),
        delta=int(
            np.sum(mla_report["fp16_onnx"]) - np.sum(prev_mla_report["fp16_onnx"])
        ),
    )
    if selected_test_type == "daily":
        with model_details[4]:
            detailed_progress_list(mla_report, prev_mla_report, "fp16_onnx")

    kpi[5].metric(
        label="Compiles",
        value=np.sum(mla_report["compiles"]),
        delta=int(np.sum(mla_report["compiles"]) - np.sum(prev_mla_report["compiles"])),
    )
    if selected_test_type == "daily":
        with model_details[5]:
            detailed_progress_list(mla_report, prev_mla_report, "compiles")

    kpi[6].metric(
        label="Assembles",
        value=np.sum(mla_report["assembles"]),
        delta=int(
            np.sum(mla_report["assembles"]) - np.sum(prev_mla_report["assembles"])
        ),
    )
    if selected_test_type == "daily":
        with model_details[6]:
            detailed_progress_list(mla_report, prev_mla_report, "assembles")

    cols = st.columns(2)
    with cols[0]:

        compiler_errors = mla_report[mla_report["compiler_error"] != "-"][
            "compiler_error"
        ]
        compiler_errors = Counter(compiler_errors)
        st.markdown("""#### Top compiler issues""")
        if len(compiler_errors) > 0:
            compiler_errors = pd.DataFrame.from_dict(
                compiler_errors, orient="index"
            ).reset_index()
            compiler_errors = compiler_errors.set_axis(
                ["error", "count"], axis=1, inplace=False
            )

            fig = px.bar(
                compiler_errors, x="count", y="error", orientation="h", height=400
            )
            st.plotly_chart(fig, use_container_width=True)
        else:
            st.markdown("""No compiler errors found :tada:""")

    with cols[1]:
        # Add parameters histogram
        all_models = [float(x) / 1000000 for x in mla_report["params"] if x != "-"]

        assembled_models = mla_report[mla_report["assembles"] == True]
        assembled_models = [
            float(x) / 1000000 for x in assembled_models["params"] if x != "-"
        ]
        hist_data = []
        group_labels = []
        if all_models != []:
            hist_data.append(all_models)
            group_labels.append("Models we tried compiling")

        if assembled_models != []:
            hist_data.append(assembled_models)
            group_labels.append("Assembled models")

        st.markdown("""#### Assembled models vs. Parameters (in millions)""")

        if len(assembled_models) > 1:

            fig = ff.create_distplot(
                hist_data,
                group_labels,
                bin_size=[25, 25],
                histnorm="",
            )
            # fig.layout.update(title="Assembled models vs. Parameters (in millions)")
            fig.layout.update(xaxis_title="Parameters in millions")
            fig.layout.update(yaxis_title="count")
            fig.update_xaxes(range=[1, 1000])
            st.plotly_chart(fig, use_container_width=True)
        else:
            st.markdown("""Need at least one assembled model to show this graph 😅""")

    if "tsp_gpu_compute_ratio" in mla_report and "tsp_gpu_e2e_ratio" in mla_report:
        cols = st.columns(2)
        with cols[0]:
            # GPU Acceleration plot
            st.markdown("""#### Speedup of GroqChip™ compared to A100 GPUs""")

            # Prepare data
            df = mla_report[
                ["model_name", "tsp_gpu_compute_ratio", "tsp_gpu_e2e_ratio"]
            ]
            df = df.sort_values(by=["model_name"])
            df = df[(df.tsp_gpu_compute_ratio != "-")]
            df = df[(df.tsp_gpu_e2e_ratio != "-")]
            df["tsp_gpu_compute_ratio"] = df["tsp_gpu_compute_ratio"].astype(float)
            df["tsp_gpu_e2e_ratio"] = df["tsp_gpu_e2e_ratio"].astype(float)

            data = [
                go.Bar(
                    x=df["model_name"],
                    y=df["tsp_gpu_compute_ratio"],
                    name="Compute only",
                ),
                go.Bar(
                    x=df["model_name"],
                    y=df["tsp_gpu_e2e_ratio"],
                    name="Compute + estimated I/O",
                ),
            ]

            layout = go.Layout(
                barmode="overlay",
                yaxis_title="Speedup compared to A100 GPU",
                colorway=colorway,
            )

            fig = dict(data=data, layout=layout)
            st.plotly_chart(fig, use_container_width=True)

            st.markdown(
                "<sup>*</sup>Estimated I/O does NOT include delays caused by Groq's runtime.",
                unsafe_allow_html=True,
            )

        with cols[1]:
            # Show stats
            st.markdown(
                f"""<br><br><br><br><br><br>
                <p style="font-family:sans-serif; font-size: 20px;text-align: center;">Average speedup of GroqChip™ considering compute only:</p>
                <p style="font-family:sans-serif; color:#3366cc; font-size: 26px;text-align: center;"> {round(df["tsp_gpu_compute_ratio"].mean(),2)}x</p>
                <p style="font-family:sans-serif; color:#3366cc; font-size: 20px;text-align: center;"> min {round(df["tsp_gpu_compute_ratio"].min(),2)}x; max {round(df["tsp_gpu_compute_ratio"].max(),2)}x</p>
                <br><br>
                <p style="font-family:sans-serif; font-size: 20px;text-align: center;">Average speedup of GroqChip™ considering compute + estimated I/O<sup>*</sup>:</p>
                <p style="font-family:sans-serif; color:#FF7F0E; font-size: 26px;text-align: center;"> {round(df["tsp_gpu_e2e_ratio"].mean(),2)}x</p>
                <p style="font-family:sans-serif; color:#FF7F0E; font-size: 20px;text-align: center;"> min {round(df["tsp_gpu_e2e_ratio"].min(),2)}x; max {round(df["tsp_gpu_e2e_ratio"].max(),2)}x</p>""",
                unsafe_allow_html=True,
            )

    st.markdown("### Detailed Data View")
    st.markdown(
        "**Model selection**: All workloads were obtained from models cards available at huggingface.co/models. Input shapes corresponds exactly to those used by the Huggingface model cards. Some of those input shapes might be small, causing the compilation process to be easier than when reasonably-sized input shapes are used.",
        unsafe_allow_html=True,
    )
    model_name = st.text_input("", placeholder="Filter model by name")
    if model_name != "":
        mla_report = mla_report[[model_name in x for x in mla_report["model_name"]]]

    # Select which columns to show
    selected_cols = list(mla_report.columns)
    # remove_cols = (
    #    "tsp_e2e_latency",
    #    "gpu_e2e_latency",
    #    "tsp_gpu_e2e_ratio",
    # )
    # for item in remove_cols:
    #    if item in selected_cols:
    #        selected_cols.remove(item)
    st.dataframe(
        mla_report[selected_cols], height=min((len(mla_report) + 1) * 35, 35 * 21)
    )