frs.py

"""
Family Resources Survey (FRS) dataset processing for PolicyEngine UK.

This module processes raw FRS survey data into PolicyEngine UK dataset format,
handling household demographics, income, benefits, and other survey variables.
The FRS is the primary source of UK household survey data used for tax-benefit
modelling and policy analysis.
"""

from functools import lru_cache
from pathlib import Path

import numpy as np
import pandas as pd
from policyengine_uk import CountryTaxBenefitSystem
from policyengine_uk.data import UKSingleYearDataset
from policyengine_uk.variables.household.income.employment_status import (
    EmploymentStatus,
)
from policyengine_uk_data.utils.datasets import (
    sum_to_entity,
    categorical,
    sum_from_positive_fields,
    sum_positive_variables,
    fill_with_mean,
    STORAGE_FOLDER,
)
from policyengine_uk_data.parameters import load_take_up_rate, load_parameter


# Canonical weeks-per-year conversion factor for annualising weekly survey
# variables. 365.25 / 7 ≈ 52.1786 accounts for leap years; using the rounded
# integer 52 would under-count by ~0.34%. Exposed at module level so sibling
# loaders (e.g. LCFS/ETB in `datasets/imputations/consumption.py`) can import
# the same value rather than re-defining `* 52` locally and drifting.
WEEKS_IN_YEAR = 365.25 / 7

LEGACY_JOBSEEKER_MIN_AGE = 18
HOURS_WORKED_WEEKS_PER_YEAR = 52
ESA_MIN_AGE = 16
ESA_HEALTH_EMPLOYMENT_STATUSES = (
    EmploymentStatus.LONG_TERM_DISABLED.name,
    EmploymentStatus.SHORT_TERM_DISABLED.name,
)
FORMULA_MODELED_EDUCATION_GRANT_VARIABLES = (
    "childcare_grant",
    "parents_learning_allowance",
    "adult_dependants_grant",
)
DISABLED_STUDENTS_ALLOWANCE_EXPENSE_INPUT = (
    "disabled_students_allowance_eligible_expenses"
)
DISABLED_STUDENTS_ALLOWANCE_FIRST_MODELED_YEAR = 2025
DISABLED_STUDENTS_ALLOWANCE_ELIGIBILITY_VARIABLES = (
    "maintenance_loan_in_england_system",
    "disabled_students_allowance_course_eligible",
    "disabled_students_allowance_has_qualifying_condition",
)
PIP_CATEGORY_SAFETY_MARGIN = 0.1


def _category_from_reported(
    reported,
    thresholds: tuple[tuple[str, float], ...],
) -> np.ndarray:
    """Convert annual reported amounts to PE-UK category inputs."""

    reported_weekly = pd.Series(reported).fillna(0).astype(float) / WEEKS_IN_YEAR
    return np.select(
        [
            reported_weekly >= rate * (1 - PIP_CATEGORY_SAFETY_MARGIN)
            for _, rate in thresholds
        ],
        [category for category, _ in thresholds],
        default="NONE",
    )


BENEFITS_IN_OWN_RIGHT_REPORTED_COLUMNS = (
    "universal_credit_reported",
    "jsa_contrib_reported",
    "jsa_income_reported",
    "esa_contrib_reported",
    "esa_income_reported",
)
NON_ADVANCED_EDUCATION_LEVELS = (
    "PRE_PRIMARY",
    "PRIMARY",
    "LOWER_SECONDARY",
    "UPPER_SECONDARY",
    "POST_SECONDARY",
)
# FRS government-training question variants use 10 or 13 for "None of these".
FRS_APPROVED_TRAINING_CODES = tuple(range(1, 10))
UNKNOWN_QUALIFYING_EDUCATION_OR_TRAINING_ENTRY_AGE = 1000


@lru_cache(maxsize=None)
def load_legacy_jobseeker_max_annual_hours(year: int) -> int:
    """Read the JSA single-claimant hours rule from policyengine-uk."""

    system = CountryTaxBenefitSystem()
    max_weekly_hours = int(system.parameters.gov.dwp.JSA.hours.single(str(year)))
    return max_weekly_hours * HOURS_WORKED_WEEKS_PER_YEAR


def derive_legacy_jobseeker_proxy(
    age,
    employment_status,
    hours_worked,
    current_education,
    employment_status_reported,
    state_pension_age,
    max_annual_hours,
) -> np.ndarray:
    """Approximate legacy JSA claimant-state from observed survey data.

    This is intentionally a proxy, not a legislative determination. It
    identifies person-level working-age adults who report being unemployed
    and working less than the legacy JSA 16-hour weekly limit. The
    ``hours_worked`` input is the annualised FRS-derived measure used in the
    dataset, so the threshold is converted to annual hours here.
    """

    age = np.asarray(age)
    employment_status = np.asarray(employment_status)
    hours_worked = np.asarray(hours_worked)
    current_education = np.asarray(current_education)
    employment_status_reported = np.asarray(employment_status_reported)
    state_pension_age = np.asarray(state_pension_age)

    return (
        employment_status_reported
        & (age >= LEGACY_JOBSEEKER_MIN_AGE)
        & (age < state_pension_age)
        & (employment_status == "UNEMPLOYED")
        & (hours_worked < max_annual_hours)
        & (current_education == "NOT_IN_EDUCATION")
    )


def derive_esa_health_condition_proxy(
    age,
    employment_status,
    employment_status_reported,
    state_pension_age,
) -> np.ndarray:
    """Approximate working-age ESA health-related claimant-state.

    This proxy relies only on person-level labour market status, not on
    current disability or incapacity benefit receipt. It is a dataset-side
    approximation for future modelling, not a direct observation of ESA
    legal entitlement or LCW/LCWRA status.
    """

    age = np.asarray(age)
    employment_status = np.asarray(employment_status)
    employment_status_reported = np.asarray(employment_status_reported)
    state_pension_age = np.asarray(state_pension_age)
    disability_labour_market_state = np.isin(
        employment_status, ESA_HEALTH_EMPLOYMENT_STATUSES
    )

    return (
        employment_status_reported
        & (age >= ESA_MIN_AGE)
        & (age < state_pension_age)
        & disability_labour_market_state
    )


def derive_esa_support_group_proxy(
    age,
    employment_status,
    hours_worked,
    esa_health_condition_proxy,
    employment_status_reported,
    state_pension_age,
) -> np.ndarray:
    """Approximate a severe-health ESA subgroup akin to support group.

    This is a stricter subset of ``esa_health_condition_proxy`` intended
    for future legacy ESA approximation work. It uses only non-receipt
    labour market signals already available in the survey.
    """

    age = np.asarray(age)
    employment_status = np.asarray(employment_status)
    hours_worked = np.asarray(hours_worked)
    esa_health_condition_proxy = np.asarray(esa_health_condition_proxy)
    employment_status_reported = np.asarray(employment_status_reported)
    state_pension_age = np.asarray(state_pension_age)
    severe_health_evidence = (employment_status == "LONG_TERM_DISABLED") & (
        hours_worked <= 0
    )

    return (
        employment_status_reported
        & (age >= ESA_MIN_AGE)
        & (age < state_pension_age)
        & esa_health_condition_proxy
        & severe_health_evidence
    )


