cps.py

from contextlib import closing, contextmanager
from importlib.resources import files
from policyengine_core.data import Dataset
from policyengine_us_data.storage import STORAGE_FOLDER, DOCS_FOLDER
import h5py
from policyengine_us_data.datasets.cps.census_cps import (
    CensusCPS,
    CensusCPS_2018,
    CensusCPS_2019,
    CensusCPS_2020,
    CensusCPS_2021,
    CensusCPS_2022,
    CensusCPS_2023,
    CensusCPS_2024,
)
from pandas import DataFrame, Series
import numpy as np
import pandas as pd
import yaml
from typing import Type
from policyengine_us_data.utils.uprating import (
    create_policyengine_uprating_factors_table,
)
from microimpute.models.qrf import QRF
import logging
from policyengine_us_data.parameters import load_take_up_rate
from policyengine_us_data.datasets.cps.takeup import (
    align_reported_ssi_disability,
    prioritize_reported_recipients,
)
from policyengine_us_data.datasets.org import (
    ORG_BOOL_VARIABLES,
    ORG_IMPUTED_VARIABLES,
    build_org_receiver_frame,
    predict_org_features,
)
from policyengine_us_data.utils.downsample import downsample_dataset_arrays
from policyengine_us_data.utils.randomness import seeded_rng
from policyengine_us_data.utils.identification import (
    _store_identification_variables,
)
from policyengine_us_data.datasets.cps.tipped_occupation import (
    derive_treasury_tipped_occupation_code,
    derive_is_tipped_occupation,
)
from policyengine_us_data.utils.takeup import (
    _sum_person_values_to_tax_units,
    _voluntary_filing_age_bin,
    _voluntary_filing_children_bin,
    _voluntary_filing_rate_by_tax_unit,
    _voluntary_filing_wage_income_bin,
    assign_takeup_with_reported_anchors,
    reported_subsidized_marketplace_by_tax_unit,
)
from policyengine_us_data.utils.asset_imputation import (
    SCF_NET_WORTH_TARGET,
    SCF_FINANCIAL_ASSET_POLICY_VARIABLES,
    SCF_HOUSEHOLD_ASSET_POLICY_VARIABLES,
    SCF_NET_WORTH_COMPONENT_VARIABLES,
    add_scf_financial_asset_targets,
    add_scf_household_asset_targets,
    add_scf_net_worth_target,
    add_scf_net_worth_component_targets,
    aggregate_person_values_to_reference_households,
    align_household_values_to_reference_households,
    build_household_vehicle_receiver,
    combine_sipp_and_scf_financial_assets,
    combine_sipp_and_scf_household_assets,
    compute_net_worth_from_components,
    rebalance_scf_net_worth_components,
    require_scf_net_worth_formula_targets,
)
from policyengine_us_data.pipeline_metadata import pipeline_node
from policyengine_us_data.pipeline_schema import PipelineNode
from policyengine_us_data.utils.source_quality import (
    cap_training_sample,
    require_columns_present,
    target_observed_source_masks,
)

ACS_RENT_TARGET_ALLOCATION_COLUMNS = {
    "rent": ["rent_is_allocated"],
    "real_estate_taxes": ["real_estate_taxes_is_allocated"],
}

CURRENT_HEALTH_COVERAGE_REPORTED_VAR_MAP = {
    "reported_has_direct_purchase_health_coverage_at_interview": "NOW_DIR",
    "reported_has_marketplace_health_coverage_at_interview": "NOW_MRK",
    "reported_has_subsidized_marketplace_health_coverage_at_interview": "NOW_MRKS",
    "reported_has_unsubsidized_marketplace_health_coverage_at_interview": "NOW_MRKUN",
    "reported_has_non_marketplace_direct_purchase_health_coverage_at_interview": (
        "NOW_NONM"
    ),
    "reported_has_employer_sponsored_health_coverage_at_interview": "NOW_GRP",
    "reported_has_medicare_health_coverage_at_interview": "NOW_MCARE",
    "reported_has_medicaid_health_coverage_at_interview": "NOW_CAID",
    "reported_has_means_tested_health_coverage_at_interview": "NOW_MCAID",
    "reported_has_chip_health_coverage_at_interview": "NOW_PCHIP",
    "reported_has_other_means_tested_health_coverage_at_interview": "NOW_OTHMT",
    "reported_has_tricare_health_coverage_at_interview": "NOW_MIL",
    "reported_has_champva_health_coverage_at_interview": "NOW_CHAMPVA",
    "reported_has_va_health_coverage_at_interview": "NOW_VACARE",
    "reported_has_indian_health_service_coverage_at_interview": "NOW_IHSFLG",
}

CURRENT_HEALTH_COVERAGE_RULE_INPUT_ALIAS_MAP = {
    "has_marketplace_health_coverage_at_interview": (
        "reported_has_marketplace_health_coverage_at_interview"
    ),
    "has_non_marketplace_direct_purchase_health_coverage_at_interview": (
        "reported_has_non_marketplace_direct_purchase_health_coverage_at_interview"
    ),
    "has_medicaid_health_coverage_at_interview": (
        "reported_has_medicaid_health_coverage_at_interview"
    ),
    "has_other_means_tested_health_coverage_at_interview": (
        "reported_has_other_means_tested_health_coverage_at_interview"
    ),
    "has_tricare_health_coverage_at_interview": (
        "reported_has_tricare_health_coverage_at_interview"
    ),
    "has_champva_health_coverage_at_interview": (
        "reported_has_champva_health_coverage_at_interview"
    ),
    "has_va_health_coverage_at_interview": (
        "reported_has_va_health_coverage_at_interview"
    ),
    "has_indian_health_service_coverage_at_interview": (
        "reported_has_indian_health_service_coverage_at_interview"
    ),
}

CPS_SSI_DISABILITY_DIFFICULTY_COLUMNS = {
    "difficulty_dressing_or_bathing": "PEDISDRS",
    "difficulty_hearing": "PEDISEAR",
    "difficulty_seeing": "PEDISEYE",
    "difficulty_doing_errands": "PEDISOUT",
    "difficulty_walking_or_climbing_stairs": "PEDISPHY",
    "difficulty_remembering_or_making_decisions": "PEDISREM",
}

# Census CPS ASEC 2024 technical documentation, PERRP:
# https://www2.census.gov/programs-surveys/cps/techdocs/cpsmar24.pdf
PERRP_UNMARRIED_PARTNER_OF_HOUSEHOLD_HEAD_CODES = {
    43: "Opposite Sex Unmarried Partner with Relatives",
    44: "Opposite Sex Unmarried Partner without Relatives",
    46: "Same Sex Unmarried Partner with Relatives",
    47: "Same Sex Unmarried Partner without Relatives",
}

ESI_POLICYHOLDER_VARIABLE = (
    "reported_owns_employer_sponsored_health_insurance_at_interview"
)
ESI_SOURCE_COLUMNS = {"NOW_OWNGRP", "NOW_HIPAID", "NOW_GRPFTYP"}


_ESI_PLAN_PRIORS_2024 = {
    # AHRQ MEPS-IC Table IV.A.1 (private sector, 2024). These plan-type
    # averages seed CPS policyholder records; national calibration later
    # aligns the aggregate to the BEA full-economy employer premium total.
    "family": {
        "total_premium": 21_207.52589669509,
        "employee_contribution": 6_490.205059544782,
    },
    "self_only": {
        "total_premium": 8_389.275834815255,
        "employee_contribution": 1_909.5781466113417,
    },
}
_HAS_CURRENT_OWN_ESI = 1
_EMPLOYER_PAYS_ALL = 1
_EMPLOYER_PAYS_SOME = 2
_ESI_FAMILY_PLAN = 1
_ESI_SELF_ONLY_PLAN = 2


def _person_column(person: DataFrame, column: str, default=0) -> np.ndarray:
    if column in person:
        return person[column].to_numpy()
    return np.full(len(person), default)


def impute_employer_sponsored_insurance_premiums(person: DataFrame) -> np.ndarray:
    """Impute annual employer-paid ESI premiums for CPS policyholders."""

    own_esi = _person_column(person, "NOW_OWNGRP").astype(int) == _HAS_CURRENT_OWN_ESI
    premium_status = _person_column(person, "NOW_HIPAID").astype(int)
    plan_type = _person_column(person, "NOW_GRPFTYP").astype(int)
    employee_paid = np.clip(person.PHIP_VAL.to_numpy(dtype=float), 0, None)

    total_premium = np.where(
        plan_type == _ESI_SELF_ONLY_PLAN,
        _ESI_PLAN_PRIORS_2024["self_only"]["total_premium"],
        _ESI_PLAN_PRIORS_2024["family"]["total_premium"],
    )
    average_employee_contribution = np.where(
        plan_type == _ESI_SELF_ONLY_PLAN,
        _ESI_PLAN_PRIORS_2024["self_only"]["employee_contribution"],
        _ESI_PLAN_PRIORS_2024["family"]["employee_contribution"],
    )
    employee_share = np.where(
        employee_paid > 0,
        employee_paid,
        average_employee_contribution,
    )
    employer_paid_when_some = np.clip(
        total_premium - employee_share,
        0,
        total_premium,
    )

    employer_paid = np.where(
        premium_status == _EMPLOYER_PAYS_ALL,
        total_premium,
        np.where(
            premium_status == _EMPLOYER_PAYS_SOME,
            employer_paid_when_some,
            0,
        ),
    )
    valid_owner_with_plan = own_esi & np.isin(
        plan_type,
        [_ESI_FAMILY_PLAN, _ESI_SELF_ONLY_PLAN],
    )
    return np.where(valid_owner_with_plan, employer_paid, 0)


@contextmanager
def _open_dataset_read_only(dataset_source):
    dataset = dataset_source(require=True)
    file_path = getattr(dataset, "file_path", None)

    if file_path is not None:
        with pd.HDFStore(file_path, mode="r") as store:
            yield store
        return

    with closing(dataset.load()) as store:
        yield store


class CPS(Dataset):
    name = "cps"
    label = "CPS"
    raw_cps: Type[CensusCPS] = None
    previous_year_raw_cps: Type[CensusCPS] = None
    data_format = Dataset.ARRAYS
    frac: float | None = 1

    def generate(self):
        """Generates the Current Population Survey dataset for PolicyEngine US microsimulations.
        Technical documentation and codebook here: https://www2.census.gov/programs-surveys/cps/techdocs/cpsmar21.pdf

        Args:
            frac (float, optional): Fraction of the dataset to keep. Defaults to 1. Example: To downsample to 25% of dataset,
                set frac=0.25.
        """
        if self.raw_cps is None:
            raise ValueError(
                f"Cannot generate {self.name}: raw_cps is not defined. "
                "For future years, use PolicyEngine's uprating at simulation time."
            )

        cps = {}

        ENTITIES = ("person", "tax_unit", "family", "spm_unit", "household")
        with _open_dataset_read_only(self.raw_cps) as raw_data:
            person, tax_unit, family, spm_unit, household = [
                raw_data[entity] for entity in ENTITIES
            ]
        _validate_raw_cps_schema(person, tax_unit, self.raw_cps.name)

        logging.info("Adding ID variables")
        add_id_variables(cps, person, tax_unit, family, spm_unit, household)
        logging.info("Adding personal variables")
        add_personal_variables(cps, person)
        logging.info("Adding personal income variables")
        add_personal_income_variables(cps, person, self.raw_cps.time_period)
        logging.info("Adding previous year income variables")
        add_previous_year_income(self, cps)
        logging.info("Adding SSN card type")
        ssn_card_type = add_ssn_card_type(
            cps,
            person,
            spm_unit,
            self.time_period,
            undocumented_target=13e6,
            undocumented_workers_target=8.3e6,
            undocumented_students_target=0.21 * 1.9e6,
        )
        logging.info("Adding taxpayer ID variables")
        _store_identification_variables(
            cps,
            person,
            ssn_card_type,
            self.time_period,
        )
        logging.info("Adding family variables")
        add_spm_variables(self, cps, spm_unit)
        logging.info("Adding household variables")
        add_household_variables(cps, household)
        logging.info("Adding rent")
        add_rent(self, cps, person, household)
        logging.info("Adding tips")
        add_tips(self, cps)
        logging.info("Adding ORG labor-market inputs")
        add_org_labor_market_inputs(cps)
        logging.info("Adding auto loan balance, interest and wealth")
        add_auto_loan_interest_and_net_worth(self, cps)
        logging.info("Added all variables")

        self.save_dataset(cps)
        logging.info("Adding takeup")
        add_takeup(self)
        logging.info("Imputing Marketplace plan benchmark ratio")
        add_marketplace_plan_benchmark_ratio(self)
        logging.info("Deriving other health insurance premiums")
        derive_other_health_insurance_premiums(self)
        logging.info("Downsampling")

        # Downsample
        if self.frac is not None and self.frac < 1.0:
            self.downsample(frac=self.frac)

    @pipeline_node(
        PipelineNode(
            id="downsample",
            label="Downsample CPS",
            node_type="library",
            description="Subsample CPS arrays for released CPS vintages while full variants skip this step.",
            source_file="policyengine_us_data/datasets/cps/cps.py",
            status="current",
            stability="stable",
            pathways=["data_build"],
            validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
        )
    )
    def downsample(self, frac: float) -> None:
        """Subsample the loaded CPS dataset and preserve downsampled arrays.

        Args:
            frac: Fraction of records to retain.
        """
        from policyengine_us import Microsimulation

        original_data: dict = self.load_dataset()
        sim = Microsimulation(dataset=self)
        sim.subsample(frac=frac)
        self.save_dataset(
            downsample_dataset_arrays(
                original_data=original_data,
                sim=sim,
                dataset_name=self.name,
            )
        )


