import streamlit as st
import folium
import numpy as np

def main():
    st.title('State Comparison')

    # Consolidate the state data and spending data into a list of dictionaries
    state_list = []
    for state, data in state_data.items():
        state_dict = {
            'state': state,
            'population': data['population'],
            'family_size': data['family_size'],
            'education_spending': spending_data[state]['education'],
            'healthcare_spending': spending_data[state]['healthcare'],
            'transportation_spending': spending_data[state]['transportation']
        }
        state_list.append(state_dict)

    # Save the data to a CSV file and provide a download link
    with open('state_data.csv', mode='w', newline='') as file:
        writer = csv.DictWriter(file, fieldnames=['state', 'population', 'family_size', 'education_spending', 'healthcare_spending', 'transportation_spending'])
        writer.writeheader()
        for state in state_list:
            writer.writerow(state)
    with open('state_data.csv', mode='rb') as file:
        b64 = base64.b64encode(file.read()).decode('utf-8')
    st.markdown(f'<a href="data:file/csv;base64,{b64}" download="state_data.csv">Download State Data CSV File</a>', unsafe_allow_html=True)

    # Create a map with pie charts for each state
    m = folium.Map(location=[37.0902, -95.7129], zoom_start=4)

    for state, data in state_data.items():
        lat, lon = get_coordinates(state)
        popup_text = f'<b>{state}</b><br>Population: {data["population"]}<br>Family Size: {data["family_size"]}'
        chart_data = [spending_data[state]['education'], spending_data[state]['healthcare'], spending_data[state]['transportation']]
        chart_labels = ['Education', 'Healthcare', 'Transportation']
        colors = ['blue', 'red', 'green']
        chart = create_pie_chart(chart_data, chart_labels, colors)
        folium.Marker(location=[lat, lon], popup=popup_text, icon=folium.Icon(color='gray', icon='info-sign')).add_to(m)
        folium.Marker(location=[lat, lon], icon=folium.Icon(color='cadetblue', icon='chart-pie'), tooltip='Click for chart').add_to(m).add_child(chart)

    st.markdown(folium_map_to_html(m), unsafe_allow_html=True)

def get_coordinates(state):
    # Replace this with your own function to get the coordinates for each state
    # This is just a placeholder
    return np.random.randint(-120, -70), np.random.randint(25, 50)

def create_pie_chart(data, labels, colors):
    chart = folium.plugins.PieChart(data, labels=labels, colors=colors, radius=30, weight=0.5)
    return chart

def folium_map_to_html(m):
    return m.get_root().render()