import streamlit as st
import pandas as pd
import numpy as np
import plotly.graph_objects as go
import plotly.express as px

# Set page config
st.set_page_config(
    page_title="GPT-4o Cost Calculator",
    page_icon="💰",
    layout="wide"
)

# Custom CSS
st.markdown("""
    <style>
    .main {
        padding: 2rem;
    }
    .stAlert {
        padding: 1rem;
        margin: 1rem 0;
    }
    </style>
""", unsafe_allow_html=True)

# Define the pricing data
pricing_data = {
    "gpt-4o-audio-preview": {
        "text_input": 2.50,
        "text_output": 10.00,
        "audio_input": 100.00,
        "audio_output": 200.00,
        "description": "Full-featured model with highest quality"
    },
    "gpt-4o-audio-preview-2024-12-17": {
        "text_input": 2.50,
        "text_output": 10.00,
        "audio_input": 40.00,
        "audio_output": 80.00,
        "description": "Updated model with optimized audio pricing"
    },
    "gpt-4o-audio-preview-2024-10-01": {
        "text_input": 2.50,
        "text_output": 10.00,
        "audio_input": 100.00,
        "audio_output": 200.00,
        "description": "Legacy model with standard pricing"
    },
    "gpt-4o-mini-audio-preview": {
        "text_input": 0.150,
        "text_output": 0.600,
        "audio_input": 10.000,
        "audio_output": 20.000,
        "description": "Cost-effective model for lighter workloads"
    },
    "gpt-4o-mini-audio-preview-2024-12-17": {
        "text_input": 0.150,
        "text_output": 0.600,
        "audio_input": 10.000,
        "audio_output": 20.000,
        "description": "Updated mini model with optimized performance"
    }
}

# Constants for calculations
TOKENS_PER_MINUTE_AUDIO = 1000  # Approximate tokens per minute of audio
WORDS_PER_MINUTE = 150  # Average speaking rate
TOKENS_PER_WORD = 1.3  # Approximate tokens per word

def calculate_costs(model, input_type, minutes):
    """Calculate costs based on input parameters"""
    pricing = pricing_data[model]
    
    if input_type == "Audio":
        tokens = minutes * TOKENS_PER_MINUTE_AUDIO
        input_cost = (tokens * pricing["audio_input"]) / 1000000
        output_cost = (tokens * pricing["audio_output"]) / 1000000
    else:  # Text
        words = minutes * WORDS_PER_MINUTE
        tokens = words * TOKENS_PER_WORD
        input_cost = (tokens * pricing["text_input"]) / 1000000
        output_cost = (tokens * pricing["text_output"]) / 1000000
    
    return {
        "tokens": tokens,
        "words": words if input_type == "Text" else None,
        "input_cost": input_cost,
        "output_cost": output_cost,
        "total_cost": input_cost + output_cost
    }

# Header
st.title("GPT-4o Cost Calculator 💰")
st.markdown("Estimate your GPT-4o API costs based on usage")

# Create tabs
tab1, tab2, tab3 = st.tabs(["📊 Pricing Reference", "🧮 Calculator", "📚 Guide"])

# Tab 1: Pricing Reference
with tab1:
    st.header("Original GPT-4o Pricing")
    
    # Create pricing table
    pricing_df = pd.DataFrame([
        {
            "Model": model,
            "Description": data["description"],
            "Text Input": f"${data['text_input']:.3f}",
            "Text Output": f"${data['text_output']:.3f}",
            "Audio Input": f"${data['audio_input']:.3f}",
            "Audio Output": f"${data['audio_output']:.3f}"
        }
        for model, data in pricing_data.items()
    ])
    
    st.dataframe(
        pricing_df,
        hide_index=True,
        use_container_width=True
    )
    
    st.caption("All prices are per 1 million tokens")

# Tab 2: Calculator
with tab2:
    st.header("Cost Calculator")
    
    # Create two columns
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("Input Parameters")
        
        # Model selection
        selected_model = st.selectbox(
            "Select Model",
            options=list(pricing_data.keys()),
            help="Choose the GPT-4o model you want to use"
        )
        
        # Input type selection
        input_type = st.radio(
            "Select Input Type",
            options=["Text", "Audio"],
            help="Choose whether you're processing text or audio",
            horizontal=True
        )
        
        # Duration input
        minutes = st.number_input(
            "Duration (minutes)",
            min_value=0.0,
            value=1.0,
            step=0.5,
            help="Enter the duration of your content in minutes"
        )
        
        # Show relevant examples
        if input_type == "Text":
            st.info(
                f"💡 For {minutes:.1f} minutes of text:\n"
                f"- Approximately {int(minutes * WORDS_PER_MINUTE):,} words\n"
                f"- Based on average speaking rate ({WORDS_PER_MINUTE} words/minute)"
            )
        else:
            st.info(
                f"💡 For {minutes:.1f} minutes of audio:\n"
                f"- Approximately {int(minutes * TOKENS_PER_MINUTE_AUDIO):,} tokens\n"
                f"- Based on audio processing requirements"
            )
    
    # Calculate costs
    costs = calculate_costs(selected_model, input_type, minutes)
    
    with col2:
        st.subheader("Cost Breakdown")
        
        # Create metrics
        col_a, col_b = st.columns(2)
        with col_a:
            st.metric(
                "Processing Cost",
                f"${costs['input_cost']:.2f}",
                f"{costs['tokens']:,.0f} tokens"
            )
        with col_b:
            st.metric(
                "Response Cost",
                f"${costs['output_cost']:.2f}",
                f"{costs['tokens']:,.0f} tokens"
            )
        
        # Total cost
        st.metric(
            "Total Estimated Cost",
            f"${costs['total_cost']:.2f}",
            f"For {minutes} minute{'s' if minutes != 1 else ''}"
        )
        
        # Create a pie chart for cost distribution
        fig = px.pie(
            values=[costs['input_cost'], costs['output_cost']],
            names=['Input Processing', 'Output Processing'],
            title='Cost Distribution'
        )
        st.plotly_chart(fig, use_container_width=True)

# Tab 3: Guide
with tab3:
    st.header("Understanding Tokens and Costs")
    
    # What are tokens?
    st.subheader("What are tokens?")
    st.markdown("""
        Tokens are the basic units that GPT-4o processes:
        - For text: ~4 characters or ¾ of a word
        - For audio: ~1 second of speech
    """)
    
    # Examples
    col_a, col_b = st.columns(2)
    
    with col_a:
        st.subheader("Text Examples")
        st.markdown("""
            | Content | Words | Tokens |
            |---------|-------|---------|
            | Short message | 20 | ~26 |
            | Email | 200 | ~260 |
            | Document page | 500 | ~650 |
        """)
    
    with col_b:
        st.subheader("Audio Examples")
        st.markdown("""
            | Content | Duration | Tokens |
            |---------|----------|---------|
            | Short clip | 1 min | ~1,000 |
            | Conversation | 5 min | ~5,000 |
            | Meeting | 30 min | ~30,000 |
        """)
    
    # Cost optimization tips
    st.subheader("Cost Optimization Tips")
    st.markdown("""
        1. **Choose the right model:**
           - Use mini models for development
           - Use full models for production
        
        2. **Optimize content length:**
           - Keep text concise
           - Trim unnecessary audio
        
        3. **Batch processing:**
           - Combine related requests
           - Process in optimal chunks
    """)

# Footer
st.markdown("---")
st.caption("""
    Note: These calculations are estimates based on average usage patterns. 
    Actual token usage and costs may vary depending on the specific content and use case.
    Prices are based on the official OpenAI GPT-4o pricing.
""")