import streamlit as st
import pandas as pd
import plotly.express as px

# Page config
st.set_page_config(
    page_title="Max Output Tokens Analysis",
    layout="wide"
)

# Load data
@st.cache_data
def load_data():
    df = pd.read_csv('data/max-tokens-by-model.csv')
    df['launch_date'] = pd.to_datetime(df['launch_date'])
    return df

df = load_data()

# Title
st.title("LLM Max Output Tokens Analysis")

# Company selection
companies = sorted(df['company'].unique())
selected_companies = st.multiselect(
    "Select companies to display:",
    options=companies,
    default=companies,
    key='company_filter'
)

# Filter data based on selection
filtered_df = df[df['company'].isin(selected_companies)]

# Create the evolution chart
fig = px.line(filtered_df, 
              x='launch_date', 
              y='max_output_tokens',
              color='company',
              hover_data=['model_name', 'max_output_tokens'],
              title='Evolution of Max Output Tokens by Company',
              labels={
                  'launch_date': 'Launch Date',
                  'max_output_tokens': 'Max Output Tokens',
                  'company': 'Company'
              },
              markers=True)  # Add markers to make trends clearer

fig.update_layout(
    hovermode='x unified',
    xaxis_title="Launch Date",
    yaxis_title="Max Output Tokens",
    yaxis_type="log",  # Using log scale for better visualization
    height=600,  # Make chart taller
    showlegend=True,
    legend=dict(
        yanchor="top",
        y=0.99,
        xanchor="left",
        x=0.01
    ),
    margin=dict(l=20, r=20, t=40, b=20)
)

fig.update_traces(
    line=dict(width=2),  # Make lines thicker
    marker=dict(size=8)  # Make markers more visible
)

# Display the chart
st.plotly_chart(fig, use_container_width=True)

# Display the data table
st.subheader("Max Output Tokens by Model")

# Prepare the data with better formatting
display_df = (
    filtered_df[['model_name', 'company', 'max_output_tokens', 'launch_date']]
    .sort_values('max_output_tokens', ascending=False)
    .assign(
        launch_date=lambda x: x['launch_date'].dt.strftime('%Y-%m-%d'),
        max_output_tokens=lambda x: x['max_output_tokens'].apply(lambda v: f"{v:,}")
    )
    .rename(columns={
        'model_name': 'Model Name',
        'company': 'Company',
        'max_output_tokens': 'Max Output Tokens',
        'launch_date': 'Launch Date'
    })
)

# Display the styled table
st.dataframe(
    display_df,
    use_container_width=True,
    hide_index=True
)

# Attribution
st.markdown("---")
st.markdown(
    "By: [Daniel Rosehill](https://danielrosehill.com) | "
    "Data sourced from public sources on February 8, 2025"
)