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_all_objectives_progress

def render_objective_progress(budget_hours=2, is_running=False, selected_objectives=None, current_step=None):
    """Render progress charts for each selected optimization objective."""

    if selected_objectives is None:
        selected_objectives = {}

    # Get all objectives progress data
    df = generate_all_objectives_progress(budget_hours)
    total_rows = len(df)

    # If optimization is "running", only show partial data based on current progress
    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)
        # Show data up to current step (1-20)
        display_rows = max(1, min(current_step, total_rows))
        df = df.iloc[:display_rows]

    # Define objective metadata
    objectives_meta = {
        'accuracy': {
            'label': 'Accuracy (%)',
            'color': '#636EFA',
            'maximize': True,
            'format': '.1f',
            'suffix': '%'
        },
        'size': {
            'label': 'Model Size (GB)',
            'color': '#EF553B',
            'maximize': False,
            'format': '.1f',
            'suffix': ' GB'
        },
        'cost': {
            'label': 'Inference Cost ($/1M)',
            'color': '#00CC96',
            'maximize': False,
            'format': '.2f',
            'suffix': ''
        },
        'throughput': {
            'label': 'Throughput (QPS)',
            'color': '#AB63FA',
            'maximize': True,
            'format': '.0f',
            'suffix': ''
        },
        'latency': {
            'label': 'Latency (ms)',
            'color': '#FFA15A',
            'maximize': False,
            'format': '.1f',
            'suffix': ' ms'
        },
        'memory': {
            'label': 'Memory Footprint (GB)',
            'color': '#19D3F3',
            'maximize': False,
            'format': '.1f',
            'suffix': ' GB'
        },
        'energy': {
            'label': 'Energy (W)',
            'color': '#FF6692',
            'maximize': False,
            'format': '.0f',
            'suffix': ' W'
        }
    }

    # Filter to only show selected objectives
    objectives_to_show = [obj for obj in objectives_meta.keys() if selected_objectives.get(obj, False)]

    if not objectives_to_show:
        return

    st.markdown("#### 📈 Objective Progress")
    st.caption("Real-time tracking of improvements for each selected optimization objective • Watch as AI Forge discovers better configurations")

    # Calculate number of rows needed (3 charts per row)
    num_objectives = len(objectives_to_show)
    num_rows = (num_objectives + 2) // 3  # Ceiling division

    for row_idx in range(num_rows):
        cols = st.columns(3)

        for col_idx in range(3):
            obj_idx = row_idx * 3 + col_idx

            if obj_idx < num_objectives:
                obj = objectives_to_show[obj_idx]
                meta = objectives_meta[obj]

                with cols[col_idx]:
                    # Create individual chart for this objective
                    fig = go.Figure()

                    fig.add_trace(
                        go.Scatter(
                            x=df['time'],
                            y=df[obj],
                            mode='lines+markers',
                            line=dict(color=meta['color'], width=2.5),
                            marker=dict(size=6, symbol='circle'),
                            fill='tozeroy',
                            fillcolor=f'rgba({int(meta["color"][1:3], 16)}, {int(meta["color"][3:5], 16)}, {int(meta["color"][5:7], 16)}, 0.1)',
                            hovertemplate=f'<b>Time:</b> %{{x:.2f}}h<br>' +
                                        f'<b>{meta["label"]}:</b> %{{y:{meta["format"]}}}{meta["suffix"]}<br>' +
                                        '<extra></extra>'
                        )
                    )

                    # Update layout for compact display
                    fig.update_layout(
                        plot_bgcolor='rgba(15, 23, 42, 0.5)',
                        paper_bgcolor='rgba(0, 0, 0, 0)',
                        height=280,
                        margin=dict(l=55, r=25, t=50, b=50),
                        title=dict(
                            text=meta['label'],
                            font=dict(size=18, color=meta['color'], family='sans-serif'),
                            x=0.5,
                            xanchor='center'
                        ),
                        xaxis=dict(
                            title='Time (h)',
                            title_font=dict(size=16, family='sans-serif'),
                            tickfont=dict(size=14, family='sans-serif'),
                            gridcolor='rgba(255, 255, 255, 0.1)',
                            showgrid=True,
                            zeroline=False,
                            color='rgba(255,255,255,0.8)',
                            range=[0, budget_hours]
                        ),
                        yaxis=dict(
                            title='',
                            tickfont=dict(size=14, family='sans-serif'),
                            gridcolor='rgba(255, 255, 255, 0.1)',
                            showgrid=True,
                            zeroline=False,
                            color='rgba(255,255,255,0.8)'
                        ),
                        font=dict(size=14, color='rgba(255,255,255,0.9)', family='sans-serif'),
                        hoverlabel=dict(
                            bgcolor='rgba(30, 41, 59, 0.95)',
                            font_size=16,
                            font_family='sans-serif'
                        )
                    )

                    st.plotly_chart(fig, use_container_width=True, key=f"progress_{obj}_{row_idx}_{col_idx}")