Code Editor for app.py

import streamlit as st
import plotly.graph_objects as go
import plotly.express as px
import pandas as pd
import us
from data import get_data

def main():
    st.title("2024 Election Forecast Visualization")
    
    # Get data
    df = get_data()
    df['percent_republican'] = df.apply(lambda x: x['pv']['rep']['median'], axis=1)
    df['percent_democrat'] = df.apply(lambda x: x['pv']['dem']['median'], axis=1)
    df['win_prob_rep'] = df.apply(lambda x: x['winprob']['rep'], axis=1)
    df['win_prob_dem'] = df.apply(lambda x: x['winprob']['dem'], axis=1)

    
    # Create visualization options
    visualizations = {
        "Win Probability (Republican)": "win_prob_rep",
        "Median Vote Share (Republican)": "percent_republican",
        "Win Probability (Democrat)": "win_prob_dem",
        "Median Vote Share (Democrat)": "percent_democrat"
    }
    
    # Create metric selector
    selected_metric = st.selectbox(
        "Select Visualization Metric",
        options=list(visualizations.keys())
    )
    
    # Create color scales for different metrics
    color_scales = {
        "Win Probability (Republican)": "Reds",
        "Median Vote Share (Republican)": "Reds",
        "Win Probability (Democrat)": "Blues",
        "Median Vote Share (Democrat)": "Blues"
    }
    
    # Create range colors for different metrics
    ranges = {
        "Win Probability (Republican)": [0, 100],
        "Median Vote Share (Republican)": [35, 65],
        "Win Probability (Democrat)": [0, 100],
        "Median Vote Share (Democrat)": [35, 65]
    }
    
    # Create the choropleth map
    fig = px.choropleth(
        df,
        locations='state',
        locationmode='USA-states',
        scope="usa",
        color=visualizations[selected_metric],
        hover_name='stateName',
        hover_data=['tipping', 'vpi', 'evs'],
        range_color=ranges[selected_metric],
        color_continuous_scale=color_scales[selected_metric],
        title=f"2024 Election Forecast - {selected_metric}"
    )
    
    # Update layout
    fig.update_layout(
        geo=dict(
            showlakes=True,
            lakecolor='rgb(255, 255, 255)'
        ),
        margin={"r":0,"t":30,"l":0,"b":0}
    )
    
    # Display the map
    st.plotly_chart(fig, use_container_width=True)
    
    # Add metrics summary
    st.subheader("Summary Statistics")
    col1, col2 = st.columns(2)
    
    with col1:
        st.metric(
            "Total Electoral Votes",
            value=df['evs'].sum()
        )
    
    with col2:
        if "Win Probability" in selected_metric:
            avg_prob = df[visualizations[selected_metric]].mean()
            st.metric(
                f"Average {selected_metric}",
                value=f"{avg_prob:.1f}%"
            )
        else:
            avg_share = df[visualizations[selected_metric]].mean()
            st.metric(
                f"Average {selected_metric}",
                value=f"{avg_share:.1f}%"
            )


if __name__ == "__main__":
    main()