app.py

import streamlit as st
import pandas as pd
from policyengine_us.variables.household.demographic.geographic.state_code import (
    StateCode,
)
import table
import heatmap

# Streamlit heading and description
st.header("Marriage Incentive Calculator")
st.write(
    "This application evaluates marriage penalties and bonuses of couples, based on state and individual employment income"
)
st.markdown(
    "This application utilizes the [`policyengine-us` Python package](https://github.com/policyengine/policyengine-us)."
)

# Streamlit inputs for state code, head income, and spouse income
statecodes = [s.value for s in StateCode]
US_TERRITORIES = {
    "GU": "Guam",
    "MP": "Northern Mariana Islands",
    "PW": "Palau",
    "PR": "Puerto Rico",
    "VI": "Virgin Islands",
    "AA": "Armed Forces Americas",
    "AE": "Armed Forces Africa/Canada/Europe/Middle East",
    "AP": "Armed Forces Pacific",
}
options = [value for value in statecodes if value not in US_TERRITORIES]
state_code = st.selectbox("State Code", options)
head_employment_income = st.number_input(
    "Head Employment Income", min_value=0, step=10000, value=0
)
spouse_employment_income = st.number_input(
    "Spouse Employment Income", min_value=0, step=10000, value=0
)
head_disability = st.checkbox("Head is disabled")
spouse_disability = st.checkbox("Spouse is disabled")
num_children = st.number_input("Number of Children", 0)
children_ages = {}
disability_status = {"head": head_disability, "spouse": spouse_disability}

for num in range(1, num_children + 1):
    children_ages[num] = st.number_input(f"Child {num} Age", 0)
    disability_status[f"child_{num}"] = st.checkbox(f"Child {num} is disabled")

submit = st.button("Calculate")


# Helper functions
def get_combined_data(programs, index, tab_name):
    categories = list(programs[0][index].keys())
    married_values = list(programs[0][index].values())
    head_values = list(programs[1][index].values())
    spouse_values = list(programs[2][index].values())
    separate_values = [x + y for x, y in zip(head_values, spouse_values)]

    return table.create_table_data(
        categories, married_values, separate_values, tab_name
    )


def display_dataframe(df, tab_name):
    st.markdown("### Current situation:")
    st.dataframe(df.drop(columns=["Tab"]), hide_index=True)
    fig = heatmap.create_heatmap_chart(
        state_code, children_ages, tab_name, disability_status
    )
    st.plotly_chart(fig, use_container_width=True)


# Main logic
if submit:
    programs = table.get_categorized_programs(
        state_code,
        head_employment_income,
        spouse_employment_income,
        children_ages,
        disability_status,
    )

    programs_list = [
        "Net income",
        "Benefits",
        "Refundable tax credits",
        "Taxes before refundable credits",
    ]
    married_programs = programs[0][:-3]
    head_separate = programs[1][:-3]
    spouse_separate = programs[2][:-3]
    separate = [x + y for x, y in zip(head_separate, spouse_separate)]

    total_data = table.create_table_data(
        programs_list, married_programs, separate, "Summary", filter_zeros=False
    )

    benefits_data = get_combined_data(programs, -2, "Benefits Breakdown")
    credits_data = get_combined_data(programs, -1, "Refundable Credits")
    taxes_data = get_combined_data(programs, -3, "Taxes before Refundable Credits")

    all_data = pd.concat([total_data, benefits_data, credits_data, taxes_data])

    tab1, tab2, tab3, tab4 = st.tabs(
        [
            "Summary",
            "Benefits Breakdown",
            "Refundable Credits",
            "Taxes before Refundable Credits",
        ]
    )

    with tab1:
        display_dataframe(all_data[all_data["Tab"] == "Summary"], "Net Income")
    with tab2:
        display_dataframe(all_data[all_data["Tab"] == "Benefits Breakdown"], "Benefits")
    with tab3:
        display_dataframe(
            all_data[all_data["Tab"] == "Refundable Credits"], "Refundable Tax Credits"
        )
    with tab4:
        display_dataframe(
            all_data[all_data["Tab"] == "Taxes before Refundable Credits"],
            "Tax Before Refundable Credits",
        )


# Add the note at the bottom
st.markdown("***")
st.markdown(
    "*We attribute all dependents to the head of household when considering unamrried filers*"
)