## LIBRARIES ###
from cProfile import label
from tkinter import font
from turtle import width
import streamlit as st
import pandas as pd
from datetime import datetime
import plotly.express as px


def read_file_to_df(file):
   return pd.read_csv(file)

def date_range(df):
    time = df.index.to_list()
    time_range = []
    for t in time:
        time_range.append(str(datetime.strptime(t, '%Y-%m-%dT%H:%M:%S.%fZ').date().month) +'/' +  str(datetime.strptime(t, '%Y-%m-%dT%H:%M:%S.%fZ').date().day))
    return time_range


if __name__ == "__main__":
    ### STREAMLIT APP CONGFIG ###
    st.set_page_config(layout="wide", page_title="HF Hub Model Usage Visualization")

    popularity = st.sidebar.radio(
        "Model popularity", ('Low', 'Moderate', 'High'), key = "popularity", index=2)
    st.header("Model Usage Visualization")

    with st.container():
        df_2021 = read_file_to_df("./assets/2021/model_init_time.csv")
        df_2021.fillna(0, inplace=True)
        df_plot = df_2021.set_index('Model').T
        df_plot.index = date_range(df_plot)
        df_plot_2021 = pd.DataFrame()
        if popularity == 'Low':
            df_plot_2021 = df_plot[df_plot.columns[(df_plot.mean(axis=0)<=5000) & (df_plot.mean(axis=0)>=3500)]]
        elif popularity == 'Moderate':
            df_plot_2021 =  df_plot[df_plot.columns[(df_plot.mean(axis=0)<=40000) & (df_plot.mean(axis=0)>=5000)]]
        else:
            df_plot_2021 = df_plot[df_plot.columns[df_plot.mean(axis=0)>=40000]]
        fig = px.line(df_plot_2021, title="Model Usage Trends in 2021", labels={"index": "Weeks", "value": "Usage", "variable": "Model"})
        st.plotly_chart(fig, use_container_width=True)

    with st.container():
        df_2022 = read_file_to_df("./assets/2022/model_init_time.csv")
        df_2022.fillna(0, inplace=True)
        df_plot = df_2022.set_index('Model').T
        df_plot.index = date_range(df_plot)
        df_plot_2022 = pd.DataFrame()
        if popularity == 'Low':
            df_plot_2022 = df_plot[df_plot.columns[(df_plot.mean(axis=0)<500) & (df_plot.mean(axis=0)>=300)]]
        elif popularity == 'Moderate':
            df_plot_2022 =  df_plot[df_plot.columns[(df_plot.mean(axis=0)<=1500) & (df_plot.mean(axis=0)>=500)]]
        else:
            df_plot_2022 = df_plot[df_plot.columns[df_plot.mean(axis=0)>=1500]]
        fig = px.line(df_plot_2022, title="Model Usage Trends in 2022", labels={"index": "Weeks", "value": "Usage", "variable": "Model"})
        st.plotly_chart(fig, use_container_width=True)