import streamlit as st
import numpy as np
import pandas as pd
import plotly.graph_objects as go
from datetime import datetime
import io

st.set_page_config(page_title="EER Simulator", layout="wide")

# --- Password Gate ---
if "authenticated" not in st.session_state:
    st.session_state.authenticated = False

if not st.session_state.authenticated:
    st.markdown("<h1 style='font-size: 2rem;'>🔒 EER Simulator Login</h1>", unsafe_allow_html=True)
    password = st.text_input("Enter password to access the app:", type="password")
    if st.button("Submit"):
        if password == st.secrets["password"]:
            st.session_state.authenticated = True
            st.rerun()
        else:
            st.error("Incorrect password. Please try again.")
    st.stop()

# Initialize session state
if 'tiers' not in st.session_state:
    st.session_state.tiers = [
        {'id': 1, 'points_required': 500, 'reward_value': 5.0, 'redemption_pct': 50.0, 'uncertainty': 5.0},
        {'id': 2, 'points_required': 2500, 'reward_value': 25.0, 'redemption_pct': 30.0, 'uncertainty': 5.0},
        {'id': 3, 'points_required': 10000, 'reward_value': 100.0, 'redemption_pct': 20.0, 'uncertainty': 5.0}
    ]
    st.session_state.next_id = 4

if 'saved_scenarios' not in st.session_state:
    st.session_state.saved_scenarios = []

if 'simulation_results' not in st.session_state:
    st.session_state.simulation_results = None

if 'scenarios_to_compare' not in st.session_state:
    st.session_state.scenarios_to_compare = []

if 'show_uniform_confirm' not in st.session_state:
    st.session_state.show_uniform_confirm = False

if 'loaded_points_per_dollar' not in st.session_state:
    st.session_state.loaded_points_per_dollar = None

if 'loaded_num_simulations' not in st.session_state:
    st.session_state.loaded_num_simulations = None

def calculate_tier_eer(reward_value, points_required, points_per_dollar):
    """Calculate EER for a single tier"""
    if points_required == 0 or points_per_dollar == 0:
        return 0
    return (reward_value / (points_required / points_per_dollar)) * 100

def get_guidance(mean_eer):
    """Get guidance based on mean EER"""
    if mean_eer >= 10:
        return {
            'category': 'High Generosity',
            'color': '#f59e0b',
            'guidance': """<p style="margin: 0 0 0.75rem 0; font-size: 0.85rem; font-weight: 600;">Key Loyalty Design Considerations:</p>
<ul style="margin: 0; padding-left: 1.5rem; font-size: 0.85rem; line-height: 1.6;">
    <li style="margin-bottom: 0.5rem;"><strong>Tier Thresholds:</strong> Make tiers harder to unlock to balance the generous earn rate and control program costs. Higher barriers ensure customers earn points more slowly relative to reward access.</li>
    <li style="margin-bottom: 0.5rem;"><strong>Reward Stacking:</strong> Limit to 1-2 offers maximum to prevent excessive point accumulation that could strain program economics.</li>
    <li style="margin-bottom: 0;"><strong>Point Expiry:</strong> Use 6-month expiry to encourage faster redemption and reduce long-term liability while maintaining urgency.</li>
</ul>"""
        }
    elif mean_eer >= 8:
        return {
            'category': 'Moderate',
            'color': '#3b82f6',
            'guidance': """<p style="margin: 0 0 0.75rem 0; font-size: 0.85rem; font-weight: 600;">Key Loyalty Design Considerations:</p>
<ul style="margin: 0; padding-left: 1.5rem; font-size: 0.85rem; line-height: 1.6;">
    <li style="margin-bottom: 0.5rem;"><strong>Tier Thresholds:</strong> Balance accessibility and aspiration by making lower tiers easier to unlock (driving early engagement) while keeping top tiers harder (maintaining exclusivity and managing costs).</li>
    <li style="margin-bottom: 0.5rem;"><strong>Reward Stacking:</strong> Allow 3-5 concurrent offers to provide flexibility and boost engagement without overwhelming program costs.</li>
    <li style="margin-bottom: 0;"><strong>Point Expiry:</strong> Use 6-12 month expiry to give customers reasonable redemption windows while managing liability.</li>
</ul>"""
        }
    else:
        return {
            'category': 'Conservative',
            'color': '#10b981',
            'guidance': """<p style="margin: 0 0 0.75rem 0; font-size: 0.85rem; font-weight: 600;">Key Loyalty Design Considerations:</p>
<ul style="margin: 0; padding-left: 1.5rem; font-size: 0.85rem; line-height: 1.6;">
    <li style="margin-bottom: 0.5rem;"><strong>Tier Thresholds:</strong> Make tiers very easy to unlock, especially the first tier, to drive immediate engagement and participation. Lower barriers compensate for conservative earn rates.</li>
    <li style="margin-bottom: 0.5rem;"><strong>Reward Stacking:</strong> Allow generous stacking (6+ offers) to accelerate point earning and maintain member motivation despite lower base earn rates.</li>
    <li style="margin-bottom: 0;"><strong>Point Expiry:</strong> Use 12-month expiry to give ample time for point accumulation and redemption, reducing frustration from conservative earning.</li>
</ul>"""
        }