def derive_receives_benefits_in_own_right(pe_person: pd.DataFrame) -> pd.Series:
    """Identify people reporting adult benefits that end QYP status."""

    return (
        pe_person[list(BENEFITS_IN_OWN_RIGHT_REPORTED_COLUMNS)].fillna(0).sum(axis=1)
        > 0
    )


def derive_is_in_non_advanced_education(
    current_education,
    is_apprentice=None,
) -> np.ndarray:
    """Identify current non-advanced education from PolicyEngine education states."""

    current_education = np.asarray(current_education)
    if is_apprentice is None:
        is_apprentice = np.zeros(len(current_education), dtype=bool)
    else:
        is_apprentice = np.asarray(is_apprentice)

    return np.isin(current_education, NON_ADVANCED_EDUCATION_LEVELS) & ~is_apprentice


def derive_is_in_approved_training_from_frs_person(
    person: pd.DataFrame,
) -> pd.Series:
    """Identify reported government training scheme participation in FRS."""

    if "train" not in person.columns:
        return pd.Series(False, index=person.index)

    train = pd.to_numeric(person.train, errors="coerce").fillna(0)
    return train.isin(FRS_APPROVED_TRAINING_CODES)


def derive_age_started_or_accepted_current_education_or_training(
    age,
    is_in_non_advanced_education,
    is_in_approved_training,
) -> np.ndarray:
    """Approximate the entry age for current QYP education or training.

    FRS observes current education/training status but not the age at which the
    current course or programme was started, enrolled on, or accepted. For
    people currently in qualifying education/training, cap the imputed entry
    age at 18 so observed 19-year-olds remain eligible for rules requiring
    entry before age 19.
    """

    age = np.asarray(age)
    in_qualifying_education_or_training = np.asarray(
        is_in_non_advanced_education
    ) | np.asarray(is_in_approved_training)

    return np.where(
        in_qualifying_education_or_training,
        np.minimum(age, 18),
        UNKNOWN_QUALIFYING_EDUCATION_OR_TRAINING_ENTRY_AGE,
    )


def derive_is_before_universal_credit_qualifying_young_person_terminal_date(
    age,
    is_in_non_advanced_education,
    is_in_approved_training,
) -> np.ndarray:
    """Approximate the UC terminal-date condition for observed 19-year-olds.

    FRS does not expose the date-of-birth and assessment-period detail needed
    to identify the exact 1 September terminal date. Use current qualifying
    education/training status as the microdata proxy for age-19 records.
    """

    age = np.asarray(age)
    in_qualifying_education_or_training = np.asarray(
        is_in_non_advanced_education
    ) | np.asarray(is_in_approved_training)

    return (age == 19) & in_qualifying_education_or_training


def add_legacy_benefit_proxies(
    pe_person: pd.DataFrame,
    employment_status_reported,
    state_pension_age,
    legacy_jobseeker_max_annual_hours,
) -> pd.DataFrame:
    """Populate person-scoped ESA/JSA proxy columns on the person frame.

    These remain person-level by design because the claimant-state inputs
    they approximate attach to individuals. Downstream benunit-level legacy
    benefit models should aggregate them explicitly rather than assuming the
    raw survey contains a benunit claimant-state field.
    """

    pe_person["legacy_jobseeker_proxy"] = derive_legacy_jobseeker_proxy(
        age=pe_person.age,
        employment_status=pe_person.employment_status,
        hours_worked=pe_person.hours_worked,
        current_education=pe_person.current_education,
        employment_status_reported=employment_status_reported,
        state_pension_age=state_pension_age,
        max_annual_hours=legacy_jobseeker_max_annual_hours,
    )
    pe_person["esa_health_condition_proxy"] = derive_esa_health_condition_proxy(
        age=pe_person.age,
        employment_status=pe_person.employment_status,
        employment_status_reported=employment_status_reported,
        state_pension_age=state_pension_age,
    )
    pe_person["esa_support_group_proxy"] = derive_esa_support_group_proxy(
        age=pe_person.age,
        employment_status=pe_person.employment_status,
        hours_worked=pe_person.hours_worked,
        esa_health_condition_proxy=pe_person.esa_health_condition_proxy,
        employment_status_reported=employment_status_reported,
        state_pension_age=state_pension_age,
    )
    return pe_person


def apply_legacy_benefit_proxies(
    pe_person: pd.DataFrame, sim, year: int, employment_status_reported
) -> pd.DataFrame:
    """Attach legacy ESA/JSA proxies using post-build simulation context."""

    state_pension_age = sim.calculate("state_pension_age", year).values
    legacy_jobseeker_max_annual_hours = load_legacy_jobseeker_max_annual_hours(year)
    return add_legacy_benefit_proxies(
        pe_person,
        employment_status_reported=employment_status_reported,
        state_pension_age=state_pension_age,
        legacy_jobseeker_max_annual_hours=legacy_jobseeker_max_annual_hours,
    )


def attach_legacy_benefit_proxies_from_frs_person(
    pe_person: pd.DataFrame, person: pd.DataFrame, sim, year: int
) -> pd.DataFrame:
    """Bridge raw FRS person fields into the proxy derivation hook."""

    employment_status_reported = person.empstati.fillna(0).to_numpy() > 0
    return apply_legacy_benefit_proxies(
        pe_person,
        sim,
        year,
        employment_status_reported=employment_status_reported,
    )


def derive_is_parent_from_frs_microdata(
    person_ids,
    person_benunit_ids,
    adult_person_ids,
    benunit_ids,
    dependent_children,
) -> np.ndarray:
    """Identify FRS adults in benefit units with dependent children.

    FRS benefit units contain either one adult or a couple plus any dependent
    children. Using the raw adult table and benefit-unit dependent-child count
    avoids ranking adults across the whole household when multiple benefit
    units share a household.
    """

    dependent_children_by_benunit = pd.Series(
        np.asarray(dependent_children, dtype=float),
        index=np.asarray(benunit_ids),
    )
    has_dependent_children = (
        pd.Series(np.asarray(person_benunit_ids))
        .map(dependent_children_by_benunit)
        .fillna(0)
        .to_numpy()
        > 0
    )
    is_adult_record = np.isin(np.asarray(person_ids), np.asarray(adult_person_ids))
    return is_adult_record & has_dependent_children


def _as_non_negative_array(values) -> np.ndarray:
    values = np.asarray(values, dtype=float)
    return np.maximum(np.nan_to_num(values, nan=0.0), 0.0)


