import pandas as pd
import numpy as np
from pathlib import Path

ROOT = Path(__file__).parent.parent.parent

input_dir = ROOT / "build/input/demographic_survey"
output_file = ROOT / "build/output/demographics_state_year.csv"

# --- Step 1: Loop over years and append ---
frames = []
for year in range(1995, 2016):
    fp = input_dir / f"demographic_survey_{year}.csv"
    df_year = pd.read_csv(fp)
    df_year["year"] = year
    frames.append(df_year)
survey = pd.concat(frames, ignore_index=True)

# --- Step 2: Filter ---
survey = survey[survey["state_fips"] != 51]
survey = survey[survey["year"] >= 2000]

# --- Step 3: Clean ---
survey["income"] = (
    survey["income"]
    .str.replace("$", "", regex=False)
    .str.replace(",", "", regex=False)
    .astype(float)
)

# --- Step 4: Collapse to state-year (weighted) ---
demographics = survey.groupby(["state_fips", "year"]).apply(
    lambda g: pd.Series({
        "population": g["weight"].sum(),
        "median_income": np.average(g["income"], weights=g["weight"]),
        "pct_urban": np.average(g["urban"], weights=g["weight"]),
    }),
    include_groups=False,
).reset_index()

assert len(demographics) == 800, f"Expected 800 rows, got {len(demographics)}"

output_file.parent.mkdir(parents=True, exist_ok=True)
demographics.to_csv(output_file, index=False)
print(f"Saved demographics_state_year.csv: {len(demographics)} rows")