def add_tier():
    """Add a new tier"""
    st.session_state.tiers.append({
        'id': st.session_state.next_id,
        'points_required': 1000,
        'reward_value': 10.0,
        'redemption_pct': 10.0,
        'uncertainty': 5.0
    })
    st.session_state.next_id += 1

def remove_tier(tier_id):
    """Remove a specific tier"""
    if len(st.session_state.tiers) > 1:
        st.session_state.tiers = [t for t in st.session_state.tiers if t['id'] != tier_id]

def apply_uniform_redemption():
    """Apply uniform redemption percentage across all tiers"""
    if not st.session_state.tiers:
        return
    
    num_tiers = len(st.session_state.tiers)
    uniform_pct = 100.0 / num_tiers
    
    for tier in st.session_state.tiers:
        tier['redemption_pct'] = round(uniform_pct, 2)
    
    st.session_state.show_uniform_confirm = False

def check_redemption_total():
    """Check if redemption percentages sum to 100"""
    total = sum(tier['redemption_pct'] for tier in st.session_state.tiers)
    return total

def load_scenario(scenario):
    """Load a saved scenario into the form"""
    # Clear existing tiers
    st.session_state.tiers = []
    st.session_state.next_id = 1
    
    # Load points per dollar
    st.session_state.loaded_points_per_dollar = scenario['results']['points_per_dollar']
    st.session_state.loaded_num_simulations = scenario['results']['num_iterations']
    
    # Load tiers from scenario
    for tier in scenario['results']['tiers']:
        st.session_state.tiers.append({
            'id': st.session_state.next_id,
            'points_required': int(tier['points_required']),
            'reward_value': float(tier['reward_value']),
            'redemption_pct': float(tier['redemption_pct']),
            'uncertainty': float(tier['uncertainty'])
        })
        st.session_state.next_id += 1

def run_simulation(points_per_dollar, num_iterations):
    """Run Monte Carlo simulation"""
    try:
        if not st.session_state.tiers:
            return None, "Please add at least one tier."
        
        # Run simulation
        results = []
        
        for i in range(int(num_iterations)):
            weighted_eer = 0
            
            # Sample redemption percentages with uncertainty
            sampled_weights = []
            for tier in st.session_state.tiers:
                min_val = max(0, tier['redemption_pct'] - tier['uncertainty'])
                max_val = min(100, tier['redemption_pct'] + tier['uncertainty'])
                sampled_weights.append(np.random.uniform(min_val, max_val))
            
            # Normalize weights to sum to 100%
            weight_sum = sum(sampled_weights)
            if weight_sum > 0:
                normalized_weights = [(w / weight_sum) * 100 for w in sampled_weights]
            else:
                normalized_weights = [0] * len(st.session_state.tiers)
            
            # Calculate weighted EER
            for tier, weight in zip(st.session_state.tiers, normalized_weights):
                tier_eer = calculate_tier_eer(tier['reward_value'], tier['points_required'], points_per_dollar)
                weighted_eer += tier_eer * (weight / 100)
            
            results.append(weighted_eer)
        
        # Calculate statistics
        mean_eer = np.mean(results)
        p10 = np.percentile(results, 10)
        p50 = np.percentile(results, 50)
        p90 = np.percentile(results, 90)
        min_eer = min(results)
        max_eer = max(results)
        
        guidance = get_guidance(mean_eer)
        
        st.session_state.simulation_results = {
            'results': results,
            'mean': mean_eer,
            'p10': p10,
            'p50': p50,
            'p90': p90,
            'min': min_eer,
            'max': max_eer,
            'points_per_dollar': points_per_dollar,
            'num_iterations': num_iterations,
            'tiers': [tier.copy() for tier in st.session_state.tiers],
            'guidance': guidance
        }
        
    except Exception as e:
        st.error(f"Error: {str(e)}")