@pipeline_node(
    PipelineNode(
        id="add_rent",
        label="Rent Imputation",
        node_type="library",
        description="Impute rent and real estate taxes using ACS donor data.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="legacy",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_rent(self, cps: h5py.File, person: DataFrame, household: DataFrame):
    cps["tenure_type"] = household.H_TENURE.map(
        {
            0: "NONE",
            1: "OWNED_WITH_MORTGAGE",
            2: "RENTED",
            3: "NONE",
        }
    ).astype("S")
    if self.file_path.exists():
        try:
            with pd.HDFStore(self.file_path, mode="r") as store:
                stale_keys = [
                    key.lstrip("/")
                    for key in store.keys()
                    if key.lstrip("/") not in cps
                ]
                if stale_keys:
                    logging.warning(
                        f"Stale H5 at {self.file_path} has {len(stale_keys)} "
                        f"extra vars before first save: {stale_keys[:5]}"
                    )
        except (BlockingIOError, OSError, ValueError) as error:
            logging.warning(
                "Unable to inspect stale H5 at %s before replacement: %s",
                self.file_path,
                error,
            )
        self.file_path.unlink()
    self.save_dataset(cps)

    from policyengine_us_data.datasets.acs.acs import ACS_2022
    from policyengine_us import Microsimulation

    acs = Microsimulation(dataset=ACS_2022)
    cps_sim = Microsimulation(dataset=self)

    PREDICTORS = [
        "is_household_head",
        "age",
        "is_male",
        "tenure_type",
        "employment_income",
        "self_employment_income",
        "social_security",
        "pension_income",
        "state_code_str",
        "household_size",
    ]
    IMPUTATIONS = ["rent", "real_estate_taxes"]
    train_df = acs.calculate_dataframe(PREDICTORS + IMPUTATIONS, map_to="person")
    # TODO(PolicyEngine/policyengine-core#482): policyengine-core 3.24.0+
    # silently drops user-supplied ETERNITY inputs on dataset reload because
    # _user_input_keys records the user-supplied period instead of the
    # canonicalized ETERNITY key. is_household_head therefore comes back as
    # all False from calculate_dataframe and the household-head filter below
    # produces an empty frame. For ACS we read it directly from the source
    # H5; for CPS we use the in-memory dict (already populated upstream in
    # add_id_variables). Remove both overrides once pyproject.toml's
    # policyengine-core upper bound is lifted.
    required_acs_flags = [
        column
        for columns in ACS_RENT_TARGET_ALLOCATION_COLUMNS.values()
        for column in columns
    ]
    with h5py.File(ACS_2022.file_path, "r") as acs_h5:
        train_df["is_household_head"] = np.asarray(
            acs_h5["is_household_head"], dtype=bool
        )
        require_columns_present(
            acs_h5,
            required_acs_flags,
            source_name="ACS_2022 artifact",
        )
        for flag_columns in ACS_RENT_TARGET_ALLOCATION_COLUMNS.values():
            for flag_column in flag_columns:
                train_df[flag_column] = np.asarray(acs_h5[flag_column], dtype=bool)
    train_df.tenure_type = train_df.tenure_type.map(
        {
            "OWNED_OUTRIGHT": "OWNED_WITH_MORTGAGE",
        },
        na_action="ignore",
    ).fillna(train_df.tenure_type)
    train_df = train_df[train_df.is_household_head].copy()
    inference_df = cps_sim.calculate_dataframe(PREDICTORS, map_to="person")
    inference_df["is_household_head"] = np.asarray(cps["is_household_head"], dtype=bool)
    mask = inference_df.is_household_head.values
    inference_df = inference_df[mask]

    logging.info("Training imputation model for rent and real estate taxes.")
    rent_target_filters = target_observed_source_masks(
        train_df,
        targets=IMPUTATIONS,
        target_allocation_flag_columns=ACS_RENT_TARGET_ALLOCATION_COLUMNS,
    )
    train_df, rent_target_filters = cap_training_sample(
        train_df,
        max_train_samples=10_000,
        seed_name="legacy_acs_rent_training_sample",
        target_filters=rent_target_filters,
    )
    fitted_model = QRF().fit(
        X_train=train_df,
        predictors=PREDICTORS,
        imputed_variables=IMPUTATIONS,
        target_filters=rent_target_filters,
    )
    logging.info("Imputing rent and real estate taxes.")
    imputed_values = fitted_model.predict(X_test=inference_df)
    logging.info("Imputation complete.")
    # ``cps["age"]`` has an integer dtype, so ``np.zeros_like(cps["age"])``
    # would silently truncate the float QRF imputations to int on assignment.
    # Force float so the imputed rent and real-estate-tax values keep
    # full precision (and do not deterministically floor toward zero).
    cps["rent"] = np.zeros(len(cps["age"]), dtype=float)
    cps["rent"][mask] = imputed_values["rent"]
    cps["pre_subsidy_rent"] = cps["rent"]
    cps["real_estate_taxes"] = np.zeros(len(cps["age"]), dtype=float)
    cps["real_estate_taxes"][mask] = imputed_values["real_estate_taxes"]


TEMPORARY_TAKEUP_SOURCE_ANCHORS = ("snap_reported", "ssi_reported")
TEMPORARY_IMPUTATION_SOURCE_VARIABLES = (
    "pension_income",
    "retirement_distributions",
)


def _drop_persisted_dataset_variables(file_path, variable_names):
    with h5py.File(file_path, "a") as dataset_file:
        for variable_name in variable_names:
            if variable_name in dataset_file:
                del dataset_file[variable_name]


@pipeline_node(
    PipelineNode(
        id="add_takeup",
        label="Benefit Takeup",
        node_type="library",
        description="Apply stochastic takeup and reported-anchor alignment for benefit programs.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_takeup(self):
    data = self.load_dataset()

    from policyengine_us import Microsimulation

    baseline = Microsimulation(dataset=self)

    n_persons = len(data["person_id"])
    n_tax_units = len(data["tax_unit_id"])
    n_spm_units = len(data["spm_unit_id"])

    # Load take-up rates
    eitc_rates_by_children = load_take_up_rate("eitc", self.time_period)
    dc_ptc_rate = load_take_up_rate("dc_ptc", self.time_period)
    snap_rate = load_take_up_rate("snap", self.time_period)
    aca_rate = load_take_up_rate("aca", self.time_period)
    medicaid_rates_by_state = load_take_up_rate("medicaid", self.time_period)
    head_start_rate = load_take_up_rate("head_start", self.time_period)
    early_head_start_rate = load_take_up_rate("early_head_start", self.time_period)
    ssi_rate = load_take_up_rate("ssi", self.time_period)
    housing_assistance_rate = load_take_up_rate("housing_assistance", self.time_period)
    voluntary_filing_rates = load_take_up_rate("voluntary_filing", self.time_period)

    # EITC: varies by number of children
    eitc_child_count = baseline.calculate("eitc_child_count").values
    potential_eitc = baseline.calculate("eitc").values
    eitc_takeup_rate = np.array(
        [eitc_rates_by_children.get(min(int(c), 3), 0.85) for c in eitc_child_count]
    )
    rng = seeded_rng("takes_up_eitc")
    data["takes_up_eitc"] = rng.random(n_tax_units) < eitc_takeup_rate

    # DC Property Tax Credit
    rng = seeded_rng("takes_up_dc_ptc")
    data["takes_up_dc_ptc"] = rng.random(n_tax_units) < dc_ptc_rate

    # SNAP: prioritize reported recipients
    rng = seeded_rng("takes_up_snap_if_eligible")
    reported_snap = data["snap_reported"] > 0
    data["takes_up_snap_if_eligible"] = prioritize_reported_recipients(
        reported_snap,
        snap_rate,
        rng.random(n_spm_units),
    )

    # ACA
    rng = seeded_rng("takes_up_aca_if_eligible")
    reported_marketplace_by_tax_unit = reported_subsidized_marketplace_by_tax_unit(
        data["person_tax_unit_id"],
        data["tax_unit_id"],
        data["reported_has_subsidized_marketplace_health_coverage_at_interview"],
    )
    data["takes_up_aca_if_eligible"] = assign_takeup_with_reported_anchors(
        rng.random(n_tax_units),
        aca_rate,
        reported_mask=reported_marketplace_by_tax_unit,
    )

    # Medicaid: state-specific rates
    state_codes = baseline.calculate("state_code_str").values
    hh_ids = data["household_id"]
    person_hh_ids = data["person_household_id"]
    hh_to_state = dict(zip(hh_ids, state_codes))
    person_states = np.array([hh_to_state.get(hh_id, "CA") for hh_id in person_hh_ids])
    medicaid_rate_by_person = np.array(
        [medicaid_rates_by_state.get(s, 0.93) for s in person_states]
    )
    rng = seeded_rng("takes_up_medicaid_if_eligible")
    data["takes_up_medicaid_if_eligible"] = assign_takeup_with_reported_anchors(
        rng.random(n_persons),
        medicaid_rate_by_person,
        reported_mask=data["has_medicaid_health_coverage_at_interview"],
        group_keys=person_states,
    )

    # Head Start
    rng = seeded_rng("takes_up_head_start_if_eligible")
    data["takes_up_head_start_if_eligible"] = rng.random(n_persons) < head_start_rate

    # Early Head Start
    rng = seeded_rng("takes_up_early_head_start_if_eligible")
    data["takes_up_early_head_start_if_eligible"] = (
        rng.random(n_persons) < early_head_start_rate
    )

    # SSI: prioritize reported recipients
    rng = seeded_rng("takes_up_ssi_if_eligible")
    reported_ssi = data["ssi_reported"] > 0
    data["takes_up_ssi_if_eligible"] = prioritize_reported_recipients(
        reported_ssi,
        ssi_rate,
        rng.random(n_persons),
    )

    # TANF
    tanf_rate = load_take_up_rate("tanf", self.time_period)
    rng = seeded_rng("takes_up_tanf_if_eligible")
    data["takes_up_tanf_if_eligible"] = rng.random(n_spm_units) < tanf_rate

    # Housing assistance: prioritize SPM-reported recipients, then fill the
    # remaining draw pool up to the national receipt rate.
    rng = seeded_rng("takes_up_housing_assistance_if_eligible")
    reported_housing_assistance = np.asarray(
        data.get(
            "receives_housing_assistance",
            np.zeros(n_spm_units, dtype=bool),
        ),
        dtype=bool,
    )
    housing_assistance_eligible = baseline.calculate(
        "is_eligible_for_housing_assistance"
    ).values
    data["takes_up_housing_assistance_if_eligible"] = prioritize_reported_recipients(
        reported_housing_assistance,
        housing_assistance_rate,
        rng.random(n_spm_units),
        eligible_mask=housing_assistance_eligible,
    )

    # WIC: resolve draws to bools using category-specific rates
    wic_categories = baseline.calculate("wic_category_str").values
    wic_takeup_rates = load_take_up_rate("wic_takeup", self.time_period)
    wic_takeup_rate_by_person = np.array(
        [wic_takeup_rates.get(c, 0) for c in wic_categories]
    )
    rng = seeded_rng("would_claim_wic")
    data["would_claim_wic"] = rng.random(n_persons) < wic_takeup_rate_by_person

    # WIC nutritional risk — fully resolved
    wic_risk_rates = load_take_up_rate("wic_nutritional_risk", self.time_period)
    wic_risk_rate_by_person = np.array(
        [wic_risk_rates.get(c, 0) for c in wic_categories]
    )
    receives_wic = baseline.calculate("receives_wic").values
    rng = seeded_rng("is_wic_at_nutritional_risk")
    imputed_risk = rng.random(n_persons) < wic_risk_rate_by_person
    data["is_wic_at_nutritional_risk"] = receives_wic | imputed_risk

    # Pregnancy: stochastically assign is_pregnant to women 15-44
    # using CDC/Census-derived state-level pregnancy rates.
    # CPS does not ask about pregnancy; calibration will fine-tune.
    from policyengine_us_data.db.etl_pregnancy import (
        get_state_pregnancy_rates,
    )

    pregnancy_rates = get_state_pregnancy_rates(
        cdc_year=self.time_period,
        acs_year=self.time_period,
    )
    national_rate = 0.041  # fallback
    pregnancy_rate_by_person = np.array(
        [pregnancy_rates.get(s, national_rate) for s in person_states]
    )
    ages = data["age"]
    is_female = data["is_female"]
    is_eligible = is_female & (ages >= 15) & (ages <= 44)
    rng = seeded_rng("is_pregnant")
    data["is_pregnant"] = is_eligible & (
        rng.random(n_persons) < pregnancy_rate_by_person
    )

    # Voluntary tax filing: some tax units file even when not required and not
    # claiming EITC. Assign rates by a simple demographic table that
    # concentrates elective filing among low-wage parents and sharply reduces
    # it among older childless households.
    claims_eitc = data["takes_up_eitc"] & (potential_eitc > 0)
    tax_unit_child_dependents = baseline.calculate("tax_unit_child_dependents").values
    tax_unit_wage_income = _sum_person_values_to_tax_units(
        data["employment_income"],
        data["person_tax_unit_id"],
        data["tax_unit_id"],
    )
    age_head = baseline.calculate("age_head").values
    voluntary_filing_rate = _voluntary_filing_rate_by_tax_unit(
        voluntary_filing_rates,
        _voluntary_filing_children_bin(tax_unit_child_dependents),
        _voluntary_filing_wage_income_bin(tax_unit_wage_income),
        _voluntary_filing_age_bin(age_head),
    )
    rng = seeded_rng("would_file_taxes_voluntarily")
    data["would_file_taxes_voluntarily"] = ~claims_eitc & (
        rng.random(n_tax_units) < voluntary_filing_rate
    )

    # --- SSI: align disability to CPS-reported receipt ---
    # CPS disability flags miss some under-65 SSI recipients, but SSI
    # requires under-65 recipients to be disabled or blind.
    reported_ssi = data["ssi_reported"] > 0
    data["is_disabled"] = align_reported_ssi_disability(
        data["is_disabled"],
        reported_ssi,
        data["age"],
    )

    for source_anchor in TEMPORARY_TAKEUP_SOURCE_ANCHORS:
        data.pop(source_anchor, None)

    self.save_dataset(data)
    _drop_persisted_dataset_variables(
        self.file_path,
        TEMPORARY_TAKEUP_SOURCE_ANCHORS,
    )


def add_marketplace_plan_benchmark_ratio(self):
    """Impute `selected_marketplace_plan_benchmark_ratio` per tax unit.

    PolicyEngine-US's ACA PTC formulas assume the selected Marketplace plan
    costs exactly the Second Lowest Cost Silver Plan (SLCSP). That under- or
    over-states household out-of-pocket premium for the ~50% of Marketplace
    enrollees who select Bronze or Gold / Platinum plans instead.

    For tax units that CPS flags as taking up Marketplace coverage, back out
    the implied plan-to-SLCSP ratio from:

        reported_net_premium   ≈ selected_plan_cost − APTC
        selected_plan_cost      = SLCSP × benchmark_ratio
        → benchmark_ratio       = (reported_net_premium + computed_PTC) / SLCSP

    The CPS private-health-premium field is net of APTC for subsidized
    enrollees, so adding back PolicyEngine's computed PTC recovers the
    sticker cost. Non-Marketplace tax units keep the default 1.0; the
    ratio is only consumed through `selected_marketplace_plan_premium_proxy`,
    which is zero when `takes_up_aca_if_eligible` is false, so the default
    value does not affect their downstream calculations.

    Ratios are clipped to [0.5, 1.5] to handle CPS reporting noise and
    Marketplace-flag false positives. Outliers usually indicate the
    household reported a private or employer premium rather than a true
    Marketplace plan.
    """
    data = self.load_dataset()

    from policyengine_us import Microsimulation

    baseline = Microsimulation(dataset=self)
    period = self.time_period

    slcsp = baseline.calculate("slcsp", map_to="tax_unit", period=period).values
    aca_ptc = baseline.calculate("aca_ptc", map_to="tax_unit", period=period).values
    reported_premium = baseline.calculate(
        "health_insurance_premiums_without_medicare_part_b",
        map_to="tax_unit",
        period=period,
    ).values
    takes_up_aca = baseline.calculate(
        "takes_up_aca_if_eligible",
        map_to="tax_unit",
        period=period,
    ).values.astype(bool)

    data["selected_marketplace_plan_benchmark_ratio"] = (
        compute_marketplace_plan_benchmark_ratio(
            reported_premium=reported_premium,
            aca_ptc=aca_ptc,
            slcsp=slcsp,
            takes_up_aca=takes_up_aca,
        )
    )

    self.save_dataset(data)


OTHER_HEALTH_INSURANCE_PREMIUM_TARGETS = {
    "other_health_insurance_premiums": {
        "reported_variable": "health_insurance_premiums_without_medicare_part_b",
        "modeled_variables": (
            "chip_premium",
            "marketplace_net_premium",
            "medicaid_premium",
        ),
    },
}


def derive_other_health_insurance_premiums(self):
    """Create other premium inputs net of baseline computed premiums.

    The model adds computed premiums back explicitly, so it needs a separate
    other-premium input for the parts of CPS-reported non-Medicare premiums
    not explained by baseline computed Marketplace, CHIP, or Medicaid
    premiums. The original CPS-reported premium inputs remain unchanged as raw
    source fields. The data package requires a policyengine-us release with
    these modeled premium variables, so missing variables fail fast instead of
    silently producing an incomplete decomposition.
    """
    from policyengine_us import Microsimulation

    data = self.load_dataset()
    baseline = Microsimulation(dataset=self)
    tbs = baseline.tax_benefit_system
    period = self.time_period
    changed = False

    for output_variable, config in OTHER_HEALTH_INSURANCE_PREMIUM_TARGETS.items():
        reported_variable = config["reported_variable"]
        premium_variables = config["modeled_variables"]

        if reported_variable not in data:
            continue

        computed_premium = np.zeros(len(data[reported_variable]), dtype=float)
        for variable in premium_variables:
            values = np.asarray(
                baseline.calculate(variable, period=period).values,
                dtype=float,
            )
            computed_premium += _premium_values_to_person(
                data=data,
                source_entity=tbs.variables[variable].entity.key,
                values=values,
            )

        data[output_variable] = compute_other_health_insurance_premiums(
            reported_premium=data[reported_variable],
            baseline_computed_premium=computed_premium,
        )
        logging.info(
            "Created %s from %s by subtracting baseline computed premiums: %s",
            output_variable,
            reported_variable,
            ", ".join(premium_variables),
        )
        changed = True

    if changed:
        self.save_dataset(data)


def compute_other_health_insurance_premiums(
    reported_premium: np.ndarray,
    baseline_computed_premium: np.ndarray,
) -> np.ndarray:
    """Return other premiums after subtracting baseline computed premiums."""
    return np.asarray(reported_premium, dtype=float) - np.asarray(
        baseline_computed_premium, dtype=float
    )


def _premium_values_to_person(
    data: dict,
    source_entity: str,
    values: np.ndarray,
) -> np.ndarray:
    """Map computed premiums to person rows for person-level premium accounting."""
    person_ids = data["person_id"]
    if source_entity == "person":
        if len(values) != len(person_ids):
            raise ValueError(
                "Person-level computed premium length does not match person rows: "
                f"got {len(values)}, expected {len(person_ids)}."
            )
        return values

    entity_id_variable = f"{source_entity}_id"
    person_entity_id_variable = f"person_{source_entity}_id"
    if entity_id_variable not in data or person_entity_id_variable not in data:
        raise ValueError(
            f"Cannot allocate {source_entity}-level premiums to people: missing "
            f"{entity_id_variable} or {person_entity_id_variable}."
        )

    entity_ids = data[entity_id_variable]
    person_entity_ids = data[person_entity_id_variable]
    if len(values) != len(entity_ids):
        raise ValueError(
            f"{source_entity}-level computed premium length does not match "
            f"{source_entity} rows: got {len(values)}, expected {len(entity_ids)}."
        )

    entity_position = {entity_id: index for index, entity_id in enumerate(entity_ids)}
    allocated = np.zeros(len(person_ids), dtype=float)
    seen_entities = set()
    for person_index, entity_id in enumerate(person_entity_ids):
        if entity_id in seen_entities:
            continue
        allocated[person_index] = values[entity_position[entity_id]]
        seen_entities.add(entity_id)
    return allocated


MARKETPLACE_PLAN_BENCHMARK_RATIO_MIN = 0.5
MARKETPLACE_PLAN_BENCHMARK_RATIO_MAX = 1.5


def compute_marketplace_plan_benchmark_ratio(
    reported_premium: np.ndarray,
    aca_ptc: np.ndarray,
    slcsp: np.ndarray,
    takes_up_aca: np.ndarray,
) -> np.ndarray:
    """Back out the Marketplace plan-to-SLCSP ratio per tax unit.

    Returns 1.0 for tax units that don't take up Marketplace coverage or that
    have no SLCSP (zero benchmark). Ratios for Marketplace takers are clipped
    to ``[MARKETPLACE_PLAN_BENCHMARK_RATIO_MIN, MARKETPLACE_PLAN_BENCHMARK_RATIO_MAX]``.

    Args:
        reported_premium: CPS-reported annual private health insurance premium
            (net of APTC for subsidized Marketplace takers), at the tax-unit
            level in USD.
        aca_ptc: PolicyEngine-computed annual advance premium tax credit at
            the tax-unit level in USD.
        slcsp: Second Lowest Cost Silver Plan annual premium at the tax-unit
            level in USD.
        takes_up_aca: Boolean array at the tax-unit level, True when the
            household is modeled as taking up Marketplace coverage.

    Returns:
        Array of benchmark ratios (unitless), same shape as the inputs.
    """
    with np.errstate(divide="ignore", invalid="ignore"):
        raw = (reported_premium + aca_ptc) / np.where(slcsp > 0, slcsp, 1.0)
    clipped = np.clip(
        raw,
        MARKETPLACE_PLAN_BENCHMARK_RATIO_MIN,
        MARKETPLACE_PLAN_BENCHMARK_RATIO_MAX,
    )
    applicable = takes_up_aca.astype(bool) & (slcsp > 0)
    return np.where(applicable, clipped, 1.0)


def uprate_cps_data(data, from_period, to_period):
    uprating = create_policyengine_uprating_factors_table()
    for variable in uprating.index.unique():
        if variable in data:
            current_index = uprating[uprating.index == variable][to_period].values[0]
            start_index = uprating[uprating.index == variable][from_period].values[0]
            growth = current_index / start_index
            data[variable] = data[variable] * growth

    return data


def _validate_raw_cps_schema(
    person: DataFrame,
    tax_unit: DataFrame,
    raw_cps_name: str,
) -> None:
    required_person_columns = {
        "CENSUS_TAX_ID",
        "PERRP",
        *ESI_SOURCE_COLUMNS,
    }
    required_tax_unit_columns = set()

    missing_person = sorted(required_person_columns - set(person.columns))
    missing_tax_unit = sorted(required_tax_unit_columns - set(tax_unit.columns))
    if not missing_person and not missing_tax_unit:
        return

    missing_parts = []
    if missing_person:
        missing_parts.append("person: " + ", ".join(missing_person))
    if missing_tax_unit:
        missing_parts.append("tax_unit: " + ", ".join(missing_tax_unit))

    raise ValueError(
        f"Raw CPS dataset {raw_cps_name} is stale and must be regenerated; "
        f"missing constructed tax-unit columns ({'; '.join(missing_parts)})."
    )


@pipeline_node(
    PipelineNode(
        id="add_id_variables",
        label="Add ID Variables",
        node_type="library",
        description="Create person, household, tax-unit, SPM-unit, family, and marital-unit IDs.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="stable",
        pathways=["data_build"],
        validation_commands=[
            "uv run pytest tests/unit/datasets/test_cps_identification.py"
        ],
    )
)
def add_id_variables(
    cps: h5py.File,
    person: DataFrame,
    tax_unit: DataFrame,
    family: DataFrame,
    spm_unit: DataFrame,
    household: DataFrame,
) -> None:
    """Add basic ID and weight variables.

    Args:
        cps (h5py.File): The CPS dataset file.
        person (DataFrame): The person table of the ASEC.
        tax_unit (DataFrame): The tax unit table created from the person table
            of the ASEC.
        family (DataFrame): The family table of the ASEC.
        spm_unit (DataFrame): The SPM unit table created from the person table
            of the ASEC.
        household (DataFrame): The household table of the ASEC.
    """
    # Add primary and foreign keys
    cps["person_id"] = person.PH_SEQ * 100 + person.P_SEQ
    cps["family_id"] = family.FH_SEQ * 10 + family.FFPOS
    cps["household_id"] = household.H_SEQ
    cps["person_tax_unit_id"] = person.TAX_ID
    cps["person_spm_unit_id"] = person.SPM_ID
    cps["tax_unit_id"] = tax_unit.TAX_ID
    cps["spm_unit_id"] = spm_unit.SPM_ID
    cps["person_household_id"] = person.PH_SEQ
    cps["person_family_id"] = person.PH_SEQ * 10 + person.PF_SEQ
    cps["is_household_head"] = person.P_SEQ == 1
    cps["household_weight"] = household.HSUP_WGT / 1e2

    # Marital units

    marital_unit_id = person.PH_SEQ * 1e6 + np.maximum(person.A_LINENO, person.A_SPOUSE)

    # marital_unit_id is not the household ID, zero padded and followed
    # by the index within household (of each person, or their spouse if
    # one exists earlier in the survey).

    marital_unit_id = Series(marital_unit_id).rank(
        method="dense"
    )  # Simplify to a natural number sequence with repetitions [0, 1, 1, 2, 3, ...]

    cps["person_marital_unit_id"] = marital_unit_id.values
    cps["marital_unit_id"] = marital_unit_id.drop_duplicates().values


@pipeline_node(
    PipelineNode(
        id="add_personal_variables",
        label="Add Personal Variables",
        node_type="library",
        description="Populate CPS personal demographics and occupation-derived inputs.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_personal_variables(cps: h5py.File, person: DataFrame) -> None:
    """Add personal demographic variables.

    Args:
        cps (h5py.File): The CPS dataset file.
        person (DataFrame): The CPS person table.
    """

    # The CPS provides age as follows:
    # 00-79 = 0-79 years of age
    # 80 = 80-84 years of age
    # 85 = 85+ years of age
    # We assign the 80 ages randomly between 80 and 84.
    # to avoid unrealistically bunching at 80. Use a local seeded
    # generator so the draw is reproducible run-to-run and does not
    # depend on whatever seeded the global RNG earlier in the process.
    age_rng = seeded_rng("age_randomization_80_84")
    cps["age"] = np.where(
        person.A_AGE == 80,
        # NB: integers is inclusive of ``low``, exclusive of ``high``.
        age_rng.integers(80, 85, len(person)),
        person.A_AGE,
    )

    # A_SEX is 1 -> male, 2 -> female.
    cps["is_female"] = person.A_SEX == 2
    # "Is...blind or does...have serious difficulty seeing even when Wearing
    #  glasses?" 1 -> Yes
    cps["is_blind"] = person.PEDISEYE == 1
    for variable, cps_column in CPS_SSI_DISABILITY_DIFFICULTY_COLUMNS.items():
        cps[variable] = person[cps_column] == 1
    cps["is_disabled"] = np.column_stack(
        [cps[variable] for variable in CPS_SSI_DISABILITY_DIFFICULTY_COLUMNS]
    ).any(axis=1)

    def children_per_parent(col: str) -> pd.DataFrame:
        """Calculate number of children in the household using parental
            pointers.

        Args:
            col (str): Either PEPAR1 and PEPAR2, which correspond to A_LINENO
            of the person's first and second parent in the household,
            respectively.
        """
        return (
            person[person[col] > 0]
            .groupby(["PH_SEQ", col])
            .size()
            .reset_index()
            .rename(columns={col: "A_LINENO", 0: "children"})
        )

    # Aggregate to parent.
    res = (
        pd.concat([children_per_parent("PEPAR1"), children_per_parent("PEPAR2")])
        .groupby(["PH_SEQ", "A_LINENO"])
        .children.sum()
        .reset_index()
    )
    tmp = person[["PH_SEQ", "A_LINENO"]].merge(
        res, on=["PH_SEQ", "A_LINENO"], how="left"
    )
    cps["own_children_in_household"] = tmp.children.fillna(0)

    for variable, cps_column in CURRENT_HEALTH_COVERAGE_REPORTED_VAR_MAP.items():
        cps[variable] = person[cps_column] == 1

    for (
        variable,
        reported_variable,
    ) in CURRENT_HEALTH_COVERAGE_RULE_INPUT_ALIAS_MAP.items():
        cps[variable] = cps[reported_variable]

    cps["reported_has_private_health_coverage_at_interview"] = person.NOW_PRIV == 1
    cps["reported_has_public_health_coverage_at_interview"] = person.NOW_PUB == 1
    cps["reported_is_insured_at_interview"] = person.NOW_COV == 1
    cps["reported_is_uninsured_at_interview"] = person.NOW_COV != 1

    coverage_families = np.column_stack(
        [
            cps["reported_has_employer_sponsored_health_coverage_at_interview"],
            cps["reported_has_marketplace_health_coverage_at_interview"],
            cps[
                "reported_has_non_marketplace_direct_purchase_health_coverage_at_interview"
            ],
            cps["reported_has_medicare_health_coverage_at_interview"],
            cps["reported_has_means_tested_health_coverage_at_interview"],
            cps["reported_has_tricare_health_coverage_at_interview"],
            cps["reported_has_champva_health_coverage_at_interview"],
            cps["reported_has_va_health_coverage_at_interview"],
            cps["reported_has_indian_health_service_coverage_at_interview"],
        ]
    )
    cps["reported_has_multiple_health_coverage_at_interview"] = (
        coverage_families.sum(axis=1) > 1
    )

    # Legacy aliases retained for compatibility until rules-side names catch up.
    cps["has_marketplace_health_coverage"] = cps[
        "reported_has_marketplace_health_coverage_at_interview"
    ]
    cps["has_esi"] = cps["reported_has_employer_sponsored_health_coverage_at_interview"]

    cps["cps_race"] = person.PRDTRACE
    cps["is_hispanic"] = person.PRDTHSP != 0

    cps["is_surviving_spouse"] = person.A_MARITL == 4
    cps["is_separated"] = person.A_MARITL == 6
    perrp = (
        person.PERRP
        if "PERRP" in person
        else pd.Series(0, index=person.index, dtype=np.int16)
    )
    cps["is_unmarried_partner_of_household_head"] = perrp.isin(
        PERRP_UNMARRIED_PARTNER_OF_HOUSEHOLD_HEAD_CODES.keys()
    )
    # High school or college/university enrollment status.
    if "A_FTPT" in person.columns:
        cps["is_full_time_college_student"] = (person.A_HSCOL == 2) & (
            person.A_FTPT == 1
        )
    else:
        cps["is_full_time_college_student"] = person.A_HSCOL == 2
    cps["detailed_occupation_recode"] = person.POCCU2
    cps["treasury_tipped_occupation_code"] = derive_treasury_tipped_occupation_code(
        person.PEIOOCC
    )
    add_overtime_occupation(cps, person)


def derive_weeks_worked(weeks_worked: Series | np.ndarray) -> Series | np.ndarray:
    return np.clip(weeks_worked, 0, 52)


@pipeline_node(
    PipelineNode(
        id="add_personal_income_variables",
        label="Add Income Variables",
        node_type="library",
        description="Populate CPS income, transfer, retirement, and QBI-related personal inputs.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_personal_income_variables(cps: h5py.File, person: DataFrame, year: int):
    """Add income variables.

    Args:
        cps (h5py.File): The CPS dataset file.
        person (DataFrame): The CPS person table.
        year (int): The CPS year
    """
    # Get income imputation parameters.
    yamlfilename = (
        files("policyengine_us_data")
        / "datasets"
        / "cps"
        / "imputation_parameters.yaml"
    )

    with open(yamlfilename, "r", encoding="utf-8") as yamlfile:
        p = yaml.safe_load(yamlfile)
    assert isinstance(p, dict)

    # Assign CPS variables.
    cps["employment_income"] = person.WSAL_VAL

    cps["weekly_hours_worked"] = person.HRSWK
    cps["hours_worked_last_week"] = person.A_HRS1
    cps["weeks_worked"] = derive_weeks_worked(person.WKSWORK)

    cps["taxable_interest_income"] = person.INT_VAL * (p["taxable_interest_fraction"])
    cps["tax_exempt_interest_income"] = person.INT_VAL * (
        1 - p["taxable_interest_fraction"]
    )
    cps["self_employment_income"] = person.SEMP_VAL
    cps["farm_operations_income"] = person.FRSE_VAL
    cps["qualified_dividend_income"] = (
        person.DIV_VAL * (p["qualified_dividend_fraction"])
    )
    cps["non_qualified_dividend_income"] = person.DIV_VAL * (
        1 - p["qualified_dividend_fraction"]
    )
    cps["rental_income"] = person.RNT_VAL

    # Classify Social Security income using CPS ASEC reason codes
    # (RESNSS1 and RESNSS2). Reason code values:
    #   1 = Retired
    #   2 = Disabled (adult or child)
    #   3 = Widowed
    #   4 = Spouse
    #   5 = Surviving child
    #   6 = Dependent child
    #   7 = On behalf of surviving/dependent/disabled child(ren)
    #   8 = Other
    is_retirement = (person.RESNSS1 == 1) | (person.RESNSS2 == 1)
    is_disability = (person.RESNSS1 == 2) | (person.RESNSS2 == 2)
    is_survivor = np.isin(person.RESNSS1, [3, 5]) | np.isin(person.RESNSS2, [3, 5])
    is_dependent = np.isin(person.RESNSS1, [4, 6, 7]) | np.isin(
        person.RESNSS2, [4, 6, 7]
    )

    # Primary classification: assign full SS_VAL to the highest-
    # priority category when someone has multiple source codes.
    cps["social_security_retirement"] = np.where(is_retirement, person.SS_VAL, 0)
    cps["social_security_disability"] = np.where(
        is_disability & ~is_retirement, person.SS_VAL, 0
    )
    cps["social_security_survivors"] = np.where(
        is_survivor & ~is_retirement & ~is_disability,
        person.SS_VAL,
        0,
    )
    cps["social_security_dependents"] = np.where(
        is_dependent & ~is_retirement & ~is_disability & ~is_survivor,
        person.SS_VAL,
        0,
    )

    # Fallback for records with SS income but no informative source
    # code: use the age-62 heuristic (retirement vs. disability).
    MINIMUM_RETIREMENT_AGE = 62
    unclassified = (
        (person.SS_VAL > 0)
        & ~is_retirement
        & ~is_disability
        & ~is_survivor
        & ~is_dependent
    )
    cps["social_security_retirement"] += np.where(
        unclassified & (person.A_AGE >= MINIMUM_RETIREMENT_AGE),
        person.SS_VAL,
        0,
    )
    cps["social_security_disability"] += np.where(
        unclassified & (person.A_AGE < MINIMUM_RETIREMENT_AGE),
        person.SS_VAL,
        0,
    )
    cps["unemployment_compensation"] = person.UC_VAL
    # Weeks looking for work during the year (Census variable LKWEEKS)
    # LKWEEKS: -1 = NIU (Not In Universe), 0 = not looking, 1-52 = weeks
    weeks_raw = person.LKWEEKS
    cps["weeks_unemployed"] = np.where(weeks_raw == -1, 0, weeks_raw)
    # Add pensions and annuities.
    cps_pensions = person.PNSN_VAL + person.ANN_VAL
    # Assume a constant fraction of pension income is taxable.
    cps["taxable_private_pension_income"] = cps_pensions * p["taxable_pension_fraction"]
    cps["tax_exempt_private_pension_income"] = cps_pensions * (
        1 - p["taxable_pension_fraction"]
    )
    # Retirement account distributions.
    RETIREMENT_CODES = {
        1: "401k",
        2: "403b",
        3: "roth_ira",
        4: "regular_ira",
        5: "keogh",
        6: "sep",  # Simplified Employee Pension plan
        7: "other_type_retirement_account",
    }
    for code, description in RETIREMENT_CODES.items():
        tmp = 0
        # The ASEC splits distributions across four variable pairs.
        for i in ["1", "2", "1_YNG", "2_YNG"]:
            tmp += (person["DST_SC" + i] == code) * person["DST_VAL" + i]
        cps[f"{description}_distributions"] = tmp
    # Allocate retirement distributions by taxability.
    for source_with_taxable_fraction in ["401k", "403b", "sep"]:
        cps[f"taxable_{source_with_taxable_fraction}_distributions"] = (
            cps[f"{source_with_taxable_fraction}_distributions"]
            * p[f"taxable_{source_with_taxable_fraction}_distribution_fraction"]
        )
        cps[f"tax_exempt_{source_with_taxable_fraction}_distributions"] = cps[
            f"{source_with_taxable_fraction}_distributions"
        ] * (1 - p[f"taxable_{source_with_taxable_fraction}_distribution_fraction"])
        del cps[f"{source_with_taxable_fraction}_distributions"]

    # Assume all regular IRA distributions are taxable,
    # and all Roth IRA distributions are not.
    cps["taxable_ira_distributions"] = cps["regular_ira_distributions"]
    cps["tax_exempt_ira_distributions"] = cps["roth_ira_distributions"]
    # Other income (OI_VAL) is a catch-all for all other income sources.
    # The code for alimony income is 20.
    alimony_income = person.OI_OFF == 20
    cps["alimony_income"] = alimony_income * person.OI_VAL
    # The code for strike benefits is 12.
    strike_benefits = person.OI_OFF == 12
    cps["strike_benefits"] = strike_benefits * person.OI_VAL
    cps["miscellaneous_income"] = np.where(
        alimony_income | strike_benefits,
        0,
        person.OI_VAL,
    )
    cps["educational_assistance"] = person.ED_VAL
    cps["financial_assistance"] = person.FIN_VAL
    cps["survivor_benefits"] = person.SRVS_VAL
    cps["child_support_received"] = person.CSP_VAL
    # CPS SSI receipt anchors SSI take-up and disability alignment inside
    # add_takeup; it is dropped before the dataset is saved.
    cps["ssi_reported"] = person.SSI_VAL
    # Allocate CPS RETCB_VAL (a single bundled retirement contribution
    # total) into account-type-specific variables using a proportional
    # split based on administrative data.
    #
    # RETCB_VAL does not distinguish account type (Census only asks
    # "how much did you contribute to retirement accounts?").  The old
    # sequential waterfall gave 401(k) first priority, which consumed
    # nearly all of RETCB_VAL and left IRA contributions at $0.
    #
    # The proportional approach uses BEA/FRED and IRS SOI shares to
    # split contributions into DC (401k) and IRA pools, then splits
    # each pool into traditional/Roth using administrative fractions.
    # See imputation_parameters.yaml for sources.
    from policyengine_us_data.utils.retirement_limits import (
        get_retirement_limits,
    )

    limits = get_retirement_limits(year)
    LIMIT_401K = limits["401k"]
    LIMIT_401K_CATCH_UP = limits["401k_catch_up"]
    LIMIT_IRA = limits["ira"]
    LIMIT_IRA_CATCH_UP = limits["ira_catch_up"]
    CATCH_UP_AGE = 50
    catch_up_eligible = person.A_AGE >= CATCH_UP_AGE
    limit_401k = LIMIT_401K + catch_up_eligible * LIMIT_401K_CATCH_UP
    limit_ira = LIMIT_IRA + catch_up_eligible * LIMIT_IRA_CATCH_UP

    retirement_contributions = person.RETCB_VAL
    has_wages = person.WSAL_VAL > 0
    has_se = person.SEMP_VAL > 0
    has_earned_income = has_wages | has_se

    # 1) Self-employed pension: cap at min(25% of SE income, dollar
    #    limit) so dual-income filers keep a remainder for 401(k)/IRA.
    se_rate = p["se_pension_contribution_rate"]
    se_dollar_cap = p["se_pension_contribution_dollar_limit"][year]
    se_pension_cap = np.minimum(person.SEMP_VAL * se_rate, se_dollar_cap)
    cps["self_employed_pension_contributions"] = np.where(
        has_se,
        np.minimum(retirement_contributions, se_pension_cap),
        0,
    )
    remaining = np.maximum(
        retirement_contributions - cps["self_employed_pension_contributions"],
        0,
    )

    # 2) Split remainder into DC and IRA pools.
    # DC (401k) requires an employer, so gated on has_wages.
    # IRA is available to anyone with earned income (wages or SE).
    dc_share = p["dc_share_of_retirement_contributions"]
    roth_dc_share = p["roth_share_of_dc_contributions"]
    trad_ira_share = p["traditional_share_of_ira_contributions"]

    dc_pool = np.where(has_wages, remaining * dc_share, 0)
    # IRA gets whatever isn't allocated to DC, for anyone with
    # earned income (including SE-only filers).
    ira_pool = np.where(has_earned_income, remaining - dc_pool, 0)

    # DC pool: split into traditional/Roth 401(k), cap at combined
    # 401(k) limit.
    dc_capped = np.minimum(dc_pool, limit_401k)
    cps["traditional_401k_contributions"] = dc_capped * (1 - roth_dc_share)
    cps["roth_401k_contributions"] = dc_capped * roth_dc_share

    # IRA pool: split into traditional/Roth IRA, cap at combined
    # IRA limit.
    ira_capped = np.minimum(ira_pool, limit_ira)
    cps["traditional_ira_contributions"] = ira_capped * trad_ira_share
    cps["roth_ira_contributions"] = ira_capped * (1 - trad_ira_share)
    # Allocate capital gains into long-term and short-term based on aggregate split.
    cps["long_term_capital_gains"] = person.CAP_VAL * (p["long_term_capgain_fraction"])
    cps["short_term_capital_gains"] = person.CAP_VAL * (
        1 - p["long_term_capgain_fraction"]
    )
    cps["receives_wic"] = person.WICYN == 1
    cps["veterans_benefits"] = person.VET_VAL
    cps["workers_compensation"] = person.WC_VAL
    # Disability income has multiple sources and values split across two pairs
    # of variables. Include everything except for worker's compensation
    # (code 1), which is defined as WC_VAL.
    WORKERS_COMP_DISABILITY_CODE = 1
    disability_benefits_1 = person.DIS_VAL1 * (
        person.DIS_SC1 != WORKERS_COMP_DISABILITY_CODE
    )
    disability_benefits_2 = person.DIS_VAL2 * (
        person.DIS_SC2 != WORKERS_COMP_DISABILITY_CODE
    )
    cps["disability_benefits"] = disability_benefits_1 + disability_benefits_2
    # Expenses.
    # "What is the annual amount of child support paid?"
    cps["child_support_expense"] = person.CHSP_VAL
    cps["health_insurance_premiums_without_medicare_part_b"] = person.PHIP_VAL
    cps[ESI_POLICYHOLDER_VARIABLE] = (
        _person_column(person, "NOW_OWNGRP").astype(int) == _HAS_CURRENT_OWN_ESI
    )
    cps["employer_sponsored_insurance_premiums"] = (
        impute_employer_sponsored_insurance_premiums(person)
    )
    cps["over_the_counter_health_expenses"] = person.POTC_VAL
    cps["other_medical_expenses"] = person.PMED_VAL
    cps["medicare_enrolled"] = person.MCARE == 1

    # Get QBI simulation parameters ---
    yamlfilename = (
        files("policyengine_us_data") / "datasets" / "puf" / "qbi_assumptions.yaml"
    )
    with open(yamlfilename, "r", encoding="utf-8") as yamlfile:
        p = yaml.safe_load(yamlfile)
    assert isinstance(p, dict)

    rng = np.random.default_rng(seed=43)
    for var, prob in p["qbi_qualification_probabilities"].items():
        cps[f"{var}_would_be_qualified"] = rng.random(len(person)) < prob


@pipeline_node(
    PipelineNode(
        id="add_spm_variables",
        label="Add SPM Variables",
        node_type="library",
        description="Populate CPS supplemental poverty measure variables.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_spm_variables(self, cps: h5py.File, spm_unit: DataFrame) -> None:
    SPM_RENAMES = dict(
        snap_reported="SPM_SNAPSUB",
        spm_unit_energy_subsidy="SPM_ENGVAL",
        spm_unit_capped_work_childcare_expenses="SPM_CAPWKCCXPNS",
        spm_unit_pre_subsidy_childcare_expenses="SPM_CHILDCAREXPNS",
    )

    for openfisca_variable, asec_variable in SPM_RENAMES.items():
        if asec_variable in spm_unit.columns:
            cps[openfisca_variable] = spm_unit[asec_variable]

    if "SPM_CAPHOUSESUB" in spm_unit.columns:
        cps["receives_housing_assistance"] = spm_unit.SPM_CAPHOUSESUB > 0

    if "SPM_TENMORTSTATUS" in spm_unit.columns:
        tenure_map = {
            1: "OWNER_WITH_MORTGAGE",
            2: "OWNER_WITHOUT_MORTGAGE",
            3: "RENTER",
        }
        cps["spm_unit_tenure_type"] = (
            spm_unit.SPM_TENMORTSTATUS.map(tenure_map).fillna("RENTER").astype("S")
        )


@pipeline_node(
    PipelineNode(
        id="add_household_variables",
        label="Add Household Variables",
        node_type="library",
        description="Populate household geography variables including state, county, and NYC flag.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="stable",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_household_variables(cps: h5py.File, household: DataFrame) -> None:
    """Populate household-level geography variables used by PolicyEngine US.

    Args:
        cps: Output CPS H5 group receiving derived household variables.
        household: Raw CPS household table.
    """
    cps["state_fips"] = household.GESTFIPS
    cps["county_fips"] = household.GTCO
    state_county_fips = cps["state_fips"] * 1e3 + cps["county_fips"]
    # Assign is_nyc here instead of as a variable formula so that it shows up
    # as toggleable in the webapp.
    # List county FIPS codes for each NYC county/borough.
    NYC_COUNTY_FIPS = [
        5,  # Bronx
        47,  # Kings (Brooklyn)
        61,  # New York (Manhattan)
        81,  # Queens
        85,  # Richmond (Staten Island)
    ]
    # Compute NYC by concatenating NY state FIPS with county FIPS.
    # For example, 36061 is Manhattan.
    NYS_FIPS = 36
    nyc_full_county_fips = [
        NYS_FIPS * 1e3 + county_fips for county_fips in NYC_COUNTY_FIPS
    ]
    cps["in_nyc"] = np.isin(state_county_fips, nyc_full_county_fips)


@pipeline_node(
    PipelineNode(
        id="add_previous_year_income",
        label="Previous-Year Income",
        node_type="library",
        description="Link CPS records across adjacent years and populate prior-year income inputs.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_previous_year_income(self, cps: h5py.File) -> None:
    if self.previous_year_raw_cps is None:
        logging.info(
            "No previous year data available for this dataset, skipping previous year income imputation."
        )
        return

    prior_year_income_sentinels = {-1, -9999}

    with (
        _open_dataset_read_only(self.raw_cps) as cps_current_year_data,
        _open_dataset_read_only(self.previous_year_raw_cps) as cps_previous_year_data,
    ):
        cps_previous_year = cps_previous_year_data.person.set_index(
            cps_previous_year_data.person.PERIDNUM
        )
        cps_current_year = cps_current_year_data.person.set_index(
            cps_current_year_data.person.PERIDNUM
        )

        previous_year_data = cps_previous_year[
            ["WSAL_VAL", "SEMP_VAL", "I_ERNVAL", "I_SEVAL"]
        ].rename(
            {
                "WSAL_VAL": "employment_income_last_year",
                "SEMP_VAL": "self_employment_income_last_year",
            },
            axis=1,
        )

        previous_year_data = previous_year_data[
            (previous_year_data.I_ERNVAL == 0) & (previous_year_data.I_SEVAL == 0)
        ]

        previous_year_data.drop(["I_ERNVAL", "I_SEVAL"], axis=1, inplace=True)

        joined_data = cps_current_year.join(previous_year_data)[
            [
                "WSAL_VAL",
                "SEMP_VAL",
                "employment_income_last_year",
                "self_employment_income_last_year",
            ]
        ].rename(
            {
                "WSAL_VAL": "current_year_employment_income",
                "SEMP_VAL": "current_year_self_employment_income",
            },
            axis=1,
        )

    invalid_previous_year_income = joined_data.employment_income_last_year.isin(
        prior_year_income_sentinels
    ) | joined_data.self_employment_income_last_year.isin(prior_year_income_sentinels)
    joined_data.loc[
        invalid_previous_year_income,
        ["employment_income_last_year", "self_employment_income_last_year"],
    ] = np.nan
    joined_data.loc[
        joined_data.current_year_employment_income.isin(prior_year_income_sentinels),
        "current_year_employment_income",
    ] = np.nan
    joined_data.loc[
        joined_data.current_year_self_employment_income.isin(
            prior_year_income_sentinels
        ),
        "current_year_self_employment_income",
    ] = np.nan

    joined_data["previous_year_income_available"] = (
        ~joined_data.employment_income_last_year.isna()
        & ~joined_data.self_employment_income_last_year.isna()
    )
    joined_data["employment_income_last_year"] = joined_data[
        "employment_income_last_year"
    ].fillna(joined_data["current_year_employment_income"])
    joined_data["self_employment_income_last_year"] = joined_data[
        "self_employment_income_last_year"
    ].fillna(joined_data["current_year_self_employment_income"])
    joined_data = joined_data.fillna(0)

    # CPS already ordered by PERIDNUM, so the join wouldn't change the order.
    cps["employment_income_last_year"] = joined_data[
        "employment_income_last_year"
    ].values
    cps["self_employment_income_last_year"] = joined_data[
        "self_employment_income_last_year"
    ].values
    cps["previous_year_income_available"] = joined_data[
        "previous_year_income_available"
    ].values


@pipeline_node(
    PipelineNode(
        id="add_ssn_card_type",
        label="Add SSN Card Type",
        node_type="library",
        description="Classify SSN card type and immigration-related CPS inputs from ASEC conditions.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_ssn_card_type(
    cps: h5py.File,
    person: pd.DataFrame,
    spm_unit: pd.DataFrame,
    time_period: int,
    undocumented_target: float = 13e6,
    undocumented_workers_target: float = 8.3e6,
    undocumented_students_target: float = 0.21 * 1.9e6,
) -> np.ndarray:
    """
    Assign SSN card type using PRCITSHP, employment status, and ASEC-UA conditions.
    Codes:
    - 0: "NONE" - Likely undocumented immigrants
    - 1: "CITIZEN" - US citizens (born or naturalized)
    - 2: "NON_CITIZEN_VALID_EAD" - Non-citizens with work/study authorization
    - 3: "OTHER_NON_CITIZEN" - Non-citizens with indicators of legal status
    """

    # Initialize CSV logging for population tracking
    population_log = []

    def select_random_subset_to_target(
        eligible_ids, current_weighted, target_weighted, random_seed=None
    ):
        """
        Randomly select subset to move current weighted population to target.

        Args:
            eligible_ids: Array of person indices eligible for selection
            current_weighted: Current weighted total
            target_weighted: Target weighted total
            random_seed: Random seed for reproducibility

        Returns:
            Array of selected person indices
        """
        if len(eligible_ids) == 0:
            return np.array([], dtype=int)

        # Calculate how much weighted population needs to be moved
        if current_weighted > target_weighted:
            excess_weighted = current_weighted - target_weighted
            # Calculate fraction to move randomly
            total_reassignable_weight = np.sum(person_weights[eligible_ids])
            share_to_move = excess_weighted / total_reassignable_weight
            share_to_move = min(share_to_move, 1.0)  # Cap at 100%
        else:
            # Calculate how much to move to reach target (for EAD case)
            needed_weighted = current_weighted - target_weighted  # Will be negative
            total_weight = np.sum(person_weights[eligible_ids])
            share_to_move = abs(needed_weighted) / total_weight
            share_to_move = min(share_to_move, 1.0)  # Cap at 100%

        if share_to_move > 0:
            if random_seed is not None:
                # Always use a local Generator so the global np.random
                # state is not clobbered for downstream helpers
                # (imputation, takeup, SIPP choice sampling, etc.).
                rng = np.random.default_rng(seed=random_seed)
                if current_weighted > target_weighted:
                    # Refinement path: draw one uniform per eligible id
                    # and mask those below ``share_to_move``.
                    random_draw = rng.random(len(eligible_ids))
                    assign_mask = random_draw < share_to_move
                    selected = eligible_ids[assign_mask]
                else:
                    # EAD path: sample ``n_to_move`` eligible ids without
                    # replacement using the local generator.
                    n_to_move = int(len(eligible_ids) * share_to_move)
                    selected = rng.choice(eligible_ids, size=n_to_move, replace=False)
            else:
                selected = np.array([], dtype=int)
        else:
            selected = np.array([], dtype=int)

        return selected

    # Get household weights for population calculations
    household_ids = cps["household_id"]
    household_weights = cps["household_weight"]
    person_household_ids = cps["person_household_id"]
    household_to_weight = dict(zip(household_ids, household_weights))
    person_weights = np.array(
        [household_to_weight.get(hh_id, 0) for hh_id in person_household_ids]
    )

    # Initialize all persons as code 0
    ssn_card_type = np.full(len(person), 0)
    initial_population = np.sum(person_weights[ssn_card_type == 0])
    print(f"Step 0 - Initial: Code 0 people: {initial_population:,.0f}")
    population_log.append(
        {
            "step": "Step 0 - Initial",
            "description": "Code 0 people",
            "population": initial_population,
        }
    )

    # ============================================================================
    # PRIMARY CLASSIFICATIONS
    # ============================================================================

    # Code 1: All US Citizens (naturalized and born)
    citizens_mask = np.isin(person.PRCITSHP, [1, 2, 3, 4])
    ssn_card_type[citizens_mask] = 1
    noncitizens = person.PRCITSHP == 5
    citizens_moved = np.sum(person_weights[citizens_mask])
    print(f"Step 1 - Citizens: Moved {citizens_moved:,.0f} people to Code 1")
    population_log.append(
        {
            "step": "Step 1 - Citizens",
            "description": "Moved to Code 1",
            "population": citizens_moved,
        }
    )

    # ============================================================================
    # ASEC UNDOCUMENTED ALGORITHM CONDITIONS
    # Remove individuals with indicators of legal status from code 0 pool
    # ============================================================================

    # paper source: https://papers.ssrn.com/sol3/papers.cfm?abstract_id=4662801
    # Helper mask: Only apply conditions to non-citizens without clear authorization
    potentially_undocumented = ~np.isin(ssn_card_type, [1, 2])

    current_code_0 = np.sum(person_weights[ssn_card_type == 0])
    print(f"\nASEC Conditions - Current Code 0 people: {current_code_0:,.0f}")
    population_log.append(
        {
            "step": "ASEC Conditions",
            "description": "Current Code 0 people",
            "population": current_code_0,
        }
    )

    # CONDITION 1: Pre-1982 Arrivals (IRCA Amnesty Eligible)
    # PEINUSYR values indicating arrival before 1982:
    # 01 = Before 1950
    # 02 = 1950–1959
    # 03 = 1960–1964
    # 04 = 1965–1969
    # 05 = 1970–1974
    # 06 = 1975–1979
    # 07 = 1980–1981
    arrived_before_1982 = np.isin(person.PEINUSYR, [1, 2, 3, 4, 5, 6, 7])
    condition_1_mask = potentially_undocumented & arrived_before_1982
    condition_1_count = np.sum(person_weights[condition_1_mask])
    print(
        f"Condition 1 - Pre-1982 arrivals: {condition_1_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 1",
            "description": "Pre-1982 arrivals qualify for Code 3",
            "population": condition_1_count,
        }
    )

    # CONDITION 2: Eligible Naturalized Citizens
    is_naturalized = person.PRCITSHP == 4
    is_adult = person.A_AGE >= 18
    # 5+ years in US (codes 8-26: 1982-2019)
    has_five_plus_years = np.isin(person.PEINUSYR, list(range(8, 27)))
    # 3+ years in US + married (codes 8-27: 1982-2021)
    has_three_plus_years = np.isin(person.PEINUSYR, list(range(8, 28)))
    is_married = person.A_MARITL.isin([1, 2]) & (person.A_SPOUSE > 0)
    eligible_naturalized = (
        is_naturalized
        & is_adult
        & (has_five_plus_years | (has_three_plus_years & is_married))
    )
    condition_2_mask = potentially_undocumented & eligible_naturalized
    condition_2_count = np.sum(person_weights[condition_2_mask])
    print(
        f"Condition 2 - Eligible naturalized citizens: {condition_2_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 2",
            "description": "Eligible naturalized citizens qualify for Code 3",
            "population": condition_2_count,
        }
    )

    # CONDITION 3: Medicare Recipients
    has_medicare = person.MCARE == 1
    condition_3_mask = potentially_undocumented & has_medicare
    condition_3_count = np.sum(person_weights[condition_3_mask])
    print(
        f"Condition 3 - Medicare recipients: {condition_3_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 3",
            "description": "Medicare recipients qualify for Code 3",
            "population": condition_3_count,
        }
    )

    # CONDITION 4: Federal Retirement Benefits
    has_federal_pension = np.isin(person.PEN_SC1, [3]) | np.isin(
        person.PEN_SC2, [3]
    )  # Federal government pension
    condition_4_mask = potentially_undocumented & has_federal_pension
    condition_4_count = np.sum(person_weights[condition_4_mask])
    print(
        f"Condition 4 - Federal retirement benefits: {condition_4_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 4",
            "description": "Federal retirement benefits qualify for Code 3",
            "population": condition_4_count,
        }
    )

    # CONDITION 5: Social Security Disability
    has_ss_disability = np.isin(person.RESNSS1, [2]) | np.isin(
        person.RESNSS2, [2]
    )  # Disabled (adult or child)
    condition_5_mask = potentially_undocumented & has_ss_disability
    condition_5_count = np.sum(person_weights[condition_5_mask])
    print(
        f"Condition 5 - Social Security disability: {condition_5_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 5",
            "description": "Social Security disability qualify for Code 3",
            "population": condition_5_count,
        }
    )

    # CONDITION 6: Indian Health Service Coverage
    has_ihs = person.IHSFLG == 1
    condition_6_mask = potentially_undocumented & has_ihs
    condition_6_count = np.sum(person_weights[condition_6_mask])
    print(
        f"Condition 6 - Indian Health Service coverage: {condition_6_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 6",
            "description": "Indian Health Service coverage qualify for Code 3",
            "population": condition_6_count,
        }
    )

    # CONDITION 7: Medicaid Recipients (simplified - no state adjustments)
    has_medicaid = person.CAID == 1
    condition_7_mask = potentially_undocumented & has_medicaid
    condition_7_count = np.sum(person_weights[condition_7_mask])
    print(
        f"Condition 7 - Medicaid recipients: {condition_7_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 7",
            "description": "Medicaid recipients qualify for Code 3",
            "population": condition_7_count,
        }
    )

    # CONDITION 8: CHAMPVA Recipients
    has_champva = person.CHAMPVA == 1
    condition_8_mask = potentially_undocumented & has_champva
    condition_8_count = np.sum(person_weights[condition_8_mask])
    print(
        f"Condition 8 - CHAMPVA recipients: {condition_8_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 8",
            "description": "CHAMPVA recipients qualify for Code 3",
            "population": condition_8_count,
        }
    )

    # CONDITION 9: Military Health Insurance
    has_military_insurance = person.MIL == 1
    condition_9_mask = potentially_undocumented & has_military_insurance
    condition_9_count = np.sum(person_weights[condition_9_mask])
    print(
        f"Condition 9 - Military health insurance: {condition_9_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 9",
            "description": "Military health insurance qualify for Code 3",
            "population": condition_9_count,
        }
    )

    # CONDITION 10: Government Employees
    is_government_worker = np.isin(person.PEIO1COW, [1, 2, 3])  # Fed/state/local gov
    is_military_occupation = person.A_MJOCC == 11  # Military occupation
    is_government_employee = is_government_worker | is_military_occupation
    condition_10_mask = potentially_undocumented & is_government_employee
    condition_10_count = np.sum(person_weights[condition_10_mask])
    print(
        f"Condition 10 - Government employees: {condition_10_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 10",
            "description": "Government employees qualify for Code 3",
            "population": condition_10_count,
        }
    )

    # CONDITION 11: Social Security Recipients
    has_social_security = person.SS_YN == 1
    condition_11_mask = potentially_undocumented & has_social_security
    condition_11_count = np.sum(person_weights[condition_11_mask])
    print(
        f"Condition 11 - Social Security recipients: {condition_11_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 11",
            "description": "Social Security recipients qualify for Code 3",
            "population": condition_11_count,
        }
    )

    # CONDITION 12: Housing Assistance
    spm_housing_map = dict(zip(spm_unit.SPM_ID, spm_unit.SPM_CAPHOUSESUB))
    has_housing_assistance = person.SPM_ID.map(spm_housing_map).fillna(0) > 0
    condition_12_mask = potentially_undocumented & has_housing_assistance
    condition_12_count = np.sum(person_weights[condition_12_mask])
    print(
        f"Condition 12 - Housing assistance: {condition_12_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 12",
            "description": "Housing assistance qualify for Code 3",
            "population": condition_12_count,
        }
    )

    # CONDITION 13: Veterans/Military Personnel
    is_veteran = person.PEAFEVER == 1
    is_current_military = person.A_MJOCC == 11
    is_military_connected = is_veteran | is_current_military
    condition_13_mask = potentially_undocumented & is_military_connected
    condition_13_count = np.sum(person_weights[condition_13_mask])
    print(
        f"Condition 13 - Veterans/Military personnel: {condition_13_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 13",
            "description": "Veterans/Military personnel qualify for Code 3",
            "population": condition_13_count,
        }
    )

    # CONDITION 14: SSI Recipients (simplified - assumes all SSI is for recipient)
    has_ssi = person.SSI_YN == 1
    condition_14_mask = potentially_undocumented & has_ssi
    condition_14_count = np.sum(person_weights[condition_14_mask])
    print(
        f"Condition 14 - SSI recipients: {condition_14_count:,.0f} people qualify for Code 3"
    )
    population_log.append(
        {
            "step": "Condition 14",
            "description": "SSI recipients qualify for Code 3",
            "population": condition_14_count,
        }
    )

    # ============================================================================
    # CONSOLIDATED ASSIGNMENT OF ASSUMED DOCUMENTED STATUS
    # ============================================================================

    # Combine all conditions that indicate legal status
    assumed_documented = (
        arrived_before_1982
        | eligible_naturalized
        | has_medicare
        | has_federal_pension
        | has_ss_disability
        | has_ihs
        | has_medicaid
        | has_champva
        | has_military_insurance
        | is_government_employee
        | has_social_security
        | has_housing_assistance
        | is_military_connected
        | has_ssi
    )

    # Apply single assignment for all conditions
    ssn_card_type[potentially_undocumented & assumed_documented] = 3
    # print(f"Step 2 - Documented indicators: Moved {np.sum(person_weights[potentially_undocumented & assumed_documented]):,.0f} people from Code 0 to Code 3")

    # Calculate undocumented workers and students after ASEC conditions
    undocumented_workers_mask = (
        (ssn_card_type == 0)
        & noncitizens
        & ((person.WSAL_VAL > 0) | (person.SEMP_VAL > 0))
    )
    undocumented_students_mask = (
        (ssn_card_type == 0) & noncitizens & (person.A_HSCOL == 2)
    )
    undocumented_workers_count = np.sum(person_weights[undocumented_workers_mask])
    undocumented_students_count = np.sum(person_weights[undocumented_students_mask])

    after_conditions_code_0 = np.sum(person_weights[ssn_card_type == 0])
    print(f"After conditions - Code 0 people: {after_conditions_code_0:,.0f}")
    print(
        f"  - Undocumented workers before adjustment: {undocumented_workers_count:,.0f} (target: {undocumented_workers_target:,.0f})"
    )
    print(
        f"  - Undocumented students before adjustment: {undocumented_students_count:,.0f} (target: {undocumented_students_target:,.0f})"
    )

    population_log.append(
        {
            "step": "After conditions",
            "description": "Code 0 people",
            "population": after_conditions_code_0,
        }
    )
    population_log.append(
        {
            "step": "Before adjustment",
            "description": "Undocumented workers",
            "population": undocumented_workers_count,
        }
    )
    population_log.append(
        {
            "step": "Target",
            "description": "Undocumented workers target",
            "population": undocumented_workers_target,
        }
    )
    population_log.append(
        {
            "step": "Before adjustment",
            "description": "Undocumented students",
            "population": undocumented_students_count,
        }
    )
    population_log.append(
        {
            "step": "Target",
            "description": "Undocumented students target",
            "population": undocumented_students_target,
        }
    )

    # ============================================================================
    # CODE 2 NON-CITIZEN WITH WORK/STUDY AUTHORIZATION
    # ============================================================================

    # Code 2: Non-citizens with work/study authorization (likely valid EAD)
    # Only consider people still in Code 0 (undocumented) after ASEC conditions
    worker_mask = (
        (ssn_card_type != 3)
        & noncitizens
        & ((person.WSAL_VAL > 0) | (person.SEMP_VAL > 0))
    )
    student_mask = (ssn_card_type != 3) & noncitizens & (person.A_HSCOL == 2)

    # Calculate target-driven worker assignment
    # Target: 8.3 million undocumented workers (from Pew Research)
    # https://www.pewresearch.org/short-reads/2024/07/22/what-we-know-about-unauthorized-immigrants-living-in-the-us/

    # Get worker IDs
    worker_ids = person[worker_mask].index

    # Use function to select workers for EAD
    total_weighted_workers = np.sum(person_weights[worker_ids])
    selected_workers = select_random_subset_to_target(
        worker_ids,
        total_weighted_workers,
        undocumented_workers_target,
        random_seed=0,
    )

    # Calculate target-driven student assignment
    # Target: 21% of 1.9 million = ~399k undocumented students (from Higher Ed Immigration Portal)
    # https://www.higheredimmigrationportal.org/research/immigrant-origin-students-in-u-s-higher-education-updated-august-2024/

    student_ids = person[student_mask].index

    # Use function to select students for EAD
    total_weighted_students = np.sum(person_weights[student_ids])
    selected_students = select_random_subset_to_target(
        student_ids,
        total_weighted_students,
        undocumented_students_target,
        random_seed=1,
    )

    # Assign code 2
    ssn_card_type[selected_workers] = 2
    ssn_card_type[selected_students] = 2
    ead_workers_moved = np.sum(person_weights[selected_workers])
    ead_students_moved = np.sum(person_weights[selected_students])
    after_ead_code_0 = np.sum(person_weights[ssn_card_type == 0])

    print(
        f"Step 3 - EAD workers: Moved {ead_workers_moved:,.0f} people from Code 0 to Code 2"
    )
    print(
        f"Step 4 - EAD students: Moved {ead_students_moved:,.0f} people from Code 0 to Code 2"
    )
    print(f"After EAD assignment - Code 0 people: {after_ead_code_0:,.0f}")

    population_log.append(
        {
            "step": "Step 3 - EAD workers",
            "description": "Moved from Code 0 to Code 2",
            "population": ead_workers_moved,
        }
    )
    population_log.append(
        {
            "step": "Step 4 - EAD students",
            "description": "Moved from Code 0 to Code 2",
            "population": ead_students_moved,
        }
    )
    population_log.append(
        {
            "step": "After EAD assignment",
            "description": "Code 0 people",
            "population": after_ead_code_0,
        }
    )

    # ============================================================================
    # PROBABILISTIC FAMILY CORRELATION ADJUSTMENT
    # ============================================================================

    # Probabilistic family correlation: Only move code 3 household members to code 0
    # if needed to hit the undocumented target. This preserves mixed-status families
    # (citizens living with undocumented) while still achieving target-driven correlation.

    # Use existing household data
    person_household_ids = cps["person_household_id"]

    # Track before state
    code_0_before = np.sum(person_weights[ssn_card_type == 0])

    # Calculate how many more undocumented people we need to hit target
    current_undocumented = code_0_before
    undocumented_needed = max(0, undocumented_target - current_undocumented)

    print(
        f"Current undocumented: {current_undocumented:,.0f}, Target: {undocumented_target:,.0f}"
    )
    print(f"Additional undocumented needed: {undocumented_needed:,.0f}")

    if undocumented_needed > 0:
        # Identify households with mixed status (code 0 + code 3 members).
        #
        # The previous implementation looped over every unique
        # household and masked against the full ``person_household_ids``
        # array on each iteration — O(uniq_households × persons) for
        # CPS 2024 that is ~3×10¹⁰ comparisons (~100k households,
        # ~300k persons).
        #
        # Vectorize via a single groupby: for each household, flag
        # whether any member is code 0 and whether any member is
        # code 3; a mixed household is one where both flags are True.
        # Then collect every code-3 person whose household is mixed.
        codes_series = pd.Series(ssn_card_type, name="ssn_card_type")
        per_household = pd.DataFrame(
            {
                "household_id": person_household_ids,
                "is_code_0": codes_series == 0,
                "is_code_3": codes_series == 3,
            }
        )
        household_flags = per_household.groupby("household_id").agg(
            has_code_0=("is_code_0", "any"),
            has_code_3=("is_code_3", "any"),
        )
        mixed_household_ids = household_flags.index[
            household_flags["has_code_0"] & household_flags["has_code_3"]
        ]
        in_mixed_household = np.isin(person_household_ids, mixed_household_ids)
        mixed_household_candidates = np.where(
            in_mixed_household & (ssn_card_type == 3)
        )[0]

        # Randomly select from eligible code 3 members in mixed households to hit target
        if len(mixed_household_candidates) > 0:
            mixed_household_candidates = np.array(mixed_household_candidates)

            # Use probabilistic selection to hit target
            selected_indices = select_random_subset_to_target(
                mixed_household_candidates,
                current_undocumented,
                undocumented_target,
                random_seed=100,  # Different seed for family correlation
            )

            if len(selected_indices) > 0:
                ssn_card_type[selected_indices] = 0
                print(
                    f"Selected {len(selected_indices)} people from {len(mixed_household_candidates)} candidates in mixed households"
                )
            else:
                print("No additional family members selected (target already reached)")
        else:
            print("No mixed-status households found for family correlation")
    else:
        print("No additional undocumented people needed - target already reached")

    # Calculate the weighted impact
    code_0_after = np.sum(person_weights[ssn_card_type == 0])
    weighted_change = code_0_after - code_0_before

    print(
        f"Step 5 - Probabilistic family correlation: Changed {weighted_change:,.0f} people from Code 3 to Code 0"
    )
    print(f"After family correlation - Code 0 people: {code_0_after:,.0f}")

    population_log.append(
        {
            "step": "Step 5 - Family correlation",
            "description": "Changed from Code 3 to Code 0",
            "population": weighted_change,
        }
    )
    population_log.append(
        {
            "step": "After family correlation",
            "description": "Code 0 people",
            "population": code_0_after,
        }
    )

    def get_arrival_year_midpoint(peinusyr):
        """
        Map PEINUSYR codes to arrival year midpoints.
        Returns a numpy array of estimated arrival years.
        """
        arrival_year_map = {
            1: 1945,  # Before 1950
            2: 1955,  # 1950-1959
            3: 1962,  # 1960-1964
            4: 1967,  # 1965-1969
            5: 1972,  # 1970-1974
            6: 1977,  # 1975-1979
            7: 1981,  # 1980-1981
            8: 1983,  # 1982-1983
            9: 1985,  # 1984-1985
            10: 1987,  # 1986-1987
            11: 1989,  # 1988-1989
            12: 1991,  # 1990-1991
            13: 1993,  # 1992-1993
            14: 1995,  # 1994-1995
            15: 1997,  # 1996-1997
            16: 1999,  # 1998-1999
            17: 2001,  # 2000-2001
            18: 2003,  # 2002-2003
            19: 2005,  # 2004-2005
            20: 2007,  # 2006-2007
            21: 2009,  # 2008-2009
            22: 2011,  # 2010-2011
            23: 2013,  # 2012-2013
            24: 2015,  # 2014-2015
            25: 2017,  # 2016-2017
            26: 2019,  # 2018-2019
            27: 2021,  # 2020-2021
            28: 2023,  # 2022-2023
            29: 2024,  # 2024-2025
        }

        # Vectorized mapping with default value of 2024
        return np.vectorize(arrival_year_map.get)(peinusyr, 2024)

    # NEW IMMIGRATION-STATUS TAGS FOR OBFBA
    birth = person.PENATVTY  # two-digit nativity flag

    # Calculate arrival years once for all logic
    arrival_years = get_arrival_year_midpoint(person.PEINUSYR)
    years_in_us = time_period - arrival_years
    age_at_entry = np.maximum(0, person.A_AGE - years_in_us)

    # start every non-citizen as LPR so no UNSET survives
    immigration_status = np.full(len(person), "LEGAL_PERMANENT_RESIDENT", dtype="U32")

    # Set citizens (SSN card type 1) to CITIZEN status
    immigration_status[ssn_card_type == 1] = "CITIZEN"

    # 1. Undocumented: SSN card type 0 who arrived 1982 or later
    arrived_before_1982 = np.isin(person.PEINUSYR, [1, 2, 3, 4, 5, 6, 7])
    undoc_mask = (ssn_card_type == 0) & (~arrived_before_1982)
    immigration_status[undoc_mask] = "UNDOCUMENTED"

    COUNTRY_CODES = {
        "COFA": {511, 512},  # Micronesia, Marshall Islands. Palau not listed
        "CUBAN_HAITIAN": {327, 332},
    }

    mask = (ssn_card_type != 0) & np.isin(birth, list(COUNTRY_CODES["COFA"]))
    immigration_status[mask] = "LEGAL_PERMANENT_RESIDENT"

    # 3. Cuban / Haitian entrants (created by Congress in 1980)
    # Only those who arrived after 1980
    CUBAN_HAITIAN_ARRIVAL_CUTOFF = 1980
    cuban_haitian_mask = (
        (ssn_card_type != 0)
        & np.isin(birth, list(COUNTRY_CODES["CUBAN_HAITIAN"]))
        & (arrival_years >= CUBAN_HAITIAN_ARRIVAL_CUTOFF)
    )
    immigration_status[cuban_haitian_mask] = "CUBAN_HAITIAN_ENTRANT"

    # DACA eligibility constants
    # Source: https://www.uscis.gov/humanitarian/consideration-of-deferred-action-for-childhood-arrivals-daca
    DACA_LATEST_ARRIVAL_YEAR = 2007  # Must have arrived before June 15, 2007
    DACA_MAX_AGE_AT_ENTRY = 16  # Must have arrived before 16th birthday
    DACA_MIN_CURRENT_AGE = 15  # Must be at least 15 years old to apply

    # 4. DACA
    daca_mask = (
        (ssn_card_type == 2)  # Temporary/unauthorized status
        & (arrival_years <= DACA_LATEST_ARRIVAL_YEAR)
        & (age_at_entry < DACA_MAX_AGE_AT_ENTRY)
        & (person.A_AGE >= DACA_MIN_CURRENT_AGE)
    )
    immigration_status[daca_mask] = "DACA"

    # 5. Recent humanitarian parole/asylee/refugee (Code 3, ≤ 5 yrs)
    recent_refugee_mask = (ssn_card_type == 3) & (years_in_us <= 5)
    immigration_status[recent_refugee_mask] = "REFUGEE"

    # 6. Temp non-qualified (Code 2 not caught by DACA rule)
    mask = (ssn_card_type == 2) & (immigration_status == "LEGAL_PERMANENT_RESIDENT")
    immigration_status[mask] = "TPS"

    # Final write (all values now in ImmigrationStatus Enum)
    # Save as immigration_status_str since that's what PolicyEngine expects
    cps["immigration_status_str"] = immigration_status.astype("S")
    # Final population summary
    print("\nFinal populations:")
    code_to_str = {
        0: "NONE",  # Likely undocumented immigrants
        1: "CITIZEN",  # US citizens
        2: "NON_CITIZEN_VALID_EAD",  # Non-citizens with work/study authorization
        3: "OTHER_NON_CITIZEN",  # Non-citizens with indicators of legal status
    }
    for code, label in code_to_str.items():
        pop = np.sum(person_weights[ssn_card_type == code])
        print(f"  Code {code} ({label}): {pop:,.0f}")
        population_log.append(
            {
                "step": "Final",
                "description": f"Code {code} ({label})",
                "population": pop,
            }
        )

    final_undocumented = np.sum(person_weights[ssn_card_type == 0])
    print(
        f"Total undocumented (Code 0): {final_undocumented:,.0f} (target: {undocumented_target:,.0f})"
    )
    population_log.append(
        {
            "step": "Final",
            "description": "Total undocumented (Code 0)",
            "population": final_undocumented,
        }
    )
    population_log.append(
        {
            "step": "Final",
            "description": "Undocumented target",
            "population": undocumented_target,
        }
    )

    # Save population log to CSV
    log_df = pd.DataFrame(population_log)
    csv_path = DOCS_FOLDER / "asec_population_log.csv"
    DOCS_FOLDER.mkdir(exist_ok=True)
    log_df.to_csv(csv_path, index=False)
    print(f"Population log saved to: {csv_path}")

    # Update documentation with actual numbers
    _update_documentation_with_numbers(log_df, DOCS_FOLDER)

    return ssn_card_type


def _update_documentation_with_numbers(log_df, docs_dir):
    """Update the documentation file with actual population numbers from CSV"""
    doc_path = docs_dir / "SSN_statuses_imputation.ipynb"

    if not doc_path.exists():
        print(f"Documentation file not found at: {doc_path}")
        return

    # Create mapping of step/description to population for easy lookup
    data_map = {}
    for _, row in log_df.iterrows():
        key = (row["step"], row["description"])
        data_map[key] = row["population"]

    # Read the documentation file
    with open(doc_path, "r", encoding="utf-8") as f:
        content = f.read()

    # Define replacements based on our logging structure
    replacements = {
        "- **Step 0 - Initial**: Code 0 people = *[Run cps.py to populate]*": lambda: (
            f"- **Step 0 - Initial**: Code 0 people = {data_map.get(('Step 0 - Initial', 'Code 0 people'), 0):,.0f}"
        ),
        "- **Step 1 - Citizens**: Moved to Code 1 = *[Run cps.py to populate]*": lambda: (
            f"- **Step 1 - Citizens**: Moved to Code 1 = {data_map.get(('Step 1 - Citizens', 'Moved to Code 1'), 0):,.0f}"
        ),
        "- **ASEC Conditions**: Current Code 0 people = *[Run cps.py to populate]*": lambda: (
            f"- **ASEC Conditions**: Current Code 0 people = {data_map.get(('ASEC Conditions', 'Current Code 0 people'), 0):,.0f}"
        ),
        "- **After conditions**: Code 0 people = *[Run cps.py to populate]*": lambda: (
            f"- **After conditions**: Code 0 people = {data_map.get(('After conditions', 'Code 0 people'), 0):,.0f}"
        ),
        "- **Before adjustment**: Undocumented workers = *[Run cps.py to populate]*": lambda: (
            f"- **Before adjustment**: Undocumented workers = {data_map.get(('Before adjustment', 'Undocumented workers'), 0):,.0f}"
        ),
        "- **Target**: Undocumented workers target = *[Run cps.py to populate]*": lambda: (
            f"- **Target**: Undocumented workers target = {data_map.get(('Target', 'Undocumented workers target'), 0):,.0f}"
        ),
        "- **Before adjustment**: Undocumented students = *[Run cps.py to populate]*": lambda: (
            f"- **Before adjustment**: Undocumented students = {data_map.get(('Before adjustment', 'Undocumented students'), 0):,.0f}"
        ),
        "- **Target**: Undocumented students target = *[Run cps.py to populate]*": lambda: (
            f"- **Target**: Undocumented students target = {data_map.get(('Target', 'Undocumented students target'), 0):,.0f}"
        ),
        "- **Step 3 - EAD workers**: Moved from Code 0 to Code 2 = *[Run cps.py to populate]*": lambda: (
            f"- **Step 3 - EAD workers**: Moved from Code 0 to Code 2 = {data_map.get(('Step 3 - EAD workers', 'Moved from Code 0 to Code 2'), 0):,.0f}"
        ),
        "- **Step 4 - EAD students**: Moved from Code 0 to Code 2 = *[Run cps.py to populate]*": lambda: (
            f"- **Step 4 - EAD students**: Moved from Code 0 to Code 2 = {data_map.get(('Step 4 - EAD students', 'Moved from Code 0 to Code 2'), 0):,.0f}"
        ),
        "- **After EAD assignment**: Code 0 people = *[Run cps.py to populate]*": lambda: (
            f"- **After EAD assignment**: Code 0 people = {data_map.get(('After EAD assignment', 'Code 0 people'), 0):,.0f}"
        ),
        "- **Step 5 - Family correlation**: Changed from Code 3 to Code 0 = *[Run cps.py to populate]*": lambda: (
            f"- **Step 5 - Family correlation**: Changed from Code 3 to Code 0 = {data_map.get(('Step 5 - Family correlation', 'Changed from Code 3 to Code 0'), 0):,.0f}"
        ),
        "- **After family correlation**: Code 0 people = *[Run cps.py to populate]*": lambda: (
            f"- **After family correlation**: Code 0 people = {data_map.get(('After family correlation', 'Code 0 people'), 0):,.0f}"
        ),
        "- **Final**: Code 0 (NONE) = *[Run cps.py to populate]*": lambda: (
            f"- **Final**: Code 0 (NONE) = {data_map.get(('Final', 'Code 0 (NONE)'), 0):,.0f}"
        ),
        "- **Final**: Code 1 (CITIZEN) = *[Run cps.py to populate]*": lambda: (
            f"- **Final**: Code 1 (CITIZEN) = {data_map.get(('Final', 'Code 1 (CITIZEN)'), 0):,.0f}"
        ),
        "- **Final**: Code 2 (NON_CITIZEN_VALID_EAD) = *[Run cps.py to populate]*": lambda: (
            f"- **Final**: Code 2 (NON_CITIZEN_VALID_EAD) = {data_map.get(('Final', 'Code 2 (NON_CITIZEN_VALID_EAD)'), 0):,.0f}"
        ),
        "- **Final**: Code 3 (OTHER_NON_CITIZEN) = *[Run cps.py to populate]*": lambda: (
            f"- **Final**: Code 3 (OTHER_NON_CITIZEN) = {data_map.get(('Final', 'Code 3 (OTHER_NON_CITIZEN)'), 0):,.0f}"
        ),
        "- **Final**: Total undocumented (Code 0) = *[Run cps.py to populate]*": lambda: (
            f"- **Final**: Total undocumented (Code 0) = {data_map.get(('Final', 'Total undocumented (Code 0)'), 0):,.0f}"
        ),
        "- **Final**: Undocumented target = *[Run cps.py to populate]*": lambda: (
            f"- **Final**: Undocumented target = {data_map.get(('Final', 'Undocumented target'), 0):,.0f}"
        ),
    }

    # Apply replacements
    for old_text, replacement_func in replacements.items():
        if old_text in content:
            content = content.replace(old_text, replacement_func())

    # Write updated content back to file
    with open(doc_path, "w", encoding="utf-8") as f:
        f.write(content)

    print(f"Documentation updated with population numbers: {doc_path}")


@pipeline_node(
    PipelineNode(
        id="add_tips",
        label="Tips And Asset Imputation",
        node_type="library",
        description="Impute tip income and household asset inputs from SIPP donor data.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="legacy",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_tips(self, cps: h5py.File):
    self.save_dataset(cps)

    existing_data = self.load_dataset()
    person_household_id = np.asarray(
        existing_data.get(
            "person_household_id",
            existing_data.get("household_id"),
        )
    )
    interest_income = existing_data.get("interest_income")
    if interest_income is None:
        interest_income = np.asarray(
            existing_data.get(
                "taxable_interest_income",
                np.zeros(len(person_household_id), dtype=np.float32),
            )
        ) + np.asarray(
            existing_data.get(
                "tax_exempt_interest_income",
                np.zeros(len(person_household_id), dtype=np.float32),
            )
        )
    dividend_income = existing_data.get("dividend_income")
    if dividend_income is None:
        dividend_income = np.asarray(
            existing_data.get(
                "qualified_dividend_income",
                np.zeros(len(person_household_id), dtype=np.float32),
            )
        ) + np.asarray(
            existing_data.get(
                "non_qualified_dividend_income",
                np.zeros(len(person_household_id), dtype=np.float32),
            )
        )
    cps = pd.DataFrame(
        {
            "person_id": np.asarray(existing_data["person_id"]),
            "household_id": person_household_id,
            "employment_income": np.asarray(existing_data["employment_income"]),
            "interest_income": np.asarray(interest_income),
            "dividend_income": np.asarray(dividend_income),
            "rental_income": np.asarray(
                existing_data.get(
                    "rental_income",
                    np.zeros(len(person_household_id), dtype=np.float32),
                )
            ),
            "social_security": np.asarray(
                existing_data.get(
                    "social_security",
                    np.zeros(len(person_household_id), dtype=np.float32),
                )
            ),
            "pension_income": np.asarray(
                existing_data.get(
                    "pension_income",
                    np.zeros(len(person_household_id), dtype=np.float32),
                )
            ),
            "retirement_distributions": np.asarray(
                existing_data.get(
                    "retirement_distributions",
                    np.zeros(len(person_household_id), dtype=np.float32),
                )
            ),
            "age": np.asarray(existing_data["age"]),
            "is_female": np.asarray(existing_data["is_female"]),
        }
    )
    household_weight = existing_data.get("household_weight")
    if household_weight is not None:
        household_weight = np.asarray(household_weight)
        if len(household_weight) == len(cps):
            cps["household_weight"] = household_weight
        else:
            household_ids = np.asarray(existing_data["household_id"])
            household_weight_map = dict(zip(household_ids, household_weight))
            cps["household_weight"] = (
                pd.Series(person_household_id)
                .map(household_weight_map)
                .fillna(0)
                .values
            )
    else:
        cps["household_weight"] = 0.0

    # Get is_married from raw CPS data (A_MARITL codes: 1,2 = married)
    # Note: is_married in policyengine-us is Family-level, but we need
    # person-level for imputation models
    with _open_dataset_read_only(self.raw_cps) as raw_data:
        raw_person = raw_data["person"]
        cps["is_married"] = raw_person.A_MARITL.isin([1, 2]).values
        cps["is_tipped_occupation"] = derive_is_tipped_occupation(
            derive_treasury_tipped_occupation_code(raw_person.PEIOOCC)
        )

    cps["is_under_18"] = cps.age < 18
    cps["is_under_6"] = cps.age < 6
    cps["count_under_18"] = (
        cps.groupby("household_id")["is_under_18"]
        .sum()
        .loc[cps.household_id.values]
        .values
    )
    cps["count_under_6"] = (
        cps.groupby("household_id")["is_under_6"]
        .sum()
        .loc[cps.household_id.values]
        .values
    )
    cps["household_size"] = (
        cps.groupby("household_id")["person_id"]
        .transform("count")
        .astype(np.float32)
        .values
    )
    cps["retirement_income"] = cps["pension_income"].fillna(0) + cps[
        "retirement_distributions"
    ].fillna(0)
    cps["non_ssi_income"] = (
        cps["employment_income"].fillna(0)
        + cps["social_security"].fillna(0)
        + cps["retirement_income"].fillna(0)
    )
    cps = pd.DataFrame(cps)

    # Impute tips

    from policyengine_us_data.datasets.sipp import get_tip_model

    model = get_tip_model()

    cps["tip_income"] = model.predict(
        X_test=cps,
        mean_quantile=0.5,
    ).tip_income.values

    # Impute SIPP liquid assets used directly by resource-tested policy rules.
    # The SCF step below applies a stable 50/50 source-model draw for the
    # overlapping bank, stock, and bond leaves.

    from policyengine_us_data.datasets.sipp import get_asset_model

    asset_model = get_asset_model()

    asset_predictions = asset_model.predict(
        X_test=cps,
        mean_quantile=0.5,
    )
    cps["bank_account_assets"] = asset_predictions.bank_account_assets.values
    cps["stock_assets"] = asset_predictions.stock_assets.values
    cps["bond_assets"] = asset_predictions.bond_assets.values

    from policyengine_us_data.datasets.sipp import (
        SSI_DISABILITY_CRITERIA_VARIABLE,
        SSI_DISABILITY_DIFFICULTY_PREDICTORS,
        get_ssi_disability_model,
        predict_ssi_disability_criteria,
        preserve_under_65_ssi_disability_criteria,
    )

    n_persons = len(cps)
    for variable in [
        *SSI_DISABILITY_DIFFICULTY_PREDICTORS,
        "social_security_disability",
    ]:
        cps[variable] = np.asarray(
            existing_data.get(variable, np.zeros(n_persons)),
        )
    disability_benefits = np.asarray(
        existing_data.get("disability_benefits", np.zeros(n_persons)),
    )
    cps["has_disability_income"] = disability_benefits > 0
    ssi_disability_model = get_ssi_disability_model()
    meets_ssi_disability_criteria = predict_ssi_disability_criteria(
        ssi_disability_model,
        cps,
    )
    meets_ssi_disability_criteria = preserve_under_65_ssi_disability_criteria(
        meets_ssi_disability_criteria,
        age=existing_data.get("age", np.full(n_persons, 65)),
        ssi_reported=existing_data.get("ssi_reported"),
        existing_meets_ssi_disability_criteria=existing_data.get(
            SSI_DISABILITY_CRITERIA_VARIABLE
        ),
    )
    cps[SSI_DISABILITY_CRITERIA_VARIABLE] = meets_ssi_disability_criteria

    from policyengine_us_data.datasets.sipp import get_vehicle_model

    vehicle_model = get_vehicle_model()
    cps["is_household_head"] = np.asarray(
        existing_data["is_household_head"],
        dtype=bool,
    )
    household_vehicle_receiver = build_household_vehicle_receiver(
        cps,
        tenure_type=existing_data.get("tenure_type"),
    )
    vehicle_predictions = vehicle_model.predict(
        X_test=household_vehicle_receiver,
        mean_quantile=0.5,
    )
    household_vehicle_data = {
        "household_vehicles_owned": np.clip(
            np.rint(vehicle_predictions.household_vehicles_owned.values),
            0,
            None,
        ).astype(np.int32),
        "household_vehicles_value": np.clip(
            vehicle_predictions.household_vehicles_value.values,
            0,
            None,
        ).astype(np.float32),
    }

    # Drop temporary columns used only for imputation
    # is_married is person-level here but policyengine-us defines it at Family
    # level, so we must not save it
    cps = cps.drop(
        columns=[
            "is_married",
            "is_under_18",
            "is_under_6",
            "is_household_head",
            "has_disability_income",
            "household_size",
            "pension_income",
            "retirement_distributions",
            "retirement_income",
            "non_ssi_income",
        ],
        errors="ignore",
    )
    self.save_dataset(cps)
    self.save_dataset(household_vehicle_data)
    _drop_persisted_dataset_variables(
        self.file_path,
        TEMPORARY_IMPUTATION_SOURCE_VARIABLES,
    )


@pipeline_node(
    PipelineNode(
        id="add_org_inputs",
        label="ORG Labor-Market Inputs",
        node_type="library",
        description="Impute hourly wage, hourly-pay status, and union coverage from CPS ORG donors.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="current",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest tests/unit/datasets/test_org.py"],
    )
)
def add_org_labor_market_inputs(cps: h5py.File) -> None:
    """Impute ORG-derived wage and union inputs onto CPS persons."""
    n_persons = len(np.asarray(cps["age"]))
    household_ids = np.asarray(cps["household_id"], dtype=np.int64)
    person_household_ids = np.asarray(
        cps["person_household_id"],
        dtype=np.int64,
    )
    household_state_fips = np.asarray(cps["state_fips"], dtype=np.float32)
    household_index = {
        int(household_id): i for i, household_id in enumerate(household_ids)
    }
    person_state_fips = np.array(
        [
            household_state_fips[household_index[int(household_id)]]
            for household_id in person_household_ids
        ],
        dtype=np.float32,
    )

    receiver = build_org_receiver_frame(
        age=cps["age"],
        is_female=cps["is_female"],
        is_hispanic=cps["is_hispanic"],
        cps_race=cps["cps_race"],
        state_fips=person_state_fips,
        employment_income=cps["employment_income"],
        weekly_hours_worked=cps["weekly_hours_worked"],
    )
    if len(receiver) != n_persons:
        raise ValueError(
            f"ORG receiver frame has {len(receiver)} rows but CPS has "
            f"{n_persons} persons"
        )
    self_employment_income = np.asarray(
        cps.get(
            "self_employment_income",
            np.zeros(len(receiver), dtype=np.float32),
        ),
        dtype=np.float32,
    )
    predictions = predict_org_features(
        receiver,
        self_employment_income=self_employment_income,
    )

    for variable in ORG_IMPUTED_VARIABLES:
        values = predictions[variable].values
        if len(values) != n_persons:
            raise ValueError(
                f"ORG prediction for '{variable}' has {len(values)} entries "
                f"but CPS has {n_persons} persons"
            )
        if variable in ORG_BOOL_VARIABLES:
            cps[variable] = values.astype(bool)
        else:
            cps[variable] = values.astype(np.float32)


def add_overtime_occupation(cps: h5py.File, person: DataFrame) -> None:
    """Add occupation categories relevant to overtime eligibility calculations.
    Based on:
    https://www.law.cornell.edu/uscode/text/29/213
    https://www.congress.gov/crs-product/IF12480
    """
    cps["has_never_worked"] = person.POCCU2 == 53
    cps["is_military"] = person.POCCU2 == 52
    cps["is_computer_scientist"] = person.POCCU2 == 8
    cps["is_farmer_fisher"] = person.POCCU2 == 41
    cps["is_executive_administrative_professional"] = person.POCCU2.isin(
        [
            1,  # Chief executives, and managers
            2,  # Compensation, human resources, and infrastructure managers
            3,  # All other managers
            5,  # Business operations specialists
            6,  # Accountants and auditors
            7,  # Financial specialists
            9,  # Mathematical science occupations
            10,  # Architects, except naval
            11,  # Surveyors, cartographers, & photogrammetrists
            12,  # Engineering technologists and technicians
            13,  # Earth scientists
            14,  # Economists
            15,  # Psychologists, and other social scientists
            16,  # Health and safety specialists
            18,  # Lawyers, judges, magistrates, and other judicial workers
            19,  # Paralegals and all other legal support workers
            25,  # Registered nurses, therapists, and specific pathologists
            26,  # Veterinarians
            27,  # Health technicians and other healthcare practitioners
            28,  # Healthcare support occupations
            29,  # First-line supervisors of protective service workers
            34,  # First-line supervisors of housekeeping and janitorial workers
            36,  # Supervisors of personal care and service workers
            38,  # First-line supervisors of retail/non-retail sales workers
            39,  # Sales and related occupations
            40,  # Office & administrative support occupations
            42,  # First-line supervisors of construction trades workers
            50,  # Supervisors of transportation and flight related workers
        ]
    )


@pipeline_node(
    PipelineNode(
        id="add_auto_loan",
        label="Auto Loan And Net Worth Imputation",
        node_type="library",
        description="Impute auto loan balance, auto loan interest, and net worth from SCF donor data.",
        source_file="policyengine_us_data/datasets/cps/cps.py",
        status="legacy",
        stability="moving",
        pathways=["data_build"],
        validation_commands=["uv run pytest validation/stage_1/test_cps.py"],
    )
)
def add_auto_loan_interest_and_net_worth(self, cps: h5py.File) -> None:
    """ "Add auto loan balance, interest and net_worth variable."""
    self.save_dataset(cps)
    full_cps_data = self.load_dataset()
    cps_data = dict(full_cps_data)

    # Access raw CPS for additional variables
    with _open_dataset_read_only(self.raw_cps) as raw_data:
        person_data = raw_data.person

        # Preprocess the CPS for imputation
        lengths = {k: len(v) for k, v in cps_data.items()}
        var_len = cps_data["person_household_id"].shape[0]
        vars_of_interest = [name for name, ln in lengths.items() if ln == var_len]
        agg_data = pd.DataFrame({n: cps_data[n] for n in vars_of_interest})
        agg_data["interest_dividend_income"] = np.sum(
            [
                agg_data["taxable_interest_income"],
                agg_data["tax_exempt_interest_income"],
                agg_data["qualified_dividend_income"],
                agg_data["non_qualified_dividend_income"],
            ],
            axis=0,
        )
        agg_data["social_security_pension_income"] = np.sum(
            [
                agg_data["tax_exempt_private_pension_income"],
                agg_data["taxable_private_pension_income"],
                agg_data["social_security_retirement"],
            ],
            axis=0,
        )

        agg = (
            agg_data.groupby("person_household_id")[
                [
                    "employment_income",
                    "interest_dividend_income",
                    "social_security_pension_income",
                ]
            ]
            .sum()
            .rename(
                columns={
                    "employment_income": "household_employment_income",
                    "interest_dividend_income": "household_interest_dividend_income",
                    "social_security_pension_income": "household_social_security_pension_income",
                }
            )
            .reset_index()
        )

    def create_scf_reference_person_mask(cps_data, raw_person_data):
        """
        Create a boolean mask identifying SCF-style reference persons.

        SCF Reference Person Definition:
        - Single adult in household without a couple
        - In households with couples: male in mixed-sex couple OR older person in same-sex couple
        """
        all_persons_data = pd.DataFrame(
            {
                "person_household_id": cps_data["person_household_id"],
                "age": cps_data["age"],
            }
        )

        # Add sex variable (PESEX=2 means female in CPS)
        all_persons_data["is_female"] = (raw_person_data.A_SEX == 2).values

        # Add marital status (A_MARITL codes: 1,2 = married with spouse present/absent)
        all_persons_data["is_married"] = raw_person_data.A_MARITL.isin([1, 2]).values

        # Define adults as age 18+
        all_persons_data["is_adult"] = all_persons_data["age"] >= 18

        # Count adults per household
        adults_per_household = (
            all_persons_data[all_persons_data["is_adult"]]
            .groupby("person_household_id")
            .size()
            .reset_index(name="n_adults")
        )
        all_persons_data = all_persons_data.merge(
            adults_per_household, on="person_household_id", how="left"
        )

        # Identify couple households (households with exactly 2 married adults)
        married_adults_per_household = (
            all_persons_data[
                (all_persons_data["is_adult"]) & (all_persons_data["is_married"])
            ]
            .groupby("person_household_id")
            .size()
        )

        couple_households = married_adults_per_household[
            (married_adults_per_household == 2)
            & (all_persons_data.groupby("person_household_id")["n_adults"].first() == 2)
        ].index

        all_persons_data["is_couple_household"] = all_persons_data[
            "person_household_id"
        ].isin(couple_households)

        def determine_reference_person(group):
            """Determine reference person for a household group."""
            adults = group[group["is_adult"]]

            if len(adults) == 0:
                # No adults - select the oldest person regardless of age
                reference_idx = group["age"].idxmax()
                result = pd.Series([False] * len(group), index=group.index)
                result[reference_idx] = True
                return result

            elif len(adults) == 1:
                # Only one adult - they are the reference person
                result = pd.Series([False] * len(group), index=group.index)
                result[adults.index[0]] = True
                return result

            elif group["is_couple_household"].iloc[0] and len(adults) == 2:
                # Couple household with 2 adults
                couple_adults = adults.copy()

                # Check if same-sex couple
                if couple_adults["is_female"].nunique() == 1:
                    # Same-sex couple - choose older person
                    reference_idx = couple_adults["age"].idxmax()
                else:
                    # Mixed-sex couple - choose male (is_female = False)
                    male_adults = couple_adults[~couple_adults["is_female"]]
                    if len(male_adults) > 0:
                        reference_idx = male_adults.index[0]
                    else:
                        # Fallback to older person
                        reference_idx = couple_adults["age"].idxmax()

                result = pd.Series([False] * len(group), index=group.index)
                result[reference_idx] = True
                return result

            else:
                # Multiple adults but not a couple household
                # Use the oldest adult as reference person
                reference_idx = adults["age"].idxmax()
                result = pd.Series([False] * len(group), index=group.index)
                result[reference_idx] = True
                return result

        # Apply the reference person logic to each household
        all_persons_data["is_scf_reference_person"] = (
            all_persons_data.groupby("person_household_id")
            .apply(determine_reference_person, include_groups=False)
            .reset_index(level=0, drop=True)
        )

        return all_persons_data["is_scf_reference_person"].values

    mask = create_scf_reference_person_mask(cps_data, person_data)
    mask_len = mask.shape[0]
    original_person_household_ids = np.asarray(cps_data["person_household_id"])

    cps_data = {
        var: data[mask] if data.shape[0] == mask_len else data
        for var, data in cps_data.items()
    }

    CPS_RACE_MAPPING = {
        1: 1,  # White only -> WHITE
        2: 2,  # Black only -> BLACK/AFRICAN-AMERICAN
        3: 5,  # American Indian, Alaskan Native only -> OTHER
        4: 4,  # Asian only -> ASIAN
        5: 5,  # Hawaiian/Pacific Islander only -> OTHER
        6: 5,  # White-Black -> OTHER
        7: 5,  # White-AI -> OTHER
        8: 5,  # White-Asian -> OTHER
        9: 3,  # White-HP -> HISPANIC
        10: 5,  # Black-AI -> OTHER
        11: 5,  # Black-Asian -> OTHER
        12: 3,  # Black-HP -> HISPANIC
        13: 5,  # AI-Asian -> OTHER
        14: 5,  # AI-HP -> OTHER
        15: 3,  # Asian-HP -> HISPANIC
        16: 5,  # White-Black-AI -> OTHER
        17: 5,  # White-Black-Asian -> OTHER
        18: 5,  # White-Black-HP -> OTHER
        19: 5,  # White-AI-Asian -> OTHER
        20: 5,  # White-AI-HP -> OTHER
        21: 5,  # White-Asian-HP -> OTHER
        22: 5,  # Black-AI-Asian -> OTHER
        23: 5,  # White-Black-AI-Asian -> OTHER
        24: 5,  # White-AI-Asian-HP -> OTHER
        25: 5,  # Other 3 race comb. -> OTHER
        26: 5,  # Other 4 or 5 race comb. -> OTHER
    }

    # Apply the mapping to recode the race values
    cps_data["cps_race"] = np.vectorize(CPS_RACE_MAPPING.get)(cps_data["cps_race"])

    lengths = {k: len(v) for k, v in cps_data.items()}
    var_len = cps_data["person_household_id"].shape[0]
    vars_of_interest = [name for name, ln in lengths.items() if ln == var_len]
    receiver_data = pd.DataFrame({n: cps_data[n] for n in vars_of_interest})

    receiver_data = receiver_data.merge(
        agg[
            [
                "person_household_id",
                "household_employment_income",
                "household_interest_dividend_income",
                "household_social_security_pension_income",
            ]
        ],
        on="person_household_id",
        how="left",
    )
    receiver_data.drop("employment_income", axis=1, inplace=True)

    receiver_data.rename(
        columns={
            "household_employment_income": "employment_income",
            "household_interest_dividend_income": "interest_dividend_income",
            "household_social_security_pension_income": "social_security_pension_income",
        },
        inplace=True,
    )

    # Add is_married variable for household heads based on raw person data
    reference_persons = person_data[mask]
    receiver_data["is_married"] = reference_persons.A_MARITL.isin([1, 2]).values

    # Impute selected auto-loan fields, SCF equivalents for overlapping
    # financial asset leaves, and SCF-only balance-sheet leaves. We compute
    # net_worth from those components rather than storing an SCF aggregate.
    from policyengine_us_data.datasets.scf.scf import SCF_2022

    scf_dataset = SCF_2022()
    scf_data = scf_dataset.load_dataset()
    scf_data = pd.DataFrame({key: scf_data[key] for key in scf_data.keys()})

    PREDICTORS = [
        "age",
        "is_female",
        "cps_race",
        "is_married",
        "own_children_in_household",
        "employment_income",
        "interest_dividend_income",
        "social_security_pension_income",
    ]
    scf_net_worth_targets = add_scf_net_worth_target(scf_data)
    scf_financial_asset_targets = add_scf_financial_asset_targets(scf_data)
    scf_household_asset_targets = add_scf_household_asset_targets(scf_data)
    scf_component_targets = add_scf_net_worth_component_targets(scf_data)
    require_scf_net_worth_formula_targets(
        scf_financial_asset_targets=scf_financial_asset_targets,
        scf_household_asset_targets=scf_household_asset_targets,
        scf_component_targets=scf_component_targets,
        scf_net_worth_targets=scf_net_worth_targets,
    )
    IMPUTED_VARIABLES = (
        list(scf_net_worth_targets)
        + [
            "auto_loan_balance",
            "auto_loan_interest",
        ]
        + list(scf_financial_asset_targets)
        + list(scf_household_asset_targets)
        + list(scf_component_targets)
    )
    weights = ["wgt"]

    donor_data = scf_data[PREDICTORS + IMPUTED_VARIABLES + weights].copy()

    from microimpute.models.qrf import QRF
    import logging
    import os

    # Set root logger level
    log_level = os.getenv("PYTHON_LOG_LEVEL", "WARNING")

    # Specifically target the microimpute logger
    logging.getLogger("microimpute").setLevel(getattr(logging, log_level))

    qrf_model = QRF()
    donor_data = donor_data.sample(frac=0.5, random_state=42).reset_index(drop=True)
    fitted_model = qrf_model.fit(
        X_train=donor_data,
        predictors=PREDICTORS,
        imputed_variables=IMPUTED_VARIABLES,
        weight_col=weights[0],
        tune_hyperparameters=False,
    )
    imputations = fitted_model.predict(X_test=receiver_data)

    for var in IMPUTED_VARIABLES:
        if var == SCF_NET_WORTH_TARGET:
            continue
        cps[var] = imputations[var]

    for scf_var, policy_var in SCF_FINANCIAL_ASSET_POLICY_VARIABLES.items():
        if scf_var not in imputations or policy_var not in full_cps_data:
            continue
        blended_values = combine_sipp_and_scf_financial_assets(
            sipp_values=full_cps_data[policy_var],
            scf_household_values=imputations[scf_var].values,
            person_household_ids=original_person_household_ids,
            reference_person_mask=mask,
            time_period=self.time_period,
        )
        full_cps_data[policy_var] = blended_values
        if policy_var in cps:
            del cps[policy_var]
        cps[policy_var] = blended_values
        if scf_var in cps:
            del cps[scf_var]

    reference_household_ids = original_person_household_ids[mask]
    for scf_var, policy_var in SCF_HOUSEHOLD_ASSET_POLICY_VARIABLES.items():
        if (
            scf_var not in imputations
            or policy_var not in full_cps_data
            or "household_id" not in full_cps_data
        ):
            continue
        blended_values = combine_sipp_and_scf_household_assets(
            sipp_household_values=full_cps_data[policy_var],
            scf_household_values=imputations[scf_var].values,
            household_ids=full_cps_data["household_id"],
            reference_household_ids=reference_household_ids,
            time_period=self.time_period,
        )
        full_cps_data[policy_var] = blended_values
        if policy_var in cps:
            del cps[policy_var]
        cps[policy_var] = blended_values
        if scf_var in cps:
            del cps[scf_var]

    net_worth_components = {}
    for variable in ("bank_account_assets", "stock_assets", "bond_assets"):
        if variable in full_cps_data:
            net_worth_components[variable] = (
                aggregate_person_values_to_reference_households(
                    full_cps_data[variable],
                    original_person_household_ids,
                    mask,
                )
            )
    if "household_vehicles_value" in full_cps_data and "household_id" in full_cps_data:
        net_worth_components["household_vehicles_value"] = (
            align_household_values_to_reference_households(
                full_cps_data["household_vehicles_value"],
                full_cps_data["household_id"],
                reference_household_ids,
            )
        )
    for variable in SCF_NET_WORTH_COMPONENT_VARIABLES:
        if variable in cps:
            net_worth_components[variable] = cps[variable]
    if SCF_NET_WORTH_TARGET in imputations:
        net_worth_components = rebalance_scf_net_worth_components(
            components=net_worth_components,
            target_net_worth=imputations[SCF_NET_WORTH_TARGET].values,
        )
        for variable in SCF_NET_WORTH_COMPONENT_VARIABLES:
            if variable not in net_worth_components:
                continue
            if variable in cps:
                del cps[variable]
            cps[variable] = net_worth_components[variable]
    if "net_worth" in cps:
        del cps["net_worth"]
    cps["net_worth"] = compute_net_worth_from_components(
        components=net_worth_components
    )

    self.save_dataset(cps)


class CPS_2019(CPS):
    name = "cps_2019"
    label = "CPS 2019"
    raw_cps = CensusCPS_2019
    previous_year_raw_cps = CensusCPS_2018
    file_path = STORAGE_FOLDER / "cps_2019.h5"
    time_period = 2019
    frac = 0.5


class CPS_2020(CPS):
    name = "cps_2020"
    label = "CPS 2020"
    raw_cps = CensusCPS_2020
    previous_year_raw_cps = CensusCPS_2019
    file_path = STORAGE_FOLDER / "cps_2020.h5"
    time_period = 2020
    frac = 0.5


class CPS_2021(CPS):
    name = "cps_2021"
    label = "CPS 2021"
    raw_cps = CensusCPS_2021
    previous_year_raw_cps = CensusCPS_2020
    file_path = STORAGE_FOLDER / "cps_2021_v1_6_1.h5"
    time_period = 2021
    frac = 0.5


class CPS_2022(CPS):
    name = "cps_2022"
    label = "CPS 2022"
    raw_cps = CensusCPS_2022
    previous_year_raw_cps = CensusCPS_2021
    file_path = STORAGE_FOLDER / "cps_2022_v1_6_1.h5"
    time_period = 2022
    frac = 0.5


class CPS_2023(CPS):
    name = "cps_2023"
    label = "CPS 2023"
    raw_cps = CensusCPS_2023
    previous_year_raw_cps = CensusCPS_2022
    file_path = STORAGE_FOLDER / "cps_2023.h5"
    time_period = 2023
    frac = 0.5


class CPS_2024(CPS):
    name = "cps_2024"
    label = "CPS 2024"
    raw_cps = CensusCPS_2024
    previous_year_raw_cps = CensusCPS_2023
    file_path = STORAGE_FOLDER / "cps_2024.h5"
    time_period = 2024
    frac = 0.5


class CPS_2024_Full(CPS):
    name = "cps_2024_full"
    label = "CPS 2024 (full)"
    raw_cps = CensusCPS_2024
    previous_year_raw_cps = CensusCPS_2023
    file_path = STORAGE_FOLDER / "cps_2024_full.h5"
    time_period = 2024


if __name__ == "__main__":
    CPS_2024_Full().generate()
    CPS_2024().generate()