Local Area Calibration Setup

import numpy as np
import pandas as pd
from collections import defaultdict

from policyengine_us import Microsimulation
from policyengine_us_data.storage import STORAGE_FOLDER
from policyengine_us_data.calibration.clone_and_assign import (
    assign_random_geography,
    GeographyAssignment,
    load_global_block_distribution,
)
from policyengine_us_data.calibration.unified_matrix_builder import (
    UnifiedMatrixBuilder,
)
from policyengine_us_data.calibration.unified_calibration import (
    rerandomize_takeup,
    SIMPLE_TAKEUP_VARS,
)
from policyengine_us_data.utils.randomness import seeded_rng
from policyengine_us_data.parameters import load_take_up_rate
from policyengine_us_data.datasets.cps.local_area_calibration.calibration_utils import (
    get_calculated_variables,
    STATE_CODES,
    get_all_cds_from_database,
)
from policyengine_us_data.datasets.cps.local_area_calibration.stacked_dataset_builder import (
    create_sparse_cd_stacked_dataset,
)

db_path = STORAGE_FOLDER / "calibration" / "policy_data.db"
db_uri = f"sqlite:///{db_path}"
dataset_path = str(STORAGE_FOLDER / "stratified_extended_cps_2024.h5")

N_CLONES = 3
SEED = 42

sim = Microsimulation(dataset=dataset_path)
hh_ids = sim.calculate("household_id", map_to="household").values
snap_values = sim.calculate("snap", map_to="household").values
n_records = len(hh_ids)

record_idx = 8629  # High SNAP ($18k), lands in TX/PA/NY with seed=42
example_hh_id = hh_ids[record_idx]
print(f"Base dataset: {n_records:,} households")
print(
    f"Example household: record_idx={record_idx}, "
    f"household_id={example_hh_id}, "
    f"SNAP=${snap_values[record_idx]:,.2f}"
)

geography = assign_random_geography(n_records, n_clones=N_CLONES, seed=SEED)
n_total = n_records * N_CLONES

print(f"Total cloned records: {n_total:,}")
print(f"Unique states: {len(np.unique(geography.state_fips))}")
print(f"Unique CDs: {len(np.unique(geography.cd_geoid))}")
print(f"Unique blocks: {len(np.unique(geography.block_geoid))}")

print(
    f"Example household (record_idx={record_idx}) across {N_CLONES} clones:\n"
)
rows = []
for c in range(N_CLONES):
    col = c * n_records + record_idx
    rows.append(
        {
            "clone": c,
            "col": col,
            "state_fips": geography.state_fips[col],
            "abbr": STATE_CODES.get(geography.state_fips[col], "??"),
            "cd_geoid": geography.cd_geoid[col],
            "block_geoid": geography.block_geoid[col],
        }
    )
pd.DataFrame(rows)

blocks, cds, states, probs = load_global_block_distribution()
print(f"Global block distribution: {len(blocks):,} blocks")
print(f"Top 5 states by total probability:")
state_prob = pd.Series(probs, index=states).groupby(level=0).sum()
top5 = state_prob.nlargest(5)
for fips, p in top5.items():
    print(f"  {STATE_CODES.get(fips, '??')} ({fips}): {p:.3%}")

clone_idx = 0
col_start = clone_idx * n_records
col_end = col_start + n_records
clone_states = geography.state_fips[col_start:col_end]

clone_sim = Microsimulation(dataset=dataset_path)
clone_sim.set_input("state_fips", 2024, clone_states.astype(np.int32))
for var in get_calculated_variables(clone_sim):
    clone_sim.delete_arrays(var)

new_snap = clone_sim.calculate("snap", map_to="household").values

orig_state = sim.calculate("state_fips", map_to="household").values[record_idx]
new_state = clone_states[record_idx]

print(f"Example household (record_idx={record_idx}):")
print(
    f"  Original state: {STATE_CODES.get(int(orig_state), '??')} "
    f"({int(orig_state)})"
)
print(
    f"  Clone 0 state:  {STATE_CODES.get(int(new_state), '??')} "
    f"({int(new_state)})"
)
print(f"  Original SNAP:  ${snap_values[record_idx]:,.2f}")
print(f"  Clone 0 SNAP:   ${new_snap[record_idx]:,.2f}")

print(f"SNAP for record_idx={record_idx} across all {N_CLONES} clones:\n")
rows = []
for c in range(N_CLONES):
    cs = geography.state_fips[c * n_records + record_idx]
    s = Microsimulation(dataset=dataset_path)
    s.set_input(
        "state_fips",
        2024,
        geography.state_fips[c * n_records : (c + 1) * n_records].astype(
            np.int32
        ),
    )
    for var in get_calculated_variables(s):
        s.delete_arrays(var)
    clone_snap = s.calculate("snap", map_to="household").values
    rows.append(
        {
            "clone": c,
            "state": STATE_CODES.get(int(cs), "??"),
            "state_fips": int(cs),
            "SNAP": f"${clone_snap[record_idx]:,.2f}",
        }
    )
pd.DataFrame(rows)

state_col_lists = defaultdict(list)
cd_col_lists = defaultdict(list)
for col in range(n_total):
    state_col_lists[int(geography.state_fips[col])].append(col)
    cd_col_lists[str(geography.cd_geoid[col])].append(col)

state_to_cols = {s: np.array(c) for s, c in state_col_lists.items()}
cd_to_cols = {cd: np.array(c) for cd, c in cd_col_lists.items()}

print(f"Unique states mapped: {len(state_to_cols)}")
print(f"Unique CDs mapped: {len(cd_to_cols)}")

state_counts = {s: len(c) for s, c in state_to_cols.items()}
sc_series = pd.Series(state_counts)
print(
    f"\nColumns per state: min={sc_series.min()}, "
    f"median={sc_series.median():.0f}, max={sc_series.max()}"
)

print(f"Example household clone visibility:\n")
for c in range(N_CLONES):
    col = c * n_records + record_idx
    state = int(geography.state_fips[col])
    cd = str(geography.cd_geoid[col])
    abbr = STATE_CODES.get(state, "??")
    print(f"Clone {c} ({abbr}, CD {cd}):")
    print(
        f"  Visible to {abbr} state targets: "
        f"col {col} in state_to_cols[{state}]? "
        f"{col in state_to_cols.get(state, [])}"
    )
    print(
        f"  Visible to CD {cd} targets: "
        f"col {col} in cd_to_cols['{cd}']? "
        f"{col in cd_to_cols.get(cd, [])}"
    )
    # Check an unrelated state
    print(
        f"  Visible to NC (37) targets: " f"{col in state_to_cols.get(37, [])}"
    )
    print()

print(f"{len(SIMPLE_TAKEUP_VARS)} takeup variables:\n")
for spec in SIMPLE_TAKEUP_VARS:
    rate_key = spec["rate_key"]
    if rate_key == "voluntary_filing":
        rate = 0.05
    else:
        rate = load_take_up_rate(rate_key, 2024)
    rate_str = (
        f"{rate:.2%}"
        if isinstance(rate, float)
        else f"dict ({len(rate)} entries)"
    )
    print(
        f"  {spec['variable']:40s} "
        f"entity={spec['entity']:10s} rate={rate_str}"
    )

block_a = "482011234567890"
block_b = "170311234567890"
var = "takes_up_snap_if_eligible"

rng_a1 = seeded_rng(var, salt=block_a)
rng_a2 = seeded_rng(var, salt=block_a)
rng_b = seeded_rng(var, salt=block_b)
rng_other = seeded_rng("takes_up_aca_if_eligible", salt=block_a)

draws_a1 = rng_a1.random(5)
draws_a2 = rng_a2.random(5)
draws_b = rng_b.random(5)
draws_other = rng_other.random(5)

print("Same block + same var (reproducible):")
print(f"  {draws_a1}")
print(f"  {draws_a2}")
print(f"  Match: {np.allclose(draws_a1, draws_a2)}")
print(f"\nDifferent block, same var:")
print(f"  {draws_b}")
print(f"  Match: {np.allclose(draws_a1, draws_b)}")
print(f"\nSame block, different var:")
print(f"  {draws_other}")
print(f"  Match: {np.allclose(draws_a1, draws_other)}")

test_sim = Microsimulation(dataset=dataset_path)
clone_0_states = geography.state_fips[:n_records]
clone_0_blocks = geography.block_geoid[:n_records]
test_sim.set_input("state_fips", 2024, clone_0_states.astype(np.int32))

before = {}
for spec in SIMPLE_TAKEUP_VARS:
    v = spec["variable"]
    vals = test_sim.calculate(v, map_to=spec["entity"]).values
    before[v] = vals.mean()

rerandomize_takeup(test_sim, clone_0_blocks, clone_0_states, 2024)

print("Takeup rates before/after re-randomization (clone 0):\n")
for spec in SIMPLE_TAKEUP_VARS:
    v = spec["variable"]
    vals = test_sim.calculate(v, map_to=spec["entity"]).values
    after = vals.mean()
    print(f"  {v:40s} before={before[v]:.3%}  after={after:.3%}")

medicaid_rates = load_take_up_rate("medicaid", 2024)
print("Medicaid takeup rates (state-specific), first 10 states:\n")
for state, rate in sorted(medicaid_rates.items())[:10]:
    print(f"  {state}: {rate:.2%}")

builder = UnifiedMatrixBuilder(
    db_uri=db_uri,
    time_period=2024,
    dataset_path=dataset_path,
)

targets_df, X_sparse, target_names = builder.build_matrix(
    geography,
    sim,
    target_filter={"domain_variables": ["snap"]},
)

print(f"Matrix shape: {X_sparse.shape}")
print(f"Non-zero entries: {X_sparse.nnz:,}")
print(f"Density: {X_sparse.nnz / (X_sparse.shape[0] * X_sparse.shape[1]):.6f}")

print(f"Example household non-zero pattern across clones:\n")
for c in range(N_CLONES):
    col = c * n_records + record_idx
    col_vec = X_sparse[:, col]
    nz_rows = col_vec.nonzero()[0]
    state = int(geography.state_fips[col])
    cd = geography.cd_geoid[col]
    abbr = STATE_CODES.get(state, "??")
    print(f"Clone {c} ({abbr}, CD {cd}): {len(nz_rows)} non-zero rows")
    for r in nz_rows:
        row = targets_df.iloc[r]
        print(
            f"  row {r}: {row['variable']} "
            f"(geo={row['geographic_id']}): "
            f"{X_sparse[r, col]:.2f}"
        )

print("Dimension mismatch:")
print(
    f"  Calibration output: ({n_records} * {N_CLONES},) "
    f"= {n_records * N_CLONES:,} (clone layout)"
)

all_cds = get_all_cds_from_database(db_uri)
n_cds = len(all_cds)
print(
    f"  Stacked builder expects: ({n_cds} * {n_records},) "
    f"= {n_cds * n_records:,} (CD layout)"
)

import os

demo_cds = ["3701", "201"]
n_demo_cds = len(demo_cds)

w = (
    np.random.default_rng(42)
    .binomial(n=1, p=0.01, size=n_demo_cds * n_records)
    .astype(float)
)

# Seed our example household into both CDs
cd_idx_3701 = demo_cds.index("3701")
w[cd_idx_3701 * n_records + record_idx] = 2.5

cd_idx_201 = demo_cds.index("201")
w[cd_idx_201 * n_records + record_idx] = 3.5

output_dir = "calibration_output"
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, "results.h5")

print(
    f"Weight vector: {len(w):,} entries "
    f"({n_demo_cds} CDs x {n_records:,} HH)"
)
print(f"Non-zero weights: {(w > 0).sum()}")
print(
    f"Example HH weight in CD 3701: {w[cd_idx_3701 * n_records + record_idx]}"
)
print(f"Example HH weight in CD 201: {w[cd_idx_201 * n_records + record_idx]}")

create_sparse_cd_stacked_dataset(
    w,
    demo_cds,
    cd_subset=demo_cds,
    dataset_path=dataset_path,
    output_path=output_path,
)

sim_after = Microsimulation(dataset=f"./{output_path}")
hh_after_df = pd.DataFrame(
    sim_after.calculate_dataframe(
        [
            "household_id",
            "congressional_district_geoid",
            "household_weight",
            "state_fips",
            "snap",
        ]
    )
)
print(f"Stacked dataset: {len(hh_after_df)} households\n")

mapping_df = pd.read_csv(
    f"{output_dir}/mappings/results_household_mapping.csv"
)
example_mapping = mapping_df.loc[
    mapping_df.original_household_id == example_hh_id
]
print(f"Example household (original_id={example_hh_id}) " f"in mapping:\n")
print(example_mapping.to_string(index=False))

new_ids = example_mapping.new_household_id
print(f"\nIn stacked dataset:\n")
print(
    hh_after_df.loc[hh_after_df.household_id.isin(new_ids)].to_string(
        index=False
    )
)

import shutil

shutil.rmtree(output_dir)
print(f"Cleaned up {output_dir}/")