def run_simulation_from_tiers(points_per_dollar, num_iterations, tiers):
    """Run Monte Carlo simulation with provided tiers (for batch processing)"""
    try:
        if not tiers:
            return None
        
        # Run simulation
        results = []
        
        for i in range(int(num_iterations)):
            weighted_eer = 0
            
            # Sample redemption percentages with uncertainty
            sampled_weights = []
            for tier in tiers:
                min_val = max(0, tier['redemption_pct'] - tier['uncertainty'])
                max_val = min(100, tier['redemption_pct'] + tier['uncertainty'])
                sampled_weights.append(np.random.uniform(min_val, max_val))
            
            # Normalize weights to sum to 100%
            weight_sum = sum(sampled_weights)
            if weight_sum > 0:
                normalized_weights = [(w / weight_sum) * 100 for w in sampled_weights]
            else:
                normalized_weights = [0] * len(tiers)
            
            # Calculate weighted EER
            for tier, weight in zip(tiers, normalized_weights):
                tier_eer = calculate_tier_eer(tier['reward_value'], tier['points_required'], points_per_dollar)
                weighted_eer += tier_eer * (weight / 100)
            
            results.append(weighted_eer)
        
        # Calculate statistics
        mean_eer = np.mean(results)
        p10 = np.percentile(results, 10)
        p50 = np.percentile(results, 50)
        p90 = np.percentile(results, 90)
        min_eer = min(results)
        max_eer = max(results)
        
        guidance = get_guidance(mean_eer)
        
        return {
            'results': results,
            'mean': mean_eer,
            'p10': p10,
            'p50': p50,
            'p90': p90,
            'min': min_eer,
            'max': max_eer,
            'points_per_dollar': points_per_dollar,
            'num_iterations': num_iterations,
            'tiers': tiers,
            'guidance': guidance
        }
        
    except Exception as e:
        return None

