# Import packages
from dash import Dash, html, dcc, callback, Input, Output
import pandas as pd
import pickle
import plotly.express as px
from graphs.model_market_share import create_plotly_stacked_area_chart
from graphs.model_characteristics import create_plotly_language_concentration_chart

# Incorporate data
df = pd.read_csv('https://raw.githubusercontent.com/plotly/datasets/master/gapminder2007.csv')

# Initialize the app
app = Dash()

# Load all pickle files in data_frames/ as loop
with open('data_frames/model_topk_df.pkl', 'rb') as f:
    model_topk_df = pickle.load(f)
with open('data_frames/model_gini_df.pkl', 'rb') as f:
    model_gini_df = pickle.load(f)
with open('data_frames/model_hhi_df.pkl', 'rb') as f:
    model_hhi_df = pickle.load(f)
with open('data_frames/language_concentration_df.pkl', 'rb') as f:
    language_concentration_df = pickle.load(f)
with open('data_frames/download_license_cumsum_df.pkl', 'rb') as f:
    license_concentration_df = pickle.load(f)

TEMP_MODEL_EVENTS = {
    # "Yolo World Mirror": "2024-03-01",
    "Llama 3": "2024-04-17",
    "Stable Cascade": "2024-02-02",
    "Stable Diffusion 3": "2024-05-30",
    # "embed/upscale": "2023-03-24",
    "DeepSeek-R1": "2025-01-20",
    "Gemma-3 12B QAT": "2025-04-15", # gemma-3-12b-it-qat-4bit
    # "Qwen": "2025-03-05",
    # "Flux RedFlux": "2025-04-12",
    # "DeepSeek-V3": "2025-03-24",
    # "bloom": "2022-05-19",
    "DALLE2-PyTorch": "2022-06-25",
    "Stable Diffusion": "2022-08-10",
    "CLIP ViT": "2021-01-05",
    "YOLOv8": "2023-04-26",
    "Sentence Transformer MiniLM v2": "2021-08-30",
}

PALETTE_0 = [
    "#335C67",
    "#FFF3B0",
    "#E09F3E",
    "#9E2A2B",
    "#540B0E"
]

fig = create_plotly_stacked_area_chart(
    model_topk_df, model_gini_df, model_hhi_df, TEMP_MODEL_EVENTS, PALETTE_0
)

LANG_SEGMENT_ORDER = [
    'Monolingual: EN', 'Monolingual: HR', 'Monolingual: M/LR', 
    'Multilingual: HR', 'Multilingual', 'Unknown',
]
fig2 = create_plotly_language_concentration_chart(
    language_concentration_df, 'time', 'metric', 'value', LANG_SEGMENT_ORDER, PALETTE_0
)

LICENSE_SEGMENT_ORDER = [
    "Open Use", "Open Use (Acceptable Use Policy)", "Open Use (Non-Commercial Only)", "Attribution", 
    "Acceptable Use Policy", "Non-Commercial Only", "Undocumented", "Undocumented (Acceptable Use Policy)",
]
fig3 = create_plotly_language_concentration_chart(
    license_concentration_df, 'period', 'status', 'percent', LICENSE_SEGMENT_ORDER, PALETTE_0
)

# Make global font family
fig.update_layout(font_family="Inter")
fig2.update_layout(font_family="Inter")
fig3.update_layout(font_family="Inter")

# App layout
app.layout = html.Div(
    [
        html.Div(children='Visualizing the Open Model Ecosystem', style={'fontSize': 28, 'fontWeight': 'bold', 'marginBottom': 10}),
        html.Div(children='An interactive dashboard to explore trends in open models on Hugging Face', style={'fontSize': 16, 'marginBottom': 20}),
        html.Hr(),
        dcc.Tabs([
            dcc.Tab(label='Model Market Share', children=[
                dcc.Graph(figure=fig, id='stacked-area-chart'),
            ]),
            dcc.Tab(label='Model Characteristics', children=[
                dcc.Graph(id='language-concentration-chart'),
                html.Div([
                    dcc.Dropdown(['Language Concentration', 'Architecture', 'License', 'Method'], 'Language Concentration', id='dropdown'),
                ]),
            ]),
        ])
    ],
    style={'fontFamily': 'Inter'}
)

# On dropdown change, update graph
@app.callback(
    Output('language-concentration-chart', 'figure'),
    [Input('dropdown', 'value')]
)
def update_graph(selected_metric):
    if selected_metric == 'Language Concentration':
        return fig2
    elif selected_metric == 'License':
        return fig3

# Run the app
if __name__ == '__main__':
    app.run(debug=True)