import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

st.set_page_config(page_title="LLM API Budget Dashboard", layout="wide")

# Title and description
st.title("LLM API Budget Dashboard")
st.markdown("This dashboard helps you budget your API calls to various LLMs based on input and output tokens.")

# Define LLM models and their costs
llm_data = {
    "GPT-4o": {"input_cost_per_m": 2.50, "output_cost_per_m": 10.00},
    "Claude 3.7 Sonnet": {"input_cost_per_m": 3.00, "output_cost_per_m": 15.00},
    "Gemini Flash 1.5-8b": {"input_cost_per_m": 0.038, "output_cost_per_m": 0.15},
    "o3-mini": {"input_cost_per_m": 1.10, "output_cost_per_m": 4.40}
}

# Convert the LLM data to a DataFrame for displaying in a table
llm_df = pd.DataFrame([
    {
        "Model": model, 
        "Input Cost ($/M tokens)": data["input_cost_per_m"], 
        "Output Cost ($/M tokens)": data["output_cost_per_m"]
    } 
    for model, data in llm_data.items()
])

# Display LLM cost info
st.subheader("LLM Cost Information")
st.dataframe(llm_df, use_container_width=True)

# Create sidebar for inputs
st.sidebar.header("Configuration")

# Token input section
st.sidebar.subheader("Token Settings")
input_tokens = st.sidebar.number_input("Input Tokens", min_value=1, value=1000, step=100)
output_tokens = st.sidebar.number_input("Output Tokens", min_value=1, value=500, step=100)

# LLM selection
st.sidebar.subheader("Select LLMs")
selected_llms = st.sidebar.multiselect("Choose LLMs", options=list(llm_data.keys()), default=list(llm_data.keys()))

# Run count settings
st.sidebar.subheader("Run Count Settings")
uniform_runs = st.sidebar.checkbox("Run all LLMs the same number of times", value=True)

if uniform_runs:
    uniform_run_count = st.sidebar.number_input("Number of runs for all LLMs", min_value=1, value=1, step=1)
    run_counts = {llm: uniform_run_count for llm in selected_llms}
else:
    st.sidebar.write("Set individual run counts for each LLM:")
    run_counts = {}
    for llm in selected_llms:
        run_counts[llm] = st.sidebar.number_input(f"Runs for {llm}", min_value=1, value=1, step=1)

# Stability test settings
st.sidebar.subheader("Stability Test Settings")
stability_test = st.sidebar.checkbox("Enable stability testing", value=False)

if stability_test:
    st.sidebar.write("Set stability iterations for selected LLMs:")
    stability_iterations = {}
    for llm in selected_llms:
        stability_enabled = st.sidebar.checkbox(f"Test stability for {llm}", value=False)
        if stability_enabled:
            iterations = st.sidebar.number_input(f"Iterations for {llm}", min_value=2, value=10, step=1)
            stability_iterations[llm] = iterations
else:
    stability_iterations = {}

# Calculate costs
results = []

for llm in selected_llms:
    base_runs = run_counts[llm]
    stability_runs = stability_iterations.get(llm, 0)
    total_runs = base_runs * (1 if stability_runs == 0 else stability_runs)
    
    total_input_tokens = input_tokens * total_runs
    total_output_tokens = output_tokens * total_runs
    
    input_cost = (total_input_tokens / 1_000_000) * llm_data[llm]["input_cost_per_m"]
    output_cost = (total_output_tokens / 1_000_000) * llm_data[llm]["output_cost_per_m"]
    total_cost = input_cost + output_cost
    
    results.append({
        "Model": llm,
        "Base Runs": base_runs,
        "Stability Test Iterations": stability_iterations.get(llm, 0),
        "Total Runs": total_runs,
        "Total Input Tokens": total_input_tokens,
        "Total Output Tokens": total_output_tokens,
        "Input Cost ($)": input_cost,
        "Output Cost ($)": output_cost,
        "Total Cost ($)": total_cost
    })

# Create DataFrame from results
results_df = pd.DataFrame(results)

# Main content
st.header("Cost Summary")
st.dataframe(results_df, use_container_width=True)

# Calculate overall totals
total_input_cost = results_df["Input Cost ($)"].sum()
total_output_cost = results_df["Output Cost ($)"].sum()
total_cost = results_df["Total Cost ($)"].sum()

# Display totals
col1, col2, col3 = st.columns(3)
col1.metric("Total Input Cost", f"${total_input_cost:.2f}")
col2.metric("Total Output Cost", f"${total_output_cost:.2f}")
col3.metric("Total API Cost", f"${total_cost:.2f}")

# Data visualization
st.header("Cost Visualization")

# Cost breakdown by model
fig1, ax1 = plt.subplots(figsize=(10, 6))
models = results_df["Model"]
input_costs = results_df["Input Cost ($)"]
output_costs = results_df["Output Cost ($)"]

x = np.arange(len(models))
width = 0.35

ax1.bar(x - width/2, input_costs, width, label='Input Cost')
ax1.bar(x + width/2, output_costs, width, label='Output Cost')

ax1.set_ylabel('Cost ($)')
ax1.set_title('Cost Breakdown by Model')
ax1.set_xticks(x)
ax1.set_xticklabels(models, rotation=45, ha='right')
ax1.legend()

fig1.tight_layout()
st.pyplot(fig1)

# Percentage of total cost by model
fig2, ax2 = plt.subplots(figsize=(8, 8))
ax2.pie(results_df["Total Cost ($)"], labels=results_df["Model"], autopct='%1.1f%%', startangle=90)
ax2.axis('equal')
ax2.set_title('Percentage of Total Cost by Model')
st.pyplot(fig2)

# Export options
st.header("Export Options")
csv = results_df.to_csv(index=False).encode('utf-8')
st.download_button(
    label="Download Results as CSV",
    data=csv,
    file_name='llm_budget_results.csv',
    mime='text/csv',
)

# Footer
st.markdown("---")
st.markdown("*Note: All costs are estimates based on the provided rates. Actual API costs may vary.*")