401k.py

def compute_tax(income):
    tax      = 0
    def tax_income(rate, top_bound):
        nonlocal income, tax
        removed  = income  - max(income - top_bound, 0)
#        print(removed, '\ttaxed at', rate, '\tfor', removed * rate)
        tax     += removed * rate
        income  -= removed

    tax_income(0.1,   10_275)
    tax_income(0.12,  41_775)
    tax_income(0.22,  89_075)
    tax_income(0.24, 170_050)
    tax_income(0.32, 215_950)
    tax_income(0.35, 539_001)
    tax_income(0.37, income)
    return tax

annual_income = 140_000
annual_401k   = 23_000
employer_match = 0.5
employer_limit = 0.04
annual_401k_growth = 0.10
projected_years = 10
retirment_withdrawl = 50_000

employer_contribution = min(annual_401k * employer_match, annual_income * employer_limit)
annual_income -= annual_401k

# PRE TAX Computation
total_income = annual_income * projected_years
total_tax    = compute_tax(annual_income - annual_401k) * projected_years
retirment_savings = 0
annual_contributions = employer_contribution + annual_401k
for year in range(projected_years):
    retirment_savings *= (1 + annual_401k_growth)
    retirment_savings += annual_contributions

retirement_tax = compute_tax(retirment_withdrawl)
pretax_networth = (total_income - total_tax) + (retirment_savings)
print("Pretax Withdrawl Tax / yr: ", int(retirement_tax))
print("PreTax Networth: ", int(pretax_networth))

# Roth 401k Computation
total_income = annual_income * projected_years
total_tax    = compute_tax(annual_income) * projected_years
retirment_savings = 0
annual_contributions = employer_contribution + annual_401k
for year in range(projected_years):
    retirment_savings *= (1 + annual_401k_growth)
    retirment_savings += annual_contributions

roth_networth = (total_income - total_tax) + (retirment_savings)
print("Roth Networth: ", int(roth_networth))