import json

import gradio as gr
import pandas as pd

with open('results.json', 'r') as file:
    results = json.load(file)

models = [key for key in results.keys()]
demo = gr.Blocks()


from random import randint, random

food_rating_data = pd.DataFrame(
    {
        "cuisine": [["Italian", "Mexican", "Chinese"][i % 3] for i in range(100)],
        "rating": [random() * 4 + 0.5 * (i % 3) for i in range(100)],
        "price": [randint(10, 50) + 4 * (i % 3) for i in range(100)],
        "wait": [random() for i in range(100)],
    }
)

df = pd.DataFrame.from_dict(results[models[0]]["main-net"], orient = "index").reset_index()
df.columns = ["Step", "Loss"]
df["Step"] = pd.to_numeric(df["Step"])
df["Test"] = "Main-net"

if "baseline" in results[models[0]]:
    df_baseline = pd.DataFrame.from_dict(results[models[0]]["baseline"], orient = "index").reset_index()
    df_baseline.columns = ["Step", "Loss"]
    df_baseline["Step"] = pd.to_numeric(df_baseline["Step"])
    df_baseline["Test"] = "Baseline"

    df = pd.concat([df, df_baseline])

def return_results(model_name):
    print(model_name)
    df = pd.DataFrame.from_dict(results[model_name]["main-net"], orient = "index").reset_index()
    df.columns = ["Step", "Loss"]
    df["Step"] = pd.to_numeric(df["Step"])
    df["Test"] = "Main-net"

    if "baseline" in results[model_name]:
        df_baseline = pd.DataFrame.from_dict(results[model_name]["baseline"], orient = "index").reset_index()
        df_baseline.columns = ["Step", "Loss"]
        df_baseline["Step"] = pd.to_numeric(df_baseline["Step"])
        df_baseline["Test"] = "Baseline"

        df = pd.concat([df, df_baseline])

    return df

with demo:
    with gr.Row():
        title = gr.Markdown(value=f"""# <p style="text-align: center;"> Subnet 38 Model Convergence</p>""")
    with gr.Row():
        dropdown_1 = gr.Dropdown(choices = models, value = models[0])
        button_1 = gr.Button("Submit")
    with gr.Row():
        chart = gr.LinePlot(df, "Step", "Loss", color="Test", x_lim = (0, max(df['Step'])))
    
    button_1.click(return_results, dropdown_1, chart)

demo.launch(debug=True, server_name="0.0.0.0", server_port=7860)