import gradio as gr
import pandas as pd
import datetime
import plotly.express as px

def split_multi_users(dfs):
    df = dfs.copy()
    df["usernames"] = df["username"].apply(lambda x: x.split(", "))
    df["count"] = 1
    new_df = []
    for row in df.to_dict(orient="records"):
        gpu_users_num = len(row["usernames"])
        for username in row["usernames"]:
            new_row = row.copy()
            new_row["count"] = 1 / gpu_users_num
            new_row["username"] = username
            new_df.append(new_row)
    df = pd.DataFrame(new_df)
    return df

def plot_now():
    dfs = pd.read_csv("hf://spaces/pluslab/PLUS_Lab_GPUs/gpus.csv")
    dfs = dfs.drop(columns=["Unnamed: 0"])
    dfs = dfs.fillna("FREE")
    dfs_plot = split_multi_users(dfs)
    fig = px.bar(
        dfs_plot, x="count", y="server", color="username",
        title=f"Last Updated {min(dfs['timestamp'])}",
        color_discrete_map={
            "FREE": "black",
        },
        text=dfs_plot['username'].astype(str) + "<br>" + dfs_plot['device'].astype(str),
    )
    fig.update_layout(
        yaxis={'categoryorder': 'array', 'categoryarray': dfs_plot["server"].unique()[::-1]},
        barcornerradius=50,
    )
    fig.update_traces(textposition='inside', insidetextanchor='middle')
    print(dfs_plot)
    return fig, dfs

def plot_history(sample=True, sampling_interval_minutes=180):
    dfh = pd.read_pickle("hf://spaces/pluslab/PLUS_Lab_GPUs/history.pkl.gz", )
    dfh = dfh.fillna("FREE")
    dfh = split_multi_users(dfh)
    dfh = dfh[["polling_timestamp", "username", "count"]]
    dfh = dfh.groupby(["polling_timestamp", "username"]).sum()
    dfh = dfh.reset_index()
    dfh = dfh.sort_values(by=["polling_timestamp", "count"], ascending=False)
    if sample:
        unique_timestamps = dfh["polling_timestamp"].unique()
        sampled_timestamps = [unique_timestamps[0]]
        for i, t in enumerate(unique_timestamps[1:]):
            diff = sampled_timestamps[-1] - t
            if diff > datetime.timedelta(minutes=sampling_interval_minutes):
                sampled_timestamps.append(t)
        dfh = dfh[dfh["polling_timestamp"].isin(sampled_timestamps)]
    fig = px.area(dfh, x="polling_timestamp", y="count", color='username', color_discrete_map={"FREE": "black",}, markers=True, line_shape='spline',)
    return fig, dfh


def plot_figs():
    fig_now, dfn = plot_now()
    try:
        fig_history, dfh = plot_history()
    except Exception as e:
        print(e)
        fig_history = None
        dfh = None
    return fig_now, dfn, fig_history

demo = gr.Interface(
    fn=plot_figs,
    inputs = [

    ],
    outputs = [
        gr.Plot(label="GPU Status", elem_classes="plotcss"),
        gr.Dataframe(label="GPU Status Details"),
        gr.Plot(label="History", elem_classes="plotcss"),
    ],
    live=True,
    flagging_options=[],
    css=".plotcss {max-width: 820px !important;}"
)

if __name__ == "__main__":
    demo.launch(debug=False)