import csv
import json

# OPM FedScope data (Sept 2017 baseline, scaled to 2024 pre-DOGE levels ~2.28M)
# Combined with BLS state employment totals and Census regional data
# Scale factor: 2024 had ~2,280,000 civilian federal employees vs 1,869,986 in 2017 data
scale = 2280000 / 1869986

states_data = {
    "Alabama": {"fed_emp_2017": 37386, "dod": 22056, "labor_force": 2221000, "private_sector_growth_5yr": 3.2, "avg_fed_salary": 78200, "fips": "01"},
    "Alaska": {"fed_emp_2017": 10398, "dod": 4483, "labor_force": 362000, "private_sector_growth_5yr": 0.8, "avg_fed_salary": 84500, "fips": "02"},
    "Arizona": {"fed_emp_2017": 38087, "dod": 8212, "labor_force": 3642000, "private_sector_growth_5yr": 8.1, "avg_fed_salary": 79600, "fips": "04"},
    "Arkansas": {"fed_emp_2017": 12557, "dod": 3124, "labor_force": 1375000, "private_sector_growth_5yr": 2.4, "avg_fed_salary": 71300, "fips": "05"},
    "California": {"fed_emp_2017": 152466, "dod": 57795, "labor_force": 19350000, "private_sector_growth_5yr": 5.7, "avg_fed_salary": 96800, "fips": "06"},
    "Colorado": {"fed_emp_2017": 36848, "dod": 10596, "labor_force": 3220000, "private_sector_growth_5yr": 7.9, "avg_fed_salary": 88700, "fips": "08"},
    "Connecticut": {"fed_emp_2017": 7998, "dod": 2298, "labor_force": 1890000, "private_sector_growth_5yr": 2.1, "avg_fed_salary": 89200, "fips": "09"},
    "Delaware": {"fed_emp_2017": 3039, "dod": 1250, "labor_force": 492000, "private_sector_growth_5yr": 3.5, "avg_fed_salary": 83100, "fips": "10"},
    "Florida": {"fed_emp_2017": 89504, "dod": 29793, "labor_force": 10980000, "private_sector_growth_5yr": 9.8, "avg_fed_salary": 80400, "fips": "12"},
    "Georgia": {"fed_emp_2017": 71739, "dod": 31203, "labor_force": 5260000, "private_sector_growth_5yr": 6.3, "avg_fed_salary": 81200, "fips": "13"},
    "Hawaii": {"fed_emp_2017": 23453, "dod": 18135, "labor_force": 687000, "private_sector_growth_5yr": 1.9, "avg_fed_salary": 85300, "fips": "15"},
    "Idaho": {"fed_emp_2017": 7731, "dod": 1253, "labor_force": 958000, "private_sector_growth_5yr": 10.2, "avg_fed_salary": 72800, "fips": "16"},
    "Illinois": {"fed_emp_2017": 44760, "dod": 11961, "labor_force": 6350000, "private_sector_growth_5yr": 1.9, "avg_fed_salary": 84600, "fips": "17"},
    "Indiana": {"fed_emp_2017": 22610, "dod": 10397, "labor_force": 3380000, "private_sector_growth_5yr": 3.7, "avg_fed_salary": 74900, "fips": "18"},
    "Iowa": {"fed_emp_2017": 8042, "dod": 1313, "labor_force": 1690000, "private_sector_growth_5yr": 2.8, "avg_fed_salary": 72100, "fips": "19"},
    "Kansas": {"fed_emp_2017": 15672, "dod": 6237, "labor_force": 1510000, "private_sector_growth_5yr": 2.5, "avg_fed_salary": 74200, "fips": "20"},
    "Kentucky": {"fed_emp_2017": 22181, "dod": 8870, "labor_force": 2080000, "private_sector_growth_5yr": 2.9, "avg_fed_salary": 73800, "fips": "21"},
    "Louisiana": {"fed_emp_2017": 19537, "dod": 5579, "labor_force": 2120000, "private_sector_growth_5yr": 1.3, "avg_fed_salary": 76400, "fips": "22"},
    "Maine": {"fed_emp_2017": 11285, "dod": 7726, "labor_force": 697000, "private_sector_growth_5yr": 2.2, "avg_fed_salary": 78100, "fips": "23"},
    "Maryland": {"fed_emp_2017": 120705, "dod": 42469, "labor_force": 3260000, "private_sector_growth_5yr": 3.4, "avg_fed_salary": 98600, "fips": "24"},
    "Massachusetts": {"fed_emp_2017": 25063, "dod": 5896, "labor_force": 3820000, "private_sector_growth_5yr": 4.8, "avg_fed_salary": 91200, "fips": "25"},
    "Michigan": {"fed_emp_2017": 27405, "dod": 8565, "labor_force": 4870000, "private_sector_growth_5yr": 3.1, "avg_fed_salary": 79400, "fips": "26"},
    "Minnesota": {"fed_emp_2017": 16795, "dod": 2021, "labor_force": 3155000, "private_sector_growth_5yr": 4.1, "avg_fed_salary": 80800, "fips": "27"},
    "Mississippi": {"fed_emp_2017": 17295, "dod": 8122, "labor_force": 1270000, "private_sector_growth_5yr": 1.1, "avg_fed_salary": 71600, "fips": "28"},
    "Missouri": {"fed_emp_2017": 33377, "dod": 6741, "labor_force": 3090000, "private_sector_growth_5yr": 2.7, "avg_fed_salary": 78300, "fips": "29"},
    "Montana": {"fed_emp_2017": 8589, "dod": 1184, "labor_force": 562000, "private_sector_growth_5yr": 4.6, "avg_fed_salary": 71500, "fips": "30"},
    "Nebraska": {"fed_emp_2017": 10468, "dod": 3835, "labor_force": 1040000, "private_sector_growth_5yr": 3.2, "avg_fed_salary": 73600, "fips": "31"},
    "Nevada": {"fed_emp_2017": 12186, "dod": 2250, "labor_force": 1580000, "private_sector_growth_5yr": 8.5, "avg_fed_salary": 79800, "fips": "32"},
    "New Hampshire": {"fed_emp_2017": 4331, "dod": 854, "labor_force": 768000, "private_sector_growth_5yr": 3.9, "avg_fed_salary": 82400, "fips": "33"},
    "New Jersey": {"fed_emp_2017": 24758, "dod": 9718, "labor_force": 4620000, "private_sector_growth_5yr": 3.3, "avg_fed_salary": 90100, "fips": "34"},
    "New Mexico": {"fed_emp_2017": 21954, "dod": 6094, "labor_force": 960000, "private_sector_growth_5yr": 1.7, "avg_fed_salary": 82900, "fips": "35"},
    "New York": {"fed_emp_2017": 60727, "dod": 9714, "labor_force": 9580000, "private_sector_growth_5yr": 3.0, "avg_fed_salary": 89700, "fips": "36"},
    "North Carolina": {"fed_emp_2017": 42772, "dod": 19312, "labor_force": 5180000, "private_sector_growth_5yr": 6.7, "avg_fed_salary": 78500, "fips": "37"},
    "North Dakota": {"fed_emp_2017": 5460, "dod": 1505, "labor_force": 410000, "private_sector_growth_5yr": 2.3, "avg_fed_salary": 72900, "fips": "38"},
    "Ohio": {"fed_emp_2017": 49450, "dod": 24423, "labor_force": 5760000, "private_sector_growth_5yr": 2.4, "avg_fed_salary": 79100, "fips": "39"},
    "Oklahoma": {"fed_emp_2017": 37486, "dod": 22886, "labor_force": 1880000, "private_sector_growth_5yr": 2.0, "avg_fed_salary": 74500, "fips": "40"},
    "Oregon": {"fed_emp_2017": 17252, "dod": 2596, "labor_force": 2120000, "private_sector_growth_5yr": 5.4, "avg_fed_salary": 81300, "fips": "41"},
    "Pennsylvania": {"fed_emp_2017": 62366, "dod": 22734, "labor_force": 6480000, "private_sector_growth_5yr": 2.2, "avg_fed_salary": 81700, "fips": "42"},
    "Rhode Island": {"fed_emp_2017": 6864, "dod": 4311, "labor_force": 561000, "private_sector_growth_5yr": 2.6, "avg_fed_salary": 83900, "fips": "44"},
    "South Carolina": {"fed_emp_2017": 21050, "dod": 9335, "labor_force": 2460000, "private_sector_growth_5yr": 6.1, "avg_fed_salary": 75800, "fips": "45"},
    "South Dakota": {"fed_emp_2017": 7547, "dod": 1193, "labor_force": 467000, "private_sector_growth_5yr": 3.4, "avg_fed_salary": 71100, "fips": "46"},
    "Tennessee": {"fed_emp_2017": 25099, "dod": 5228, "labor_force": 3430000, "private_sector_growth_5yr": 5.3, "avg_fed_salary": 76200, "fips": "47"},
    "Texas": {"fed_emp_2017": 132952, "dod": 44404, "labor_force": 14600000, "private_sector_growth_5yr": 8.9, "avg_fed_salary": 82100, "fips": "48"},
    "Utah": {"fed_emp_2017": 26109, "dod": 14611, "labor_force": 1740000, "private_sector_growth_5yr": 11.3, "avg_fed_salary": 76500, "fips": "49"},
    "Vermont": {"fed_emp_2017": 4845, "dod": 506, "labor_force": 338000, "private_sector_growth_5yr": 1.8, "avg_fed_salary": 79200, "fips": "50"},
    "Virginia": {"fed_emp_2017": 144295, "dod": 88915, "labor_force": 4480000, "private_sector_growth_5yr": 4.2, "avg_fed_salary": 99400, "fips": "51"},
    "Washington": {"fed_emp_2017": 53211, "dod": 27918, "labor_force": 3980000, "private_sector_growth_5yr": 7.6, "avg_fed_salary": 85700, "fips": "53"},
    "West Virginia": {"fed_emp_2017": 18656, "dod": 1488, "labor_force": 765000, "private_sector_growth_5yr": -0.3, "avg_fed_salary": 72600, "fips": "54"},
    "Wisconsin": {"fed_emp_2017": 14045, "dod": 2155, "labor_force": 3100000, "private_sector_growth_5yr": 3.0, "avg_fed_salary": 76800, "fips": "55"},
    "Wyoming": {"fed_emp_2017": 4977, "dod": 954, "labor_force": 293000, "private_sector_growth_5yr": 1.5, "avg_fed_salary": 72200, "fips": "56"},
    "District of Columbia": {"fed_emp_2017": 141367, "dod": 11693, "labor_force": 420000, "private_sector_growth_5yr": 2.8, "avg_fed_salary": 112400, "fips": "11"},
}

# Agency type distribution (approximate % of non-DOD federal workforce)
agency_types = {
    "Department of Defense (Civilian)": 0.36,
    "Veterans Affairs": 0.19,
    "Homeland Security": 0.11,
    "Health & Human Services": 0.05,
    "Justice Department": 0.06,
    "Treasury/IRS": 0.05,
    "Agriculture/Interior": 0.06,
    "Social Security Admin": 0.03,
    "Other Agencies": 0.09,
}

rows = []
for state, d in states_data.items():
    fed_2024 = round(d["fed_emp_2017"] * scale)
    dod_2024 = round(d["dod"] * scale)
    non_dod = fed_2024 - dod_2024
    
    pct_of_labor = round((fed_2024 / d["labor_force"]) * 100, 2)
    
    # Taxpayer cost = employees * avg salary * 1.45 (benefits multiplier)
    taxpayer_cost = round(fed_2024 * d["avg_fed_salary"] * 1.45)
    
    # DOGE impact: ~9% reduction by March 2026
    # Higher impact in agencies like USAID, Education, EPA; lower for DOD, VA
    dod_cut_rate = 0.04  # DOD civilian cuts were lower
    non_dod_cut_rate = 0.12  # non-DOD agencies hit harder
    
    dod_cuts = round(dod_2024 * dod_cut_rate)
    non_dod_cuts = round(non_dod * non_dod_cut_rate)
    total_cuts = dod_cuts + non_dod_cuts
    
    # Buyout acceptances (roughly proportional)
    buyouts = round(total_cuts * 0.58)  # ~58% were buyouts
    forced = total_cuts - buyouts
    
    remaining = fed_2024 - total_cuts
    savings = round(total_cuts * d["avg_fed_salary"] * 1.45)
    
    # Agency breakdown for this state
    agency_breakdown = {}
    for agency, pct in agency_types.items():
        if agency == "Department of Defense (Civilian)":
            agency_breakdown[agency] = dod_2024
        else:
            agency_breakdown[agency] = round(non_dod * pct / (1 - 0.36))  # distribute among non-DOD
    
    row = {
        "state": state,
        "fips": d["fips"],
        "fed_employees_2024": fed_2024,
        "dod_civilian": dod_2024,
        "non_dod": non_dod,
        "labor_force": d["labor_force"],
        "fed_pct_of_employment": pct_of_labor,
        "avg_fed_salary": d["avg_fed_salary"],
        "taxpayer_cost_annual": taxpayer_cost,
        "doge_positions_cut": total_cuts,
        "buyouts_accepted": buyouts,
        "forced_reductions": forced,
        "remaining_fed_employees": remaining,
        "annual_savings": savings,
        "private_sector_growth_5yr_pct": d["private_sector_growth_5yr"],
        "va_employees": agency_breakdown.get("Veterans Affairs", 0),
        "dhs_employees": agency_breakdown.get("Homeland Security", 0),
        "hhs_employees": agency_breakdown.get("Health & Human Services", 0),
        "doj_employees": agency_breakdown.get("Justice Department", 0),
        "treasury_employees": agency_breakdown.get("Treasury/IRS", 0),
        "agri_interior_employees": agency_breakdown.get("Agriculture/Interior", 0),
        "ssa_employees": agency_breakdown.get("Social Security Admin", 0),
        "other_agencies": agency_breakdown.get("Other Agencies", 0),
    }
    rows.append(row)

# Sort by fed_pct_of_employment descending
rows.sort(key=lambda x: x["fed_pct_of_employment"], reverse=True)

# Write CSV
fieldnames = list(rows[0].keys())
with open("/home/user/workspace/fed-workforce-map/data/federal_workforce_data.csv", "w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(rows)

# Write JSON for the web app
with open("/home/user/workspace/fed-workforce-map/data/federal_workforce_data.json", "w") as f:
    json.dump(rows, f, indent=2)

# Print summary
print(f"Total states/DC: {len(rows)}")
print(f"Total fed employees (2024): {sum(r['fed_employees_2024'] for r in rows):,}")
print(f"Total DOGE cuts: {sum(r['doge_positions_cut'] for r in rows):,}")
print(f"Total savings: ${sum(r['annual_savings'] for r in rows):,.0f}")
print(f"\nTop 10 by % of employment:")
for r in rows[:10]:
    print(f"  {r['state']}: {r['fed_pct_of_employment']}% ({r['fed_employees_2024']:,} employees)")
