import streamlit as st
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import time
from mock_data.sample_data import generate_optimization_progress

def render_progress_chart(budget_hours=2, is_running=False, enable_animation=True, current_step=None):
    """Render the optimization progress chart with dual axes.

    Args:
        budget_hours: Estimated hours for optimization
        is_running: Whether optimization is currently running
        enable_animation: If True, uses 10-second animation with sleep (for local demos)
                         If False, instant updates (for Hugging Face Spaces)
        current_step: Current step in the optimization (1-20), None to use session state
    """

    # Get mock data based on budget hours
    df = generate_optimization_progress(budget_hours)
    total_rows = len(df)  # Store original length before slicing

    # If optimization is "running", only show partial data
    if is_running:
        # Use provided current_step or get from session state
        if current_step is None:
            current_step = st.session_state.get('progress_step', 1)
        # Always show at least 1 row so chart is visible from the start
        display_rows = max(1, min(current_step, total_rows))
        df = df.iloc[:display_rows]
    else:
        # Show all data when not running
        if 'progress_step' in st.session_state:
            del st.session_state.progress_step
            if 'last_update' in st.session_state:
                del st.session_state.last_update

    # Create figure with secondary y-axis
    fig = make_subplots(specs=[[{"secondary_y": True}]])

    # Add accuracy trace with improved styling
    fig.add_trace(
        go.Scatter(
            x=df['search_time'],
            y=df['accuracy'],
            name='Accuracy',
            mode='lines+markers',
            line=dict(color='#636EFA', width=3),
            marker=dict(size=8, symbol='circle'),
            fill='tozeroy',
            fillcolor='rgba(99, 110, 250, 0.1)',
            hovertemplate='<b>Time:</b> %{x:.2f}h<br>' +
                          '<b>Accuracy:</b> %{y:.1f}%<br>' +
                          '<extra></extra>'
        ),
        secondary_y=False
    )

    # Add model size trace with improved styling
    fig.add_trace(
        go.Scatter(
            x=df['search_time'],
            y=df['model_size_gb'],
            name='Model Size',
            mode='lines+markers',
            line=dict(color='#EF553B', width=3),
            marker=dict(size=8, symbol='diamond'),
            hovertemplate='<b>Time:</b> %{x:.2f}h<br>' +
                          '<b>Size:</b> %{y:.1f} GB<br>' +
                          '<extra></extra>'
        ),
        secondary_y=True
    )

    # Update axes with dark theme styling
    fig.update_xaxes(
        title_text="<b>Optimization Time (hours)</b>",
        title_font=dict(size=18, family='sans-serif'),
        tickfont=dict(size=16, family='sans-serif'),
        range=[0, budget_hours],
        gridcolor='rgba(255, 255, 255, 0.1)',
        showgrid=True,
        zeroline=False,
        color='rgba(255,255,255,0.9)'
    )

    fig.update_yaxes(
        title_text="<b>Accuracy (%)</b>",
        secondary_y=False,
        range=[70, 92],
        title_font=dict(color='#636EFA', size=18, family='sans-serif'),
        tickfont=dict(size=16, family='sans-serif'),
        gridcolor='rgba(255, 255, 255, 0.1)',
        showgrid=True,
        zeroline=False,
        color='rgba(255,255,255,0.9)'
    )

    fig.update_yaxes(
        title_text="<b>Model Size (GB)</b>",
        secondary_y=True,
        range=[0, 40],
        title_font=dict(color='#EF553B', size=18, family='sans-serif'),
        tickfont=dict(size=16, family='sans-serif'),
        gridcolor='rgba(255, 255, 255, 0.05)',
        showgrid=False,
        zeroline=False,
        color='rgba(255,255,255,0.9)'
    )

    # Update layout with premium dark theme styling
    fig.update_layout(
        plot_bgcolor='rgba(15, 23, 42, 0.5)',
        paper_bgcolor='rgba(0, 0, 0, 0)',
        height=550,
        margin=dict(l=80, r=90, t=50, b=70),
        hovermode='x unified',
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="right",
            x=1,
            bgcolor='rgba(30, 41, 59, 0.8)',
            bordercolor='rgba(99, 102, 241, 0.3)',
            borderwidth=1,
            font=dict(color='rgba(255,255,255,0.9)', size=15, family='sans-serif')
        ),
        font=dict(size=16, color='rgba(255,255,255,0.9)', family='sans-serif'),
        hoverlabel=dict(
            bgcolor='rgba(30, 41, 59, 0.95)',
            font_size=18,
            font_family='sans-serif'
        )
    )

    # Display the chart with premium container
    st.markdown("""
    <div style="background: rgba(255, 255, 255, 0.02); border-radius: 15px; padding: 1.5rem;
                border: 1px solid rgba(99, 102, 241, 0.2); box-shadow: 0 8px 32px rgba(0, 0, 0, 0.2);">
    """, unsafe_allow_html=True)

    st.plotly_chart(fig, use_container_width=True, key="progress_chart")

    st.markdown("</div>", unsafe_allow_html=True)