def allocate_reported_education_grants(
    reported_grants, grant_capacities: dict[str, np.ndarray]
) -> dict[str, np.ndarray]:
    """Split aggregate FRS education grants across modelled grant capacity.

    The FRS reports several direct education grants in one aggregate field. When
    several modelled grants are plausible for the same person, allocate the
    reported amount proportionally to each grant's modelled capacity and keep any
    excess in the generic ``education_grants`` residual.
    """

    reported_grants = _as_non_negative_array(reported_grants)
    capacities = {
        variable: _as_non_negative_array(capacity)
        for variable, capacity in grant_capacities.items()
    }
    total_capacity = np.zeros_like(reported_grants, dtype=float)
    for variable, capacity in capacities.items():
        if capacity.shape != reported_grants.shape:
            raise ValueError(
                f"{variable} capacity has shape {capacity.shape}, "
                f"expected {reported_grants.shape}."
            )
        total_capacity += capacity

    allocation_fraction = np.divide(
        reported_grants,
        total_capacity,
        out=np.zeros_like(reported_grants, dtype=float),
        where=total_capacity > 0,
    )
    allocation_fraction = np.minimum(allocation_fraction, 1)

    allocations = {}
    allocated_total = np.zeros_like(reported_grants, dtype=float)
    for variable, capacity in capacities.items():
        allocation = capacity * allocation_fraction
        allocations[variable] = allocation
        allocated_total += allocation

    allocations["education_grants"] = np.maximum(reported_grants - allocated_total, 0)
    return allocations


def calculate_disabled_students_allowance_reported_grant_capacity(
    sim, year: int, maximum: float
) -> np.ndarray:
    if year < DISABLED_STUDENTS_ALLOWANCE_FIRST_MODELED_YEAR:
        return np.zeros_like(
            np.asarray(
                sim.calculate(
                    DISABLED_STUDENTS_ALLOWANCE_ELIGIBILITY_VARIABLES[0], year
                )
            ),
            dtype=float,
        )

    eligible = None
    for variable in DISABLED_STUDENTS_ALLOWANCE_ELIGIBILITY_VARIABLES:
        variable_eligible = np.asarray(sim.calculate(variable, year), dtype=bool)
        eligible = (
            variable_eligible if eligible is None else eligible & variable_eligible
        )
    equivalent_support = np.asarray(
        sim.calculate("disabled_students_allowance_receives_equivalent_support", year),
        dtype=bool,
    )
    return np.where(eligible & ~equivalent_support, float(maximum), 0.0)


def split_reported_education_grants(
    pe_person: pd.DataFrame, sim, year: int, dsa_maximum: float
) -> pd.DataFrame:
    """Move specific modelled grants out of the generic education-grant residual.

    PLA, ADG, and Childcare Grant remain formula-driven because they are
    calibration targets. Their modelled capacity is only used to avoid also
    counting the same reported FRS grant amount in the generic residual.
    DSA lacks a modelled amount signal, so its allocation seeds eligible
    expenses directly where the DSA parameter is available.
    """

    grant_capacities = {
        variable: sim.calculate(variable, year)
        for variable in FORMULA_MODELED_EDUCATION_GRANT_VARIABLES
    }
    grant_capacities[DISABLED_STUDENTS_ALLOWANCE_EXPENSE_INPUT] = (
        calculate_disabled_students_allowance_reported_grant_capacity(
            sim, year, dsa_maximum
        )
    )
    allocations = allocate_reported_education_grants(
        pe_person["education_grants"], grant_capacities
    )

    pe_person["education_grants"] = allocations["education_grants"]
    pe_person[DISABLED_STUDENTS_ALLOWANCE_EXPENSE_INPUT] = allocations[
        DISABLED_STUDENTS_ALLOWANCE_EXPENSE_INPUT
    ]

    return pe_person


def create_frs(
    raw_frs_folder: str,
    year: int,
) -> UKSingleYearDataset:
    """
    Process raw FRS data into PolicyEngine UK dataset format.

    Transforms the Family Resources Survey microdata from raw tab-delimited
    files into a structured PolicyEngine UK dataset with person, benefit unit,
    and household-level variables mapped to the appropriate tax-benefit system
    variables.

    Args:
        raw_frs_folder: Path to folder containing raw FRS .tab files.
        year: Survey year for the dataset.

    Returns:
        UKSingleYearDataset with processed FRS data ready for policy simulation.
    """
    raw_folder = Path(raw_frs_folder)
    if not raw_folder.exists():
        raise FileNotFoundError(f"Raw folder {raw_folder} does not exist.")

    frs = {}
    # Store SALSAC values before numeric conversion (for salary sacrifice
    # imputation)
    job_salsac_raw = None
    for file in raw_folder.glob("*.tab"):
        table_name = file.stem
        # Read raw data first
        df_raw = pd.read_csv(file, sep="\t")
        df_raw.columns = df_raw.columns.str.lower()

        # Preserve SALSAC column from job table before numeric conversion
        # SALSAC indicates salary sacrifice participation:
        # '1' = Yes, '2' = No, ' ' or blank = skip/not asked
        if table_name == "job" and "salsac" in df_raw.columns:
            job_salsac_raw = df_raw["salsac"].copy()

        # Make numeric where possible
        df = df_raw.apply(pd.to_numeric, errors="coerce")

        # Standardise column names to lower case (already done above)
        # df.columns = df.columns.str.lower()

        # Edit ID variables for simplicity
        if "sernum" in df.columns:
            df.rename(columns={"sernum": "household_id"}, inplace=True)

        if "benunit" in df.columns:
            # In the tables, benunit is the index of the benefit unit *within* the household.
            df.rename(columns={"benunit": "benunit_id"}, inplace=True)
            df["benunit_id"] = (df["household_id"] * 1e2 + df["benunit_id"]).astype(int)

        if "person" in df.columns:
            df.rename(columns={"person": "person_id"}, inplace=True)
            df["person_id"] = (df["household_id"] * 1e3 + df["person_id"]).astype(int)

        frs[table_name] = df

    # Combine adult and child tables for convenience

    frs["person"] = pd.concat([frs["adult"], frs["child"]]).sort_index().fillna(0)

    person = frs["person"]
    benunit = frs["benunit"]
    household = frs["househol"]
    household = household.set_index("household_id")
    pension = frs["pension"]
    oddjob = frs["oddjob"]
    account = frs["accounts"]
    job = frs["job"]
    # Add raw SALSAC column to job table for salary sacrifice imputation
    # SALSAC values: '1' = Yes (participates), '2' = No, ' '/blank = not asked
    if job_salsac_raw is not None:
        job["salsac_raw"] = job_salsac_raw.values
    benefits = frs["benefits"]
    maintenance = frs["maint"]
    pen_prov = frs["penprov"]
    childcare = frs["chldcare"]
    extchild = frs["extchild"]
    mortgage = frs["mortgage"]

    pe_person = pd.DataFrame()
    pe_benunit = pd.DataFrame()
    pe_household = pd.DataFrame()

    # Add primary and foreign keys
    pe_person["person_id"] = person.person_id
    pe_person["person_benunit_id"] = person.benunit_id
    pe_person["person_household_id"] = person.household_id
    pe_benunit["benunit_id"] = benunit.benunit_id
    pe_household["household_id"] = person.household_id.sort_values().unique()

    # Add grossing weights
    pe_household["household_weight"] = household.gross4.values

    # Add basic personal variables
    age = person.age80 + person.age
    pe_person["age"] = age
    # birth_year should be calculated from age and period in the model,
    # not stored as static data (see PolicyEngine/policyengine-uk#1352)
    # Age fields are AGE80 (top-coded) and AGE in the adult and child tables, respectively.
    pe_person["gender"] = np.where(person.sex == 1, "MALE", "FEMALE")
    pe_person["hours_worked"] = np.maximum(person.tothours, 0) * 52
    pe_person["is_household_head"] = person.hrpid == 1
    pe_person["is_benunit_head"] = person.uperson == 1
    dependent_children = (
        benunit.depchldb
        if "depchldb" in benunit
        else frs["child"]
        .groupby("benunit_id")
        .size()
        .reindex(benunit.benunit_id)
        .fillna(0)
        .to_numpy()
    )
    pe_person["is_parent"] = derive_is_parent_from_frs_microdata(
        person_ids=pe_person.person_id,
        person_benunit_ids=pe_person.person_benunit_id,
        adult_person_ids=frs["adult"].person_id,
        benunit_ids=pe_benunit.benunit_id,
        dependent_children=dependent_children,
    )
    MARITAL = [
        "MARRIED",
        "SINGLE",
        "SINGLE",
        "WIDOWED",
        "SEPARATED",
        "DIVORCED",
    ]
    pe_person["marital_status"] = categorical(
        person.marital, 2, range(1, 7), MARITAL
    ).fillna("SINGLE")

    # Add education levels
    if "fted" in person.columns:
        fted = person.fted
    else:
        fted = person.educft  # Renamed in FRS 2022-23
    typeed2 = person.typeed2

    def determine_education_level(fted_val, typeed2_val, age_val):
        # By default, not in education
        if fted_val in (2, -1, 0):
            return "NOT_IN_EDUCATION"
        # In pre-primary
        elif typeed2_val == 1:
            return "PRE_PRIMARY"
        # In primary education
        elif (
            typeed2_val in (2, 4)
            or (typeed2_val in (3, 8) and age_val < 11)
            or (typeed2_val == 0 and fted_val == 1 and age_val > 5 and age_val < 11)
        ):
            return "PRIMARY"
        # In lower secondary
        elif (
            typeed2_val in (5, 6)
            or (typeed2_val in (3, 8) and age_val >= 11 and age_val <= 16)
            or (typeed2_val == 0 and fted_val == 1 and age_val <= 16)
        ):
            return "LOWER_SECONDARY"
        # In upper secondary
        elif (
            typeed2_val == 7
            or (typeed2_val in (3, 8) and age_val > 16)
            or (typeed2_val == 0 and fted_val == 1 and age_val > 16)
        ):
            return "UPPER_SECONDARY"
        # In post-secondary
        elif typeed2_val in (7, 8) and age_val >= 19:
            return "POST_SECONDARY"
        # In tertiary
        elif typeed2_val == 9 or (typeed2_val == 0 and fted_val == 1 and age_val >= 19):
            return "TERTIARY"
        else:
            return "NOT_IN_EDUCATION"

    # Apply the function to determine education level
    pe_person["current_education"] = pd.Series(
        [determine_education_level(f, t, a) for f, t, a in zip(fted, typeed2, age)],
        index=pe_person.index,
    )
    pe_person["is_in_non_advanced_education"] = derive_is_in_non_advanced_education(
        pe_person.current_education
    )
    pe_person["is_in_approved_training"] = (
        derive_is_in_approved_training_from_frs_person(person)
    )
    pe_person["age_started_or_accepted_current_education_or_training"] = (
        derive_age_started_or_accepted_current_education_or_training(
            age,
            pe_person.is_in_non_advanced_education,
            pe_person.is_in_approved_training,
        )
    )
    pe_person["is_before_universal_credit_qualifying_young_person_terminal_date"] = (
        derive_is_before_universal_credit_qualifying_young_person_terminal_date(
            age,
            pe_person.is_in_non_advanced_education,
            pe_person.is_in_approved_training,
        )
    )

    # Add highest education from EDUCQUAL (highest qualification achieved)
    # Codes from FRS ADT_324X classification; unmapped codes default to UPPER_SECONDARY
    EDUCQUAL_MAP = {
        1: "NOT_COMPLETED_PRIMARY",
        2: "LOWER_SECONDARY",  # GCSE D-G / CSE 2-5
        3: "LOWER_SECONDARY",  # GCSE A-C / O-level A-C
        4: "UPPER_SECONDARY",  # AS-level
        5: "UPPER_SECONDARY",  # A-level (1 subject)
        6: "UPPER_SECONDARY",  # A-level (2 subjects)
        7: "UPPER_SECONDARY",  # A-level (3+ subjects)
        8: "LOWER_SECONDARY",  # Scottish Standard/Ordinary Grade
        9: "UPPER_SECONDARY",  # Scottish Higher Grade
        10: "UPPER_SECONDARY",  # Scottish 6th Year Studies
        11: "POST_SECONDARY",  # HNC/HND
        12: "POST_SECONDARY",  # City & Guilds advanced / BTEC National
        13: "UPPER_SECONDARY",  # City & Guilds craft / BTEC General
        14: "POST_SECONDARY",  # ONC/OND / BTEC National (lower)
        15: "UPPER_SECONDARY",  # City & Guilds foundation
        16: "POST_SECONDARY",  # RSA advanced
        17: "TERTIARY",  # First/foundation degree
        18: "TERTIARY",  # Second degree
        19: "TERTIARY",  # Higher degree (Masters/PhD)
        20: "TERTIARY",  # PGCE / teaching qualification
        21: "TERTIARY",  # Nursing/paramedical qualification
        66: "UPPER_SECONDARY",  # NVQ/SVQ Level 1
        67: "UPPER_SECONDARY",  # NVQ/SVQ Level 2
        68: "UPPER_SECONDARY",  # NVQ/SVQ Level 3
        69: "POST_SECONDARY",  # NVQ/SVQ Level 4
        70: "TERTIARY",  # NVQ/SVQ Level 5
    }
    # Codes 22-65 and 71-85 are further vocational/professional qualifications;
    # treat as POST_SECONDARY. Codes 86-87 are catch-alls; treat as UPPER_SECONDARY.
    for code in range(22, 66):
        EDUCQUAL_MAP[code] = "POST_SECONDARY"
    for code in range(71, 86):
        EDUCQUAL_MAP[code] = "POST_SECONDARY"
    EDUCQUAL_MAP[86] = "UPPER_SECONDARY"
    EDUCQUAL_MAP[87] = "UPPER_SECONDARY"

    educqual = pd.to_numeric(person.educqual, errors="coerce")
    pe_person["highest_education"] = educqual.map(EDUCQUAL_MAP).fillna(
        "UPPER_SECONDARY"
    )

    # Add employment status
    EMPLOYMENTS = [
        "CHILD",
        "FT_EMPLOYED",
        "PT_EMPLOYED",
        "FT_SELF_EMPLOYED",
        "PT_SELF_EMPLOYED",
        "UNEMPLOYED",
        "RETIRED",
        "STUDENT",
        "CARER",
        "LONG_TERM_DISABLED",
        "SHORT_TERM_DISABLED",
    ]
    pe_person["employment_status"] = categorical(
        person.empstati, 1, range(12), EMPLOYMENTS
    ).fillna("LONG_TERM_DISABLED")

    REGIONS = [
        "NORTH_EAST",
        "NORTH_WEST",
        "YORKSHIRE",
        "EAST_MIDLANDS",
        "WEST_MIDLANDS",
        "EAST_OF_ENGLAND",
        "LONDON",
        "SOUTH_EAST",
        "SOUTH_WEST",
        "WALES",
        "SCOTLAND",
        "NORTHERN_IRELAND",
        "UNKNOWN",
    ]
    pe_household["region"] = categorical(
        household.gvtregno, 14, [1, 2] + list(range(4, 15)), REGIONS
    ).values
    TENURES = [
        "RENT_FROM_COUNCIL",
        "RENT_FROM_HA",
        "RENT_PRIVATELY",
        "RENT_PRIVATELY",
        "OWNED_OUTRIGHT",
        "OWNED_WITH_MORTGAGE",
    ]
    pe_household["tenure_type"] = categorical(
        household.ptentyp2, 3, range(1, 7), TENURES
    ).values
    frs["num_bedrooms"] = household.bedroom6
    ACCOMMODATIONS = [
        "HOUSE_DETACHED",
        "HOUSE_SEMI_DETACHED",
        "HOUSE_TERRACED",
        "FLAT",
        "CONVERTED_HOUSE",
        "MOBILE",
        "OTHER",
    ]
    pe_household["accommodation_type"] = categorical(
        household.typeacc, 1, range(1, 8), ACCOMMODATIONS
    ).values

    # Impute Council Tax

    # Only ~25% of household report Council Tax bills - use
    # these to build a model to impute missing values
    CT_valid = household.ctannual > 0

    # Find the mean reported Council Tax bill for a given
    # (region, CT band, is-single-person-household) triplet
    region = household.gvtregno[CT_valid]
    band = household.ctband[CT_valid]
    single_person = (household.adulth == 1)[CT_valid]
    ctannual = household.ctannual[CT_valid]

    # Build the table
    ct_mean = ctannual.groupby([region, band, single_person], dropna=False).mean()
    ct_mean = ct_mean.replace(-1, ct_mean.mean())

    # For every household consult the table to find the imputed
    # Council Tax bill
    pairs = household.set_index(
        [household.gvtregno, household.ctband, (household.adulth == 1)]
    )
    hh_CT_mean = pd.Series(index=pairs.index)
    has_mean = pairs.index.isin(ct_mean.index)
    hh_CT_mean[has_mean] = ct_mean[pairs.index[has_mean]].values
    hh_CT_mean[~has_mean] = 0
    ct_imputed = hh_CT_mean

    # For households which originally reported Council Tax,
    # use the reported value. Otherwise, use the imputed value
    council_tax = pd.Series(
        np.where(
            # 2018 FRS uses blanks for missing values, 2019 FRS
            # uses -1 for missing values
            (household.ctannual < 0) | household.ctannual.isna(),
            np.maximum(ct_imputed, 0).values,
            household.ctannual,
        )
    )
    pe_household["council_tax"] = council_tax.fillna(0)
    BANDS = ["A", "B", "C", "D", "E", "F", "G", "H", "I"]
    # Band 1 is the most common
    pe_household["council_tax_band"] = (
        categorical(household.ctband, 1, range(1, 10), BANDS).fillna("D").values
    )
    # Domestic rates variables are all weeklyised, unlike Council Tax variables (despite the variable name suggesting otherwise)
    if year < 2021:
        DOMESTIC_RATES_VARIABLE = "rtannual"
    else:
        DOMESTIC_RATES_VARIABLE = "niratlia"
    pe_household["domestic_rates"] = (
        np.select(
            [
                household[DOMESTIC_RATES_VARIABLE] >= 0,
                household.rt2rebam >= 0,
                True,
            ],
            [
                household[DOMESTIC_RATES_VARIABLE],
                household.rt2rebam,
                0,
            ],
        )
        * 52
    ).astype(float)

    WEEKS_IN_YEAR = 365.25 / 7

    pe_person["employment_income"] = np.maximum(0, person.inearns) * WEEKS_IN_YEAR

    pension_payment = sum_to_entity(
        pension.penpay * (pension.penpay > 0),
        pension.person_id,
        person.person_id,
    )
    pension_tax_paid = sum_to_entity(
        (pension.ptamt * ((pension.ptinc == 2) & (pension.ptamt > 0))),
        pension.person_id,
        person.person_id,
    )
    pension_deductions_removed = sum_to_entity(
        pension.poamt
        * (((pension.poinc == 2) | (pension.penoth == 1)) & (pension.poamt > 0)),
        pension.person_id,
        person.person_id,
    )

    pe_person["private_pension_income"] = (
        pension_payment + pension_tax_paid + pension_deductions_removed
    ) * WEEKS_IN_YEAR

    pe_person["self_employment_income"] = np.maximum(0, person.seincam2) * WEEKS_IN_YEAR

    INVERTED_BASIC_RATE = 1.25

    pe_person["tax_free_savings_income"] = np.maximum(
        0,
        sum_to_entity(
            account.accint * (account.account == 21),
            account.person_id,
            person.person_id,
        )
        * WEEKS_IN_YEAR,
    )
    taxable_savings_interest = (
        sum_to_entity(
            (account.accint * np.where(account.acctax == 1, INVERTED_BASIC_RATE, 1))
            * (account.account.isin((1, 3, 5, 27, 28))),
            account.person_id,
            person.person_id,
        )
        * WEEKS_IN_YEAR
    )
    pe_person["savings_interest_income"] = np.maximum(
        0,
        taxable_savings_interest + pe_person["tax_free_savings_income"].values,
    )
    pe_person["dividend_income"] = np.maximum(
        0,
        sum_to_entity(
            (account.accint * np.where(account.invtax == 1, INVERTED_BASIC_RATE, 1))
            * (
                ((account.account == 6) & (account.invtax == 1))  # GGES
                | account.account.isin((7, 8))  # Stocks/shares/UITs
            ),
            account.person_id,
            person.index,
        )
        * 52,
    )
    is_head = person.hrpid == 1
    household_property_income = (
        household.tentyp2.isin((5, 6)) * household.subrent
    )  # Owned and subletting
    persons_household_property_income = (
        pd.Series(
            household_property_income[person.household_id].values,
            index=person.person_id,
        )
        .fillna(0)
        .values
    )
    pe_person["property_income"] = (
        np.maximum(
            0,
            is_head * persons_household_property_income + person.cvpay + person.royyr1,
        )
        * WEEKS_IN_YEAR
    )
    maintenance_to_self = np.maximum(
        pd.Series(np.where(person.mntus1 == 2, person.mntusam1, person.mntamt1)).fillna(
            0
        ),
        0,
    )
    maintenance_from_dwp = person.mntamt2
    pe_person["maintenance_income"] = (
        sum_positive_variables([maintenance_to_self, maintenance_from_dwp])
        * WEEKS_IN_YEAR
    )

    odd_job_income = sum_to_entity(
        oddjob.ojamt * (oddjob.ojnow == 1), oddjob.person_id, person.person_id
    )

    MISC_INCOME_FIELDS = [
        "allpay2",
        "royyr2",
        "royyr3",
        "royyr4",
        "chamtern",
        "chamttst",
    ]

    pe_person["miscellaneous_income"] = (
        odd_job_income + sum_from_positive_fields(person, MISC_INCOME_FIELDS)
    ) * WEEKS_IN_YEAR

    PRIVATE_TRANSFER_INCOME_FIELDS = [
        "apamt",
        "apdamt",
        "pareamt",
        "allpay2",
        "allpay3",
        "allpay4",
    ]

    pe_person["private_transfer_income"] = (
        sum_from_positive_fields(person, PRIVATE_TRANSFER_INCOME_FIELDS) * WEEKS_IN_YEAR
    )

    pe_person["lump_sum_income"] = person.redamt

    pe_person["student_loan_repayments"] = person.slrepamt * WEEKS_IN_YEAR

    BENEFIT_CODES = dict(
        child_benefit=3,
        income_support=19,
        housing_benefit=94,
        attendance_allowance=12,
        dla_sc=1,
        dla_m=2,
        iidb=15,
        carers_allowance=13,
        sda=10,
        afcs=8,
        ssmg=22,
        pension_credit=4,
        child_tax_credit=91,
        working_tax_credit=90,
        state_pension=5,
        winter_fuel_allowance=62,
        incapacity_benefit=17,
        universal_credit=95,
        pip_m=97,
        pip_dl=96,
    )
    for benefit, code in BENEFIT_CODES.items():
        pe_person[benefit + "_reported"] = (
            sum_to_entity(
                benefits.benamt * (benefits.benefit == code),
                benefits.person_id.values,
                person.person_id,
            )
            * WEEKS_IN_YEAR
        )

    pe_person["jsa_contrib_reported"] = (
        sum_to_entity(
            benefits.benamt * (benefits.var2.isin((1, 3))) * (benefits.benefit == 14),
            benefits.person_id,
            person.person_id,
        )
        * WEEKS_IN_YEAR
    )
    pe_person["jsa_income_reported"] = (
        sum_to_entity(
            benefits.benamt * (benefits.var2.isin((2, 4))) * (benefits.benefit == 14),
            benefits.person_id,
            person.person_id,
        )
        * WEEKS_IN_YEAR
    )
    pe_person["esa_contrib_reported"] = (
        sum_to_entity(
            benefits.benamt * (benefits.var2.isin((1, 3))) * (benefits.benefit == 16),
            benefits.person_id,
            person.person_id,
        )
        * WEEKS_IN_YEAR
    )
    pe_person["esa_income_reported"] = (
        sum_to_entity(
            benefits.benamt * (benefits.var2.isin((2, 4))) * (benefits.benefit == 16),
            benefits.person_id,
            person.person_id,
        )
        * WEEKS_IN_YEAR
    )
    pe_person["receives_benefits_in_own_right"] = derive_receives_benefits_in_own_right(
        pe_person
    )

    pe_person["bsp_reported"] = (
        sum_to_entity(
            benefits.benamt * (benefits.benefit.isin((6, 9))),
            benefits.person_id,
            person.person_id,
        )
        * WEEKS_IN_YEAR
    )

    pe_person["winter_fuel_allowance_reported"] /= WEEKS_IN_YEAR

    pe_person["statutory_sick_pay"] = person.sspadj * WEEKS_IN_YEAR
    pe_person["statutory_maternity_pay"] = person.smpadj * WEEKS_IN_YEAR

    pe_person["student_loans"] = np.maximum(person.tuborr, 0)
    if "adema" not in person.columns:
        person["adema"] = person.eduma
        person["ademaamt"] = person.edumaamt
    pe_person["adult_ema"] = fill_with_mean(person, "adema", "ademaamt")
    pe_person["child_ema"] = fill_with_mean(person, "chema", "chemaamt")

    pe_person["access_fund"] = np.maximum(person.accssamt, 0) * WEEKS_IN_YEAR

    pe_person["education_grants"] = np.maximum(
        person[["grtdir1", "grtdir2"]].sum(axis=1), 0
    )

    pe_person["council_tax_benefit_reported"] = np.maximum(
        (person.hrpid == 1)
        * pd.Series(
            household.ctrebamt[person.household_id.values].values,
            index=person.person_id,
        )
        .fillna(0)
        .values
        * WEEKS_IN_YEAR,
        0,
    )

    pe_person["healthy_start_vouchers"] = person.heartval * WEEKS_IN_YEAR

    pe_person["free_school_breakfasts"] = person.fsbval * WEEKS_IN_YEAR
    pe_person["free_school_fruit_veg"] = person.fsfvval * WEEKS_IN_YEAR
    pe_person["free_school_meals"] = person.fsmval * WEEKS_IN_YEAR

    pe_person["maintenance_expenses"] = (
        pd.Series(
            np.where(maintenance.mrus == 2, maintenance.mruamt, maintenance.mramt)
        )
        .groupby(maintenance.person_id)
        .sum()
        .reindex(person.person_id)
        .fillna(0)
        .values
        * WEEKS_IN_YEAR
    )
    pe_household["rent"] = household.hhrent.fillna(0).values * WEEKS_IN_YEAR
    pe_household["mortgage_interest_repayment"] = (
        household.mortint.fillna(0).values * WEEKS_IN_YEAR
    )
    mortgage_capital = np.where(mortgage.rmort == 1, mortgage.rmamt, mortgage.borramt)
    mortgage_capital_repayment = sum_to_entity(
        mortgage_capital / mortgage.mortend,
        mortgage.household_id,
        household.index,
    )
    pe_household["mortgage_capital_repayment"] = mortgage_capital_repayment

    pe_person["childcare_expenses"] = (
        sum_to_entity(
            childcare.chamt * (childcare.cost == 1) * (childcare.registrd == 1),
            childcare.person_id,
            person.person_id,
        )
        * 52
    )

    pe_person["personal_pension_contributions"] = np.maximum(
        0,
        sum_to_entity(
            pen_prov.penamt[pen_prov.stemppen.isin((5, 6))],
            pen_prov.person_id,
            person.person_id,
        ).clip(0, pen_prov.penamt.quantile(0.95))
        * WEEKS_IN_YEAR,
    )
    pe_person["employee_pension_contributions"] = np.maximum(
        0,
        sum_to_entity(job.deduc1.fillna(0), job.person_id, person.person_id)
        * WEEKS_IN_YEAR,
    )
    pe_person["employer_pension_contributions"] = (
        pe_person["employee_pension_contributions"] * 3
    )  # Rough estimate based on aggregates.
    # Salary sacrifice pension contributions from FRS Job table (SPNAMT field)
    # SPNAMT represents employer pension contributions made via salary sacrifice
    # arrangements where employees forego salary in exchange for increased pension
    # contributions. This is separate from regular employee pension contributions
    # (deduc1) and provides tax advantages for both employer and employee.
    # Uses same pattern as employee_pension_contributions (deduc1) without outlier
    # clipping, as job-level data is generally cleaner than pension provider data.
    # Source: https://datacatalogue.ukdataservice.ac.uk/datasets/dataset/630d4a8d-ba6a-82b3-f33d-c713c66efcb3
    # Note: Values are annualized from weekly amounts reported in the survey.
    pe_person["pension_contributions_via_salary_sacrifice"] = np.maximum(
        0,
        sum_to_entity(job.spnamt.fillna(0), job.person_id, person.person_id)
        * WEEKS_IN_YEAR,
    )

    # Salary sacrifice participation indicator from SALSAC field
    # Used for imputation: 1 = Yes, 0 = No, -1 = not asked (skip)
    # This allows distinguishing between explicit No responses and
    # respondents who were not asked the question (imputation candidates)
    if "salsac_raw" in job.columns:
        salsac_numeric = (
            job["salsac_raw"].map({"1": 1, "2": 0, " ": -1}).fillna(-1).astype(int)
        )
        # Aggregate to person level: take max (any job with SS = person has SS)
        pe_person["salary_sacrifice_reported"] = np.clip(
            sum_to_entity(
                (salsac_numeric == 1).astype(int),
                job.person_id,
                person.person_id,
            ),
            0,
            1,
        )
        # Track if person was asked about SS in any job (for imputation)
        pe_person["salary_sacrifice_asked"] = np.clip(
            sum_to_entity(
                (salsac_numeric >= 0).astype(int),
                job.person_id,
                person.person_id,
            ),
            0,
            1,
        )
    else:
        # If SALSAC not available, mark all as not asked
        pe_person["salary_sacrifice_reported"] = 0
        pe_person["salary_sacrifice_asked"] = 0

    pe_household["housing_service_charges"] = (
        pd.DataFrame(
            [
                household[f"chrgamt{i}"] * (household[f"chrgamt{i}"] > 0)
                for i in range(1, 10)
            ]
        )
        .sum()
        .values
        * WEEKS_IN_YEAR
    )
    pe_household["structural_insurance_payments"] = (
        household.struins.values * WEEKS_IN_YEAR
    )
    pe_household["water_and_sewerage_charges"] = (
        pd.Series(
            np.where(
                household.gvtregno == 12,
                household.csewamt + household.cwatamtd,
                household.watsewrt,
            )
        )
        .fillna(0)
        .values
        * WEEKS_IN_YEAR
    )

    pe_household["external_child_payments"] = sum_to_entity(
        extchild.nhhamt * WEEKS_IN_YEAR,
        extchild.household_id,
        household.index,
    )

    benefit = CountryTaxBenefitSystem().parameters(year).gov.dwp

    pe_person["aa_category"] = _category_from_reported(
        pe_person["attendance_allowance_reported"],
        (
            ("HIGHER", benefit.attendance_allowance.higher),
            ("LOWER", benefit.attendance_allowance.lower),
        ),
    )
    pe_person["dla_sc_category"] = _category_from_reported(
        pe_person["dla_sc_reported"],
        (
            ("HIGHER", benefit.dla.self_care.higher),
            ("MIDDLE", benefit.dla.self_care.middle),
            ("LOWER", benefit.dla.self_care.lower),
        ),
    )
    pe_person["dla_m_category"] = _category_from_reported(
        pe_person["dla_m_reported"],
        (
            ("HIGHER", benefit.dla.mobility.higher),
            ("LOWER", benefit.dla.mobility.lower),
        ),
    )
    pe_person["pip_dl_category"] = _category_from_reported(
        pe_person["pip_dl_reported"],
        (
            ("ENHANCED", benefit.pip.daily_living.enhanced),
            ("STANDARD", benefit.pip.daily_living.standard),
        ),
    )
    pe_person["pip_m_category"] = _category_from_reported(
        pe_person["pip_m_reported"],
        (
            ("ENHANCED", benefit.pip.mobility.enhanced),
            ("STANDARD", benefit.pip.mobility.standard),
        ),
    )

    has_pip = (pe_person["pip_dl_category"] != "NONE") | (
        pe_person["pip_m_category"] != "NONE"
    )
    has_dla = (pe_person["dla_sc_category"] != "NONE") | (
        pe_person["dla_m_category"] != "NONE"
    )
    pe_person["is_disabled_for_benefits"] = has_dla | has_pip

    pe_person["is_enhanced_disabled_for_benefits"] = (
        pe_person["dla_sc_category"] == "HIGHER"
    )

    # Child Tax Credit Regulations 2002 s. 8
    paragraph_3 = pe_person["dla_sc_category"] == "HIGHER"
    paragraph_4 = pe_person["pip_dl_category"] == "ENHANCED"
    paragraph_5 = pe_person.afcs_reported > 0
    pe_person["is_severely_disabled_for_benefits"] = (
        paragraph_3 | paragraph_4 | paragraph_5
    )

    pe_person = pe_person.drop(
        columns=[
            "attendance_allowance_reported",
            "dla_sc_reported",
            "dla_m_reported",
            "pip_dl_reported",
            "pip_m_reported",
        ],
        errors="ignore",
    )

    dataset = UKSingleYearDataset(
        person=pe_person,
        benunit=pe_benunit,
        household=pe_household,
        fiscal_year=year,
    )

    # Randomly select broad rental market areas from regions.
    from policyengine_uk import Microsimulation

    sim = Microsimulation(dataset=dataset)
    region = sim.populations["benunit"].household("region", dataset.time_period)
    lha_category = sim.calculate("LHA_category", year)
    brma = np.empty(len(region), dtype=object)

    # Sample from a random BRMA in the region, weighted by the number of observations in each BRMA
    lha_list_of_rents = pd.read_csv(STORAGE_FOLDER / "lha_list_of_rents.csv.gz")
    lha_list_of_rents = lha_list_of_rents.copy()

    for possible_region in lha_list_of_rents.region.unique():
        for possible_lha_category in lha_list_of_rents.lha_category.unique():
            lor_mask = (lha_list_of_rents.region == possible_region) & (
                lha_list_of_rents.lha_category == possible_lha_category
            )
            mask = (region == possible_region) & (lha_category == possible_lha_category)
            brma[mask] = lha_list_of_rents[lor_mask].brma.sample(
                n=len(region[mask]), replace=True
            )

    # Convert benunit-level BRMAs to household-level BRMAs (pick a random one)

    df = pd.DataFrame(
        {
            "brma": brma,
            "household_id": sim.populations["benunit"].household(
                "household_id", sim.dataset.time_period
            ),
        }
    )

    df = df.groupby("household_id").brma.aggregate(lambda x: x.sample(n=1).iloc[0])
    brmas = df[sim.calculate("household_id")].values

    pe_household["brma"] = brmas

    # Dataset-side claimant-state approximations for future legacy ESA/JSA
    # modelling. These are explicit proxies based on observed survey
    # conditions, not legislative determinations.
    pe_person = attach_legacy_benefit_proxies_from_frs_person(
        pe_person, person, sim, year
    )

    if (pe_person["education_grants"] > 0).any():
        student_support_dataset = UKSingleYearDataset(
            person=pe_person,
            benunit=pe_benunit,
            household=pe_household,
            fiscal_year=year,
        )
        student_support_sim = Microsimulation(dataset=student_support_dataset)
        dsa_maximum = student_support_sim.tax_benefit_system.parameters(
            year
        ).gov.dfe.disabled_students_allowance.maximum
        pe_person = split_reported_education_grants(
            pe_person, student_support_sim, year, dsa_maximum
        )

    # Generate stochastic take-up decisions
    # All randomness is generated here in the data package using take-up rates
    # stored in YAML parameter files. This keeps the country package purely
    # deterministic.

    generator = np.random.default_rng(seed=100)

    # Load take-up rates from parameter files
    child_benefit_rate = load_take_up_rate("child_benefit", year)
    pension_credit_rate = load_take_up_rate("pension_credit", year)
    universal_credit_rate = load_take_up_rate("universal_credit", year)
    marriage_allowance_rate = load_take_up_rate("marriage_allowance", year)
    child_benefit_opts_out_rate = load_take_up_rate("child_benefit_opts_out_rate", year)
    tfc_rate = load_take_up_rate("tax_free_childcare", year)
    extended_childcare_rate = load_take_up_rate("extended_childcare", year)
    universal_childcare_rate = load_take_up_rate("universal_childcare", year)
    targeted_childcare_rate = load_take_up_rate("targeted_childcare", year)
    scp_under_6_rate = load_take_up_rate("scp_under_6", year)
    scp_6_plus_rate = load_take_up_rate("scp_6_plus", year)

    # Generate take-up decisions by comparing random draws to take-up rates,
    # anchored to reported receipts where the FRS captures them. Respondents
    # who report positive receipt of a benefit are assigned takeup=True with
    # certainty; the remaining non-reporters are filled probabilistically to
    # hit the aggregate target rate. See policyengine_uk_data/utils/takeup.py.
    from policyengine_uk_data.utils.takeup import (
        assign_takeup_with_reported_anchors,
    )

    def _reported_benunit_mask(person_column: str) -> np.ndarray:
        reporter_benunits = set(
            pe_person.loc[pe_person[person_column] > 0, "person_benunit_id"].values
        )
        return pe_benunit["benunit_id"].isin(reporter_benunits).values

    # Person-level
    pe_person["would_claim_marriage_allowance"] = (
        generator.random(len(pe_person)) < marriage_allowance_rate
    )

    # Benefit unit-level — anchor on any adult in the benefit unit having
    # reported positive receipt in the FRS benefits table.
    pe_benunit["would_claim_child_benefit"] = assign_takeup_with_reported_anchors(
        generator.random(len(pe_benunit)),
        child_benefit_rate,
        reported_mask=_reported_benunit_mask("child_benefit_reported"),
    )
    pe_benunit["child_benefit_opts_out"] = (
        generator.random(len(pe_benunit)) < child_benefit_opts_out_rate
    )
    pe_benunit["would_claim_pc"] = assign_takeup_with_reported_anchors(
        generator.random(len(pe_benunit)),
        pension_credit_rate,
        reported_mask=_reported_benunit_mask("pension_credit_reported"),
    )
    pe_benunit["would_claim_uc"] = assign_takeup_with_reported_anchors(
        generator.random(len(pe_benunit)),
        universal_credit_rate,
        reported_mask=_reported_benunit_mask("universal_credit_reported"),
    )
    pe_benunit["would_claim_tfc"] = generator.random(len(pe_benunit)) < tfc_rate
    pe_benunit["would_claim_extended_childcare"] = (
        generator.random(len(pe_benunit)) < extended_childcare_rate
    )
    pe_benunit["would_claim_universal_childcare"] = (
        generator.random(len(pe_benunit)) < universal_childcare_rate
    )
    pe_benunit["would_claim_targeted_childcare"] = (
        generator.random(len(pe_benunit)) < targeted_childcare_rate
    )

    # Scottish Child Payment take-up at child level:
    # 97% for children under 6, 85% for children 6+
    # Source: gov.scot take-up rates publication Nov 2024
    # Note: We apply rates to all ages (not just <16) so that policy reforms
    # extending the age limit will have reasonable takeup assumptions.
    ages = pe_person["age"]
    scp_rate_per_person = np.where(ages < 6, scp_under_6_rate, scp_6_plus_rate)
    pe_person["would_claim_scp"] = (
        generator.random(len(pe_person)) < scp_rate_per_person
    )

    # Generate other stochastic variables using rates from parameter files
    tv_ownership_rate = load_parameter("stochastic", "tv_ownership_rate", year)
    tv_evasion_rate = load_parameter("stochastic", "tv_licence_evasion_rate", year)
    first_time_buyer_rate = load_parameter("stochastic", "first_time_buyer_rate", year)

    # Household-level: TV ownership
    pe_household["household_owns_tv"] = (
        generator.random(len(pe_household)) < tv_ownership_rate
    )

    # Household-level: TV licence evasion
    pe_household["would_evade_tv_licence_fee"] = (
        generator.random(len(pe_household)) < tv_evasion_rate
    )

    # Household-level: First home purchase
    pe_household["main_residential_property_purchased_is_first_home"] = (
        generator.random(len(pe_household)) < first_time_buyer_rate
    )

    # Person-level: Tie-breaking for higher earner (uniform random)
    pe_person["higher_earner_tie_break"] = generator.random(len(pe_person))

    # Person-level: Private school attendance random draw
    pe_person["attends_private_school_random_draw"] = generator.random(len(pe_person))

    # Generate extended childcare hours usage values with mean 15.019 and sd
    # 4.972
    extended_hours_values = generator.normal(15.019, 4.972, len(pe_benunit))
    # Clip values to be between 0 and 30 hours
    extended_hours_values = np.clip(extended_hours_values, 0, 30)

    # Add the maximum extended childcare hours usage
    pe_benunit["maximum_extended_childcare_hours_usage"] = extended_hours_values

    # Add marital status at the benefit unit level

    pe_benunit["is_married"] = frs["benunit"].famtypb2.isin([5, 7])

    # Stochastically set property_purchased based on UK housing transaction rate.
    # Previously defaulted to True in policyengine-uk, causing all households
    # to be charged SDLT as if they just bought their property (£370bn total).
    #
    # Sources:
    # - Transactions: HMRC 2024 - 1.1m/year
    #   https://www.gov.uk/government/statistics/monthly-property-transactions-completed-in-the-uk-with-value-40000-or-above
    # - Households: ONS 2024 - 28.6m
    #   https://www.ons.gov.uk/peoplepopulationandcommunity/birthsdeathsandmarriages/families/bulletins/familiesandhouseholds/2024
    # - Rate: 1.1m / 28.6m = 3.85%
    #
    # Verification against official SDLT revenue (2024-25):
    # - Official SDLT: £13.9bn (https://www.gov.uk/government/statistics/uk-stamp-tax-statistics)
    # - With fix (3.85%): £15.7bn (close to official)
    # - Without fix (100%): £370bn (26x too high)
    PROPERTY_PURCHASE_RATE = 0.0385
    pe_household["property_purchased"] = (
        np.random.random(len(pe_household)) < PROPERTY_PURCHASE_RATE
    )

    dataset = UKSingleYearDataset(
        person=pe_person,
        benunit=pe_benunit,
        household=pe_household,
        fiscal_year=year,
    )

    return dataset


if __name__ == "__main__":
    frs = create_frs(
        raw_frs_folder=STORAGE_FOLDER / "frs_2022_23",
        year=2022,
    )
    frs.save(STORAGE_FOLDER / "frs_2022.h5")