def process_csv_upload(uploaded_file):
    """Process uploaded CSV and create scenarios"""
    try:
        # Read CSV
        df = pd.read_csv(uploaded_file)
        
        # Validate required columns
        required_cols = ['scenario_name', 'points_per_dollar', 'num_simulations', 'tier_number', 
                        'points_required', 'reward_value', 'redemption_pct']
        
        missing_cols = [col for col in required_cols if col not in df.columns]
        if missing_cols:
            return False, f"Missing required columns: {', '.join(missing_cols)}"
        
        # Add uncertainty column if not present
        if 'uncertainty' not in df.columns:
            df['uncertainty'] = 5.0
        
        # Group by scenario
        scenarios_data = df.groupby('scenario_name')
        
        processed_count = 0
        error_count = 0
        errors = []
        
        # Process each scenario
        for scenario_name, scenario_df in scenarios_data:
            try:
                # Get scenario-level data
                points_per_dollar = int(scenario_df['points_per_dollar'].iloc[0])
                num_simulations = int(scenario_df['num_simulations'].iloc[0])
                
                # Build tiers
                tiers = []
                for _, row in scenario_df.iterrows():
                    tiers.append({
                        'points_required': int(row['points_required']),
                        'reward_value': float(row['reward_value']),
                        'redemption_pct': float(row['redemption_pct']),
                        'uncertainty': float(row['uncertainty'])
                    })
                
                # Validate redemption percentages sum to ~100
                total_redemption = sum(t['redemption_pct'] for t in tiers)
                if total_redemption < 99.9 or total_redemption > 100.1:
                    errors.append(f"{scenario_name}: Redemption % sums to {total_redemption:.1f}% (must be 100%)")
                    error_count += 1
                    continue
                
                # Run simulation
                result = run_simulation_from_tiers(points_per_dollar, num_simulations, tiers)
                
                if result:
                    # Save to scenarios
                    scenario = {
                        'name': scenario_name,
                        'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                        'results': result
                    }
                    st.session_state.saved_scenarios.append(scenario)
                    processed_count += 1
                else:
                    errors.append(f"{scenario_name}: Simulation failed")
                    error_count += 1
                    
            except Exception as e:
                errors.append(f"{scenario_name}: {str(e)}")
                error_count += 1
        
        # Build result message
        if processed_count > 0 and error_count == 0:
            return True, f"Successfully processed {processed_count} scenario(s)!"
        elif processed_count > 0 and error_count > 0:
            error_msg = "\n".join(errors)
            return True, f"Processed {processed_count} scenario(s). {error_count} failed:\n{error_msg}"
        else:
            error_msg = "\n".join(errors)
            return False, f"All scenarios failed:\n{error_msg}"
            
    except Exception as e:
        return False, f"Error reading CSV: {str(e)}"

def generate_csv_template():
    """Generate a CSV template for download"""
    template_data = {
        'scenario_name': ['Example Scenario 1', 'Example Scenario 1', 'Example Scenario 1', 
                         'Example Scenario 2', 'Example Scenario 2', 'Example Scenario 2'],
        'points_per_dollar': [1, 1, 1, 2, 2, 2],
        'num_simulations': [10000, 10000, 10000, 10000, 10000, 10000],
        'tier_number': [1, 2, 3, 1, 2, 3],
        'points_required': [500, 2500, 10000, 500, 2500, 10000],
        'reward_value': [5.0, 25.0, 100.0, 10.0, 50.0, 200.0],
        'redemption_pct': [50.0, 30.0, 20.0, 50.0, 30.0, 20.0],
        'uncertainty': [5.0, 5.0, 5.0, 5.0, 5.0, 5.0]
    }
    
    df = pd.DataFrame(template_data)
    return df.to_csv(index=False)

# Header
st.markdown("<h1 style='font-size: 2rem;'>Effective Earn Rate (EER) Simulator</h1>", unsafe_allow_html=True)
st.markdown("<p style='font-size: 0.95rem;'>Model uncertainty in redemption patterns to understand EER distribution and inform program design decisions</p>", unsafe_allow_html=True)

# CSV Upload Section
with st.expander("📤 Batch Upload Scenarios (CSV)", expanded=False):
    st.markdown("Upload a CSV file to automatically create and run multiple scenarios.")
    
    col1, col2 = st.columns([2, 1])
    with col1:
        uploaded_file = st.file_uploader("Choose a CSV file", type=['csv'], key="csv_upload")
    with col2:
        # Download template button
        template_csv = generate_csv_template()
        st.download_button(
            label="📥 Download Template",
            data=template_csv,
            file_name="eer_scenarios_template.csv",
            mime="text/csv",
            help="Download a template CSV to see the required format"
        )
    
    if uploaded_file is not None:
        # Show preview
        try:
            preview_df = pd.read_csv(uploaded_file)
            st.markdown("**Preview:**")
            st.dataframe(preview_df.head(10), use_container_width=True)
            
            # Reset file pointer for processing
            uploaded_file.seek(0)
            
            if st.button("Process CSV and Create Scenarios", type="primary"):
                with st.spinner("Processing scenarios..."):
                    success, message = process_csv_upload(uploaded_file)
                    if success:
                        st.success(message)
                        st.rerun()
                    else:
                        st.error(message)
        except Exception as e:
            st.error(f"Error reading file: {str(e)}")
    
    st.markdown("""
    **CSV Format (long format):**
    - `scenario_name`: Name of the scenario
    - `points_per_dollar`: Earn rate (whole number)
    - `num_simulations`: Number of simulations (1000, 10000, or 50000)
    - `tier_number`: Tier number (1, 2, 3, etc.)
    - `points_required`: Points needed for reward (whole number)
    - `reward_value`: Dollar value of reward
    - `redemption_pct`: Expected redemption percentage
    - `uncertainty`: Uncertainty range (optional, defaults to 5.0)
    
    **Note:** Each scenario should have multiple rows (one per tier). Redemption percentages must sum to 100% per scenario.
    """)

