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)
    
    # Add margin calculations
    df['vote_margin'] = df['percent_democrat'] - df['percent_republican'] 
    df['win_prob_margin'] = df['win_prob_dem'] - df['win_prob_rep']
    
    # Create visualization options
    visualizations = {
        "Vote Margin (Red-Blue)": "vote_margin",
        "Win Probability Margin (Red-Blue)": "win_prob_margin",
        "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 = {
        "Vote Margin (Red-Blue)": "RdBu",
        "Win Probability Margin (Red-Blue)": "RdBu",
        "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 = {
        "Vote Margin (Red-Blue)": [-30, 30],
        "Win Probability Margin (Red-Blue)": [-100, 100],
        "Win Probability (Republican)": [0, 100],
        "Median Vote Share (Republican)": [35, 65],
        "Win Probability (Democrat)": [0, 100],
        "Median Vote Share (Democrat)": [35, 65]
    }
    
    # Add custom hover text
    def create_hover_text(row, metric):
        if "Margin" in metric:
            if row[visualizations[metric]] > 0:
                lead = f"R+{abs(row[visualizations[metric]]):.1f}"
            else:
                lead = f"D+{abs(row[visualizations[metric]]):.1f}"
            return f"<b>{row['stateName']}</b><br>" + \
                   f"Margin: {lead}<br>" + \
                   f"Electoral Votes: {row['evs']}<br>" + \
                   f"Tipping Point: {row['tipping']:.1f}%<br>" + \
                   f"Voter Power Index: {row['vpi']:.2f}"
        else:
            return f"<b>{row['stateName']}</b><br>" + \
                   f"{metric}: {row[visualizations[metric]]:.1f}%<br>" + \
                   f"Electoral Votes: {row['evs']}<br>" + \
                   f"Tipping Point: {row['tipping']:.1f}%<br>" + \
                   f"Voter Power Index: {row['vpi']:.2f}"

    df['hover_text'] = df.apply(lambda row: create_hover_text(row, selected_metric), axis=1)
    
    # Create the choropleth map
    fig = go.Figure(data=go.Choropleth(
        locations=df['state'],
        z=df[visualizations[selected_metric]],
        locationmode='USA-states',
        colorscale=color_scales[selected_metric],
        zmin=ranges[selected_metric][0],
        zmax=ranges[selected_metric][1],
        text=df['hover_text'],
        hoverinfo='text',
        colorbar_title=selected_metric
    ))
    
    # Update layout
    fig.update_layout(
        title=f"2024 Election Forecast - {selected_metric}",
        geo=dict(
            scope='usa',
            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, col3 = st.columns(3)
    
    with col1:
        st.metric(
            "Total Electoral Votes",
            value=df['evs'].sum()
        )
    
    with col2:
        if "Margin" in selected_metric:
            avg_margin = df[visualizations[selected_metric]].mean()
            lead = "R+" if avg_margin > 0 else "D+"
            st.metric(
                f"Average {selected_metric}",
                value=f"{lead}{abs(avg_margin):.1f}"
            )
        else:
            avg_metric = df[visualizations[selected_metric]].mean()
            st.metric(
                f"Average {selected_metric}",
                value=f"{avg_metric:.1f}%"
            )
    
    with col3:
        tipping_state = df.loc[df['tipping'].idxmax()]
        st.metric(
            "Tipping Point State",
            value=f"{tipping_state['stateName']}"
        )

if __name__ == "__main__":
    main()