Clear cached UK calculations between Policy and Dynamic modifiers #717

"""Real public-data UK runs refresh each Policy/Dynamic operation boundary."""

import importlib

import numpy as np
import pandas as pd
import pytest
from microdf import MicroDataFrame

from policyengine.core import Dynamic, Policy, Simulation
from policyengine.tax_benefit_models.common.reform import (
    compile_reform_to_dynamic,
    compile_reform_to_policy,
)
from policyengine.tax_benefit_models.uk.datasets import PolicyEngineUKDataset, UKYearData
from policyengine.tax_benefit_models.uk.model import uk_latest

YEAR = 2027
BENEFIT = "child_benefit"
RESPONSE = "employment_income_behavioral_response"
INCOME = "income_elasticity_lsr"
SUBSTITUTION = "substitution_elasticity_lsr"
INCOME_PARAMETER = "gov.simulation.labour_supply_responses.income_elasticity"
SUBSTITUTION_PARAMETER = "gov.simulation.labour_supply_responses.substitution_elasticity"
ELDEST = "gov.hmrc.child_benefit.amount.eldest"
ELASTICITIES = {INCOME_PARAMETER: -0.05, SUBSTITUTION_PARAMETER: 0.15}
EXTRA = {"person": [RESPONSE, INCOME, SUBSTITUTION, "relative_wage_change"]}


def values(microsim, variable):
    return np.asarray(microsim.calculate(variable, YEAR)).copy()


@pytest.fixture
def public_dataset(monkeypatch, tmp_path):
    """Two actual public synthetic records; no materializer, engine or formula stub."""
    def frame(data, weight):
        return MicroDataFrame(pd.DataFrame(data), weights=weight)

    dataset = PolicyEngineUKDataset(
        name="public-modifier-family", description="One synthetic parent and child",
        year=YEAR, filepath=None,
        data=UKYearData(
            person=frame({
                "person_id": [101, 102],
                "person_benunit_id": [11, 11],
                "person_household_id": [1, 1],
                "person_weight": [1.0, 1.0],
                "age": [30, 2], "gender": ["FEMALE", "FEMALE"],
                "employment_income": [25_000.0, 0.0],
                "self_employment_income": [0.0, 0.0],
                "hours_worked": [1_500.0, 0.0],
                "months_since_last_birthday": [0, 0],
                "attends_private_school": [False, False],
            }, "person_weight"),
            benunit=frame({"benunit_id": [11], "benunit_weight": [1.0]}, "benunit_weight"),
            household=frame({
                "household_id": [1], "household_weight": [1.0],
                "region": ["LONDON"], "tenure_type": ["RENT_PRIVATELY"],
                "rent": [12_000.0], "council_tax": [1_500.0],
            }, "household_weight"),
        ),
    )
    original = {
        entity: pd.DataFrame(table).copy(deep=True)
        for entity, table in dataset.data.entity_data.items()
    }

    def unexpected_load(self):
        raise AssertionError("A supplied public in-memory dataset must not load from disk")

    monkeypatch.setattr(PolicyEngineUKDataset, "load", unexpected_load)
    model_module = importlib.import_module("policyengine.tax_benefit_models.uk.model")
    monkeypatch.setattr(model_module, "_output_dataset_filepath", lambda simulation: tmp_path / f"{simulation.id}.h5")
    yield dataset
    for entity, table in dataset.data.entity_data.items():
        pd.testing.assert_frame_equal(pd.DataFrame(table), original[entity])


def policy(parameters=None, callback=None):
    result = compile_reform_to_policy(parameters, year=YEAR, model_version=uk_latest) if parameters else Policy(name="Public callback")
    result.simulation_modifier = callback
    return result


def dynamic(parameters=None, callback=None):
    result = compile_reform_to_dynamic(parameters, year=YEAR, model_version=uk_latest) if parameters else Dynamic(name="Public dynamic")
    result.simulation_modifier = callback
    return result


def run(dataset, *, policy_value=None, dynamic_value=None):
    simulation = Simulation(
        dataset=dataset, tax_benefit_model_version=uk_latest,
        policy=policy_value, dynamic=dynamic_value, extra_variables=EXTRA,
    )
    simulation.run()
    assert simulation.output_dataset is not None
    return simulation


def benefit_parameter():
    from policyengine_uk import CountryTaxBenefitSystem

    amount = CountryTaxBenefitSystem().parameters.get_child(ELDEST)(str(YEAR))
    return {ELDEST: amount + 1_000 / 52}


@pytest.mark.parametrize("with_dynamic_callback", [False, True], ids=["parameter-only", "immediate-consumer"])
def test_late_dynamic_parameters_refresh_policy_warmed_benefit(public_dataset, with_dynamic_callback):
    captured = {}
    events = []

    def warm(microsim):
        events.append("policy")
        captured["before"] = values(microsim, BENEFIT)
        assert captured["before"][0] > 0
        values(microsim, "household_net_income")
        return microsim

    def consume(microsim):
        events.append("dynamic")
        np.testing.assert_allclose(values(microsim, BENEFIT), captured["before"] + 1_000, atol=0.01)
        return microsim

    actual = run(
        public_dataset, policy_value=policy(callback=warm),
        dynamic_value=dynamic(benefit_parameter(), consume if with_dynamic_callback else None),
    )
    cold = run(public_dataset, dynamic_value=dynamic(benefit_parameter()))
    expected = cold.output_dataset.data.benunit[BENEFIT]
    np.testing.assert_allclose(expected, captured["before"] + 1_000, atol=0.01)
    np.testing.assert_allclose(actual.output_dataset.data.benunit[BENEFIT], expected, atol=0.01)
    np.testing.assert_allclose(actual.output_dataset.data.household["household_net_income"], cold.output_dataset.data.household["household_net_income"], atol=0.01, rtol=1e-5)
    assert events == (["policy", "dynamic"] if with_dynamic_callback else ["policy"])