st.markdown("---")

# Global settings
col1, col2 = st.columns(2)
with col1:
    # Use loaded value if available, otherwise use default
    default_ppd = st.session_state.loaded_points_per_dollar if st.session_state.loaded_points_per_dollar else 1
    points_per_dollar = st.number_input("Points Per Dollar (Earn Rate)", min_value=1, value=default_ppd, step=1, key="ppd_input")
    # Clear loaded value after using it
    if st.session_state.loaded_points_per_dollar:
        st.session_state.loaded_points_per_dollar = None
with col2:
    # Use loaded value if available, otherwise use default
    default_sims = st.session_state.loaded_num_simulations if st.session_state.loaded_num_simulations else 10000
    num_iterations = st.selectbox("Number of Simulations", options=[1000, 10000, 50000], index=[1000, 10000, 50000].index(default_sims), key="sims_input")
    # Clear loaded value after using it
    if st.session_state.loaded_num_simulations:
        st.session_state.loaded_num_simulations = None

st.markdown("---")

st.markdown("<h2 style='font-size: 1.3rem; margin-bottom: 0.5rem;'>Redemption Tiers</h2>", unsafe_allow_html=True)

# Display tiers in compact format
for idx, tier in enumerate(st.session_state.tiers):
    cols = st.columns([0.8, 1.2, 1.2, 1.2, 1.2, 1.2, 0.6])
    
    with cols[0]:
        st.markdown(f"<p style='margin-top: 1.8rem; font-weight: 600; font-size: 0.9rem;'>Tier {idx + 1}</p>", unsafe_allow_html=True)
    
    with cols[1]:
        new_points = st.text_input(
            "Points Required",
            value=str(int(tier['points_required'])),
            key=f"points_{tier['id']}",
            label_visibility="visible"
        )
        try:
            tier['points_required'] = int(new_points)
        except:
            tier['points_required'] = 1
    
    with cols[2]:
        new_reward = st.text_input(
            "Reward Value ($)",
            value=str(tier['reward_value']),
            key=f"reward_{tier['id']}",
            label_visibility="visible"
        )
        try:
            tier['reward_value'] = float(new_reward)
        except:
            tier['reward_value'] = 0.0
    
    with cols[3]:
        new_redemption = st.text_input(
            "Redemption %",
            value=str(tier['redemption_pct']),
            key=f"redemption_{tier['id']}",
            label_visibility="visible"
        )
        try:
            tier['redemption_pct'] = float(new_redemption)
        except:
            tier['redemption_pct'] = 0.0
    
    with cols[4]:
        new_uncertainty = st.text_input(
            "Uncertainty (±%)",
            value=str(tier['uncertainty']),
            key=f"uncertainty_{tier['id']}",
            label_visibility="visible"
        )
        try:
            tier['uncertainty'] = float(new_uncertainty)
        except:
            tier['uncertainty'] = 0.0
    
    with cols[5]:
        tier_eer = calculate_tier_eer(tier['reward_value'], tier['points_required'], points_per_dollar)
        st.markdown("<p style='font-size: 0.75rem; margin-bottom: 0.2rem; color: #666;'>Tier EER</p>", unsafe_allow_html=True)
        st.markdown(f"<p style='font-size: 0.95rem; font-weight: 600; margin-top: 0;'>{tier_eer:.2f}%</p>", unsafe_allow_html=True)
    
    with cols[6]:
        if len(st.session_state.tiers) > 1:
            st.markdown("<p style='font-size: 0.75rem; margin-bottom: 0.2rem; opacity: 0;'>.</p>", unsafe_allow_html=True)
            if st.button("🗑️", key=f"remove_{tier['id']}", help="Remove this tier"):
                remove_tier(tier['id'])
                st.rerun()

# Add tier and uniform buttons
col1, col2 = st.columns(2)
with col1:
    if st.button("➕ Add Tier", key="add_tier"):
        add_tier()
        st.rerun()
with col2:
    if st.button("Apply Uniform Redemption Rates", key="uniform_btn"):
        st.session_state.show_uniform_confirm = True
        st.rerun()

# Confirmation dialog for uniform distribution
if st.session_state.show_uniform_confirm:
    st.warning("⚠️ This will overwrite all current redemption percentages. Are you sure?")
    col1, col2, col3 = st.columns([1, 1, 4])
    with col1:
        if st.button("Yes, Apply", type="primary"):
            apply_uniform_redemption()
            st.rerun()
    with col2:
        if st.button("Cancel"):
            st.session_state.show_uniform_confirm = False
            st.rerun()

st.markdown("")

# Check redemption total and show warning/disable button
redemption_total = check_redemption_total()
is_valid = 99.9 <= redemption_total <= 100.1

if not is_valid:
    st.markdown(f'<p style="color: #dc2626; font-weight: bold; font-size: 0.95rem;">⚠️ Warning: Redemption percentages sum to {redemption_total:.1f}%. They must total 100% to run simulation.</p>', unsafe_allow_html=True)

# Run simulation button
if st.button("▶️ Run Simulation", type="primary", use_container_width=True, disabled=not is_valid):
    run_simulation(points_per_dollar, num_iterations)

# Display results if available
if st.session_state.simulation_results:
    st.markdown("---")
    st.markdown("<h2 style='font-size: 1.3rem;'>Results</h2>", unsafe_allow_html=True)
    
    results = st.session_state.simulation_results
    
    # Save scenario section
    col1, col2 = st.columns([3, 1])
    with col1:
        scenario_name = st.text_input("Scenario Name", placeholder="Enter name to save...", key="scenario_name_input")
    with col2:
        st.write("")
        st.write("")
        if st.button("💾 Save Scenario", disabled=not scenario_name):
            scenario = {
                'name': scenario_name,
                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                'results': results
            }
            st.session_state.saved_scenarios.append(scenario)
            st.success(f"Saved '{scenario_name}'!")
            st.rerun()
    
    # Statistics
    cols = st.columns(5)
    with cols[0]:
        st.metric("Mean EER", f"{results['mean']:.2f}%")
    with cols[1]:
        st.metric("10th Percentile", f"{results['p10']:.2f}%")
    with cols[2]:
        st.metric("Median", f"{results['p50']:.2f}%")
    with cols[3]:
        st.metric("90th Percentile", f"{results['p90']:.2f}%")
    with cols[4]:
        st.metric("Range", f"{results['min']:.2f}% - {results['max']:.2f}%")
    
    # Guidance box with key considerations
    guidance = results['guidance']
    st.markdown(f"""
    <div style="padding: 1rem; background-color: {guidance['color']}20; border-radius: 0.5rem; margin: 1rem 0;">
        <h3 style="color: {guidance['color']}; margin: 0 0 0.5rem 0; font-size: 1.5rem; font-weight: 700;">{guidance['category']} EER ({results['mean']:.2f}%)</h3>
        {guidance['guidance']}
    </div>
    """, unsafe_allow_html=True)
    
    # Distribution chart
    st.markdown("<h3 style='font-size: 1.1rem;'>EER Distribution</h3>", unsafe_allow_html=True)
    
    fig = go.Figure()
    
    # Create histogram
    fig.add_trace(go.Histogram(
        x=results['results'],
        nbinsx=50,
        name='EER Distribution',
        marker_color='rgba(99, 102, 241, 0.7)',
        hovertemplate='EER: %{x:.2f}%<br>Count: %{y}<extra></extra>'
    ))
    
    # Add reference lines for zones
    fig.add_vline(x=6, line_dash="dash", line_color="#10b981", line_width=2)
    fig.add_vline(x=7.99, line_dash="dash", line_color="#10b981", line_width=2)
    fig.add_vline(x=8, line_dash="dash", line_color="#3b82f6", line_width=2)
    fig.add_vline(x=9.99, line_dash="dash", line_color="#3b82f6", line_width=2)
    fig.add_vline(x=10, line_dash="dash", line_color="#f59e0b", line_width=2)
    fig.add_vline(x=12, line_dash="dash", line_color="#f59e0b", line_width=2)
    
    # Add shaded regions
    fig.add_vrect(x0=6, x1=7.99, fillcolor="#10b981", opacity=0.1, line_width=0)
    fig.add_vrect(x0=8, x1=9.99, fillcolor="#3b82f6", opacity=0.1, line_width=0)
    fig.add_vrect(x0=10, x1=12, fillcolor="#f59e0b", opacity=0.1, line_width=0)
    
    fig.update_layout(
        xaxis_title="Effective Earn Rate (%)",
        yaxis_title="Frequency",
        showlegend=False,
        height=500,
        annotations=[
            dict(x=7, y=1, yref="paper", text="Conservative<br>6-7.99%", showarrow=False, yshift=10),
            dict(x=9, y=1, yref="paper", text="Moderate<br>8-9.99%", showarrow=False, yshift=10),
            dict(x=11, y=1, yref="paper", text="High Generosity<br>10-12%", showarrow=False, yshift=10)
        ]
    )
    
    st.plotly_chart(fig, use_container_width=True)
    
    # Legend
    st.markdown("""
    <div style="display: flex; justify-content: center; gap: 2rem; margin-top: 1rem;">
        <div style="display: flex; align-items: center; gap: 0.5rem;">
            <div style="width: 1rem; height: 1rem; background-color: #10b981; border-radius: 0.25rem;"></div>
            <span style="font-size: 0.9rem;">Conservative (6-7.99%)</span>
        </div>
        <div style="display: flex; align-items: center; gap: 0.5rem;">
            <div style="width: 1rem; height: 1rem; background-color: #3b82f6; border-radius: 0.25rem;"></div>
            <span style="font-size: 0.9rem;">Moderate (8-9.99%)</span>
        </div>
        <div style="display: flex; align-items: center; gap: 0.5rem;">
            <div style="width: 1rem; height: 1rem; background-color: #f59e0b; border-radius: 0.25rem;"></div>
            <span style="font-size: 0.9rem;">High Generosity (10-12%)</span>
        </div>
    </div>
    """, unsafe_allow_html=True)

# Display saved scenarios
if st.session_state.saved_scenarios:
    st.markdown("---")
    st.markdown("<h2 style='font-size: 1.3rem;'>Saved Scenarios</h2>", unsafe_allow_html=True)
    
    # Multi-select for comparison
    scenario_names = [s['name'] for s in st.session_state.saved_scenarios]
    selected_scenarios = st.multiselect(
        "Select scenarios to compare",
        options=scenario_names,
        key="scenario_comparison"
    )
    
    if selected_scenarios:
        st.markdown("<h3 style='font-size: 1.1rem;'>Scenario Comparison</h3>", unsafe_allow_html=True)
        
        # Create comparison chart
        fig_compare = go.Figure()
        
        colors = ['#6366f1', '#f59e0b', '#10b981', '#ec4899', '#8b5cf6', '#06b6d4']
        
        for idx, scenario_name in enumerate(selected_scenarios):
            scenario = next(s for s in st.session_state.saved_scenarios if s['name'] == scenario_name)
            res = scenario['results']
            
            fig_compare.add_trace(go.Histogram(
                x=res['results'],
                nbinsx=50,
                name=scenario_name,
                marker_color=colors[idx % len(colors)],
                opacity=0.6,
                hovertemplate=f'{scenario_name}<br>EER: %{{x:.2f}}%<br>Count: %{{y}}<extra></extra>'
            ))
        
        # Add reference lines
        fig_compare.add_vline(x=6, line_dash="dash", line_color="#10b981", line_width=1)
        fig_compare.add_vline(x=7.99, line_dash="dash", line_color="#10b981", line_width=1)
        fig_compare.add_vline(x=8, line_dash="dash", line_color="#3b82f6", line_width=1)
        fig_compare.add_vline(x=9.99, line_dash="dash", line_color="#3b82f6", line_width=1)
        fig_compare.add_vline(x=10, line_dash="dash", line_color="#f59e0b", line_width=1)
        fig_compare.add_vline(x=12, line_dash="dash", line_color="#f59e0b", line_width=1)
        
        fig_compare.update_layout(
            xaxis_title="Effective Earn Rate (%)",
            yaxis_title="Frequency",
            barmode='overlay',
            height=500,
            showlegend=True,
            legend=dict(yanchor="top", y=0.99, xanchor="right", x=0.99)
        )
        
        st.plotly_chart(fig_compare, use_container_width=True)
    
    st.markdown("---")
    
    # Display individual scenarios
    for scenario in st.session_state.saved_scenarios:
        with st.expander(f"📊 {scenario['name']} - {scenario['timestamp']}"):
            res = scenario['results']
            guidance = get_guidance(res['mean'])
            
            cols = st.columns(5)
            with cols[0]:
                st.metric("Mean EER", f"{res['mean']:.2f}%")
            with cols[1]:
                st.metric("P10", f"{res['p10']:.2f}%")
            with cols[2]:
                st.metric("Median", f"{res['p50']:.2f}%")
            with cols[3]:
                st.metric("P90", f"{res['p90']:.2f}%")
            with cols[4]:
                st.metric("Pts/$", f"{res['points_per_dollar']}")
            
            st.markdown(f"<p style='font-size: 0.9rem;'><strong>Category:</strong> {guidance['category']}</p>", unsafe_allow_html=True)
            st.markdown(f"<p style='font-size: 0.9rem;'><strong>Tiers:</strong> {len(res['tiers'])}</p>", unsafe_allow_html=True)
            
            if st.button("✏️ Edit Scenario", key=f"edit_{scenario['name']}"):
                load_scenario(scenario)
                st.success(f"Loaded '{scenario['name']}' into form above!")
                st.rerun()

# Custom CSS for better styling
st.markdown("""
<style>
    /* Remove +/- buttons from number inputs */
    input[type=number]::-webkit-inner-spin-button,
    input[type=number]::-webkit-outer-spin-button {
        -webkit-appearance: none;
        margin: 0;
    }
    input[type=number] {
        -moz-appearance: textfield;
    }
    
    /* Tighter spacing for inputs */
    .stTextInput > div > div > input {
        padding: 0.25rem 0.5rem;
        font-size: 0.9rem;
    }
    
    .stTextInput > label {
        font-size: 0.75rem;
        margin-bottom: 0.2rem;
    }
    
    /* Style remove button */
    button[key^="remove_"] {
        background-color: #fee2e2 !important;
        color: #dc2626 !important;
        border: 1px solid #dc2626 !important;
        font-size: 1.1rem !important;
        padding: 0.3rem 0.6rem !important;
        height: 2.2rem !important;
        margin-top: 1.4rem !important;
    }
    button[key^="remove_"]:hover {
        background-color: #fecaca !important;
    }
    
    /* Reduce overall font size differences */
    h1 {
        font-size: 2rem !important;
    }
    h2 {
        font-size: 1.3rem !important;
    }
    h3 {
        font-size: 1.1rem !important;
    }
    p {
        font-size: 0.95rem !important;
    }
    
    /* Compact metric styling */
    [data-testid="stMetricValue"] {
        font-size: 1.2rem !important;
    }
    
    [data-testid="stMetricLabel"] {
        font-size: 0.85rem !important;
    }
</style>
""", unsafe_allow_html=True)