Model/tests/domain/modelling/test_optimiser.py

"""Behaviour of the Optimiser core: a grouped-knapsack MILP over per-Option
role-1 scores (ADR-0016). Picks at most one Option per Recommendation (disjoint
groups, no cross-group constraints) to maximise total SAP gain subject to the
Scenario budget. This is the warm-start *signal* — the truthful figure comes
from the whole-package re-score + repair (a later slice); here we test the
selection with synthetic scores and no calculator.
"""

from __future__ import annotations

from typing import Sequence

from datatypes.epc.domain.epc_property_data import (
    BuildingPartIdentifier,
    EpcPropertyData,
)
from domain.modelling.optimisation.optimiser import (
    MeasureDependency,
    OptimisedPackage,
    ScoredOption,
    optimise,
    optimise_package,
)
from domain.modelling.scoring.package_scorer import Score
from domain.modelling.recommendation import Cost, MeasureOption
from domain.modelling.simulation import (
    BuildingPartOverlay,
    EpcSimulation,
    VentilationOverlay,
)
from tests.domain.sap10_calculator.worksheet._elmhurst_worksheet_000490 import (
    build_epc,
)


def _scored(measure_type: str, *, gain: float, cost: float) -> ScoredOption:
    return ScoredOption(
        option=MeasureOption(
            measure_type=measure_type,
            description=measure_type,
            overlay=EpcSimulation(),
            cost=Cost(total=cost, contingency_rate=0.0),
        ),
        sap_gain=gain,
    )


# Distinguishable overlays so the stub scorer can attribute a true gain per
# measure (wall / roof / floor) regardless of the role-1 signal.
_WALL_OVERLAY = EpcSimulation(
    building_parts={
        BuildingPartIdentifier.MAIN: BuildingPartOverlay(wall_insulation_type=2)
    }
)
_ROOF_OVERLAY = EpcSimulation(
    building_parts={
        BuildingPartIdentifier.MAIN: BuildingPartOverlay(roof_insulation_thickness=300)
    }
)
_FLOOR_OVERLAY = EpcSimulation(
    building_parts={
        BuildingPartIdentifier.MAIN: BuildingPartOverlay(floor_insulation_thickness=100)
    }
)


def _scored_overlay(
    measure_type: str, *, gain: float, cost: float, overlay: EpcSimulation
) -> ScoredOption:
    return ScoredOption(
        option=MeasureOption(
            measure_type=measure_type,
            description=measure_type,
            overlay=overlay,
            cost=Cost(total=cost, contingency_rate=0.0),
        ),
        sap_gain=gain,
    )


class _StubScorer:
    """A deterministic stand-in for PackageScorer: the package SAP is a base
    plus a fixed *true* gain per measure present (by overlay field), decoupled
    from the role-1 signal — so the repair loop is exercised without the
    calculator (ADR-0016)."""

    def __init__(self, *, base: float, wall: float, roof: float, floor: float) -> None:
        self._base = base
        self._wall = wall
        self._roof = roof
        self._floor = floor

    def score(
        self, baseline: EpcPropertyData, simulations: Sequence[EpcSimulation]
    ) -> Score:
        sap = self._base
        for sim in simulations:
            part = sim.building_parts[BuildingPartIdentifier.MAIN]
            if part.wall_insulation_type is not None:
                sap += self._wall
            if part.roof_insulation_thickness is not None:
                sap += self._roof
            if part.floor_insulation_thickness is not None:
                sap += self._floor
        return Score(sap_continuous=sap, co2_kg_per_yr=0.0, primary_energy_kwh_per_yr=0.0)


def _selected_types(selection: list[ScoredOption]) -> set[str]:
    return {scored.option.measure_type for scored in selection}


def test_grouped_knapsack_maximises_gain_within_budget() -> None:
    # Arrange — wall group has two mutually-exclusive options; roof + floor one
    # each. EWI has the best gain but is unaffordable alongside the rest.
    groups: list[list[ScoredOption]] = [
        [
            _scored("external_wall_insulation", gain=10.0, cost=8000.0),
            _scored("cavity_wall_insulation", gain=6.0, cost=1000.0),
        ],
        [_scored("loft_insulation", gain=4.0, cost=1500.0)],
        [_scored("suspended_floor_insulation", gain=3.0, cost=2000.0)],
    ]

    # Act
    selection: list[ScoredOption] = optimise(groups, budget=5000.0)

    # Assert — cavity + loft + floor (cost 4500, gain 13) beats any package
    # containing the 8000 EWI option within the 5000 budget.
    assert _selected_types(selection) == {
        "cavity_wall_insulation",
        "loft_insulation",
        "suspended_floor_insulation",
    }


def test_picks_at_most_one_option_per_group() -> None:
    # Arrange — both wall options are individually affordable.
    groups: list[list[ScoredOption]] = [
        [
            _scored("external_wall_insulation", gain=10.0, cost=2000.0),
            _scored("cavity_wall_insulation", gain=6.0, cost=1000.0),
        ],
    ]

    # Act
    selection: list[ScoredOption] = optimise(groups, budget=10000.0)

    # Assert — never both treatments of the same wall; the higher-gain one wins.
    assert len(selection) == 1
    assert _selected_types(selection) == {"external_wall_insulation"}


def test_no_budget_picks_the_best_option_in_every_group() -> None:
    # Arrange
    groups: list[list[ScoredOption]] = [
        [
            _scored("external_wall_insulation", gain=10.0, cost=8000.0),
            _scored("cavity_wall_insulation", gain=6.0, cost=1000.0),
        ],
        [_scored("loft_insulation", gain=4.0, cost=1500.0)],
    ]

    # Act — None budget = unconstrained.
    selection: list[ScoredOption] = optimise(groups, budget=None)

    # Assert
    assert _selected_types(selection) == {
        "external_wall_insulation",
        "loft_insulation",
    }


def test_budget_too_small_for_any_option_selects_nothing() -> None:
    # Arrange
    groups: list[list[ScoredOption]] = [
        [_scored("cavity_wall_insulation", gain=6.0, cost=1000.0)],
        [_scored("loft_insulation", gain=4.0, cost=1500.0)],
    ]

    # Act
    selection: list[ScoredOption] = optimise(groups, budget=500.0)

    # Assert — nothing affordable; selecting none is the optimum.
    assert selection == []


def test_no_groups_selects_nothing() -> None:
    # Act / Assert
    assert optimise([], budget=10000.0) == []


def test_within_budget_partial_selection_prefers_the_higher_gain_option() -> None:
    # Arrange — only one of the two fits the budget; pick the affordable best.
    groups: list[list[ScoredOption]] = [
        [_scored("external_wall_insulation", gain=10.0, cost=8000.0)],
        [_scored("loft_insulation", gain=4.0, cost=1500.0)],
    ]

    # Act
    selection: list[ScoredOption] = optimise(groups, budget=2000.0)

    # Assert — EWI is unaffordable; loft alone is the best within £2000.
    assert _selected_types(selection) == {"loft_insulation"}


def test_repair_adds_an_untreated_group_option_to_close_the_undershoot() -> None:
    # Arrange — role-1 under-counts roof (signal 0 → warm-start skips it), but
    # its true re-scored gain (+4) is what closes the target.
    groups: list[list[ScoredOption]] = [
        [_scored_overlay("cavity_wall_insulation", gain=10.0, cost=1000.0, overlay=_WALL_OVERLAY)],
        [_scored_overlay("loft_insulation", gain=0.0, cost=1000.0, overlay=_ROOF_OVERLAY)],
        [_scored_overlay("suspended_floor_insulation", gain=8.0, cost=1000.0, overlay=_FLOOR_OVERLAY)],
    ]
    scorer = _StubScorer(base=40.0, wall=5.0, roof=4.0, floor=3.0)

    # Act
    package: OptimisedPackage = optimise_package(
        groups=groups,
        scorer=scorer,
        baseline_epc=build_epc(),
        budget=5000.0,
        target_sap=50.0,
    )

    # Assert — warm-start took wall+floor (re-score 48 < 50); repair added the
    # roof (true +4) to reach 52, the truthful package total.
    types = {scored.option.measure_type for scored in package.selected}
    assert "loft_insulation" in types
    assert types == {
        "cavity_wall_insulation",
        "suspended_floor_insulation",
        "loft_insulation",
    }
    assert abs(package.score.sap_continuous - 52.0) <= 1e-9


def test_no_target_returns_the_warm_start_package_without_repair() -> None:
    # Arrange
    groups: list[list[ScoredOption]] = [
        [_scored_overlay("cavity_wall_insulation", gain=10.0, cost=1000.0, overlay=_WALL_OVERLAY)],
    ]
    scorer = _StubScorer(base=40.0, wall=5.0, roof=4.0, floor=3.0)

    # Act
    package: OptimisedPackage = optimise_package(
        groups=groups,
        scorer=scorer,
        baseline_epc=build_epc(),
        budget=None,
        target_sap=None,
    )

    # Assert — no target → no repair; warm-start package re-scored as the truth.
    assert {s.option.measure_type for s in package.selected} == {
        "cavity_wall_insulation"
    }
    assert abs(package.score.sap_continuous - 45.0) <= 1e-9


def test_repair_stops_when_no_affordable_improving_option_remains() -> None:
    # Arrange — the only untreated-group option costs more than the budget left.
    groups: list[list[ScoredOption]] = [
        [_scored_overlay("cavity_wall_insulation", gain=10.0, cost=1000.0, overlay=_WALL_OVERLAY)],
        [_scored_overlay("loft_insulation", gain=0.0, cost=5000.0, overlay=_ROOF_OVERLAY)],
    ]
    scorer = _StubScorer(base=40.0, wall=5.0, roof=4.0, floor=3.0)

    # Act
    package: OptimisedPackage = optimise_package(
        groups=groups,
        scorer=scorer,
        baseline_epc=build_epc(),
        budget=1000.0,
        target_sap=50.0,
    )

    # Assert — wall only (re-score 45 < 50); roof unaffordable, so repair stops
    # at the best achievable package rather than overspending.
    assert {s.option.measure_type for s in package.selected} == {
        "cavity_wall_insulation"
    }
    assert abs(package.score.sap_continuous - 45.0) <= 1e-9


# --- Measure Dependency injection (ADR-0016) -------------------------------

_VENT_OVERLAY = EpcSimulation(
    ventilation=VentilationOverlay(
        mechanical_ventilation_kind="EXTRACT_OR_PIV_OUTSIDE"
    )
)


class _VentStubScorer:
    """A stub that adds a fixed gain per wall overlay present and a fixed
    (negative) `vent` contribution when a ventilation overlay is present —
    so the Measure Dependency's effect on the truthful package total and the
    repair decision is exercised without the calculator."""

    def __init__(self, *, base: float, wall: float, roof: float, vent: float) -> None:
        self._base = base
        self._wall = wall
        self._roof = roof
        self._vent = vent

    def score(
        self, baseline: EpcPropertyData, simulations: Sequence[EpcSimulation]
    ) -> Score:
        sap = self._base
        for sim in simulations:
            if sim.ventilation is not None:
                sap += self._vent
            for part in sim.building_parts.values():
                if part.wall_insulation_type is not None:
                    sap += self._wall
                if part.roof_insulation_thickness is not None:
                    sap += self._roof
        return Score(sap_continuous=sap, co2_kg_per_yr=0.0, primary_energy_kwh_per_yr=0.0)


def _ventilation_dependency(*, cost: float) -> MeasureDependency:
    """A forced 'fabric requires ventilation' edge for the tests."""
    return MeasureDependency(
        triggers=frozenset({"cavity_wall_insulation", "external_wall_insulation"}),
        required=ScoredOption(
            option=MeasureOption(
                measure_type="mechanical_ventilation",
                description="mechanical_ventilation",
                overlay=_VENT_OVERLAY,
                cost=Cost(total=cost, contingency_rate=0.0),
            ),
            sap_gain=0.0,
        ),
    )


def test_dependency_injected_when_a_trigger_measure_is_selected() -> None:
    # Arrange — the wall is selected, so its ventilation dependency must be
    # injected before the re-score; ventilation never competes in the pool.
    groups: list[list[ScoredOption]] = [
        [_scored_overlay("cavity_wall_insulation", gain=10.0, cost=1000.0, overlay=_WALL_OVERLAY)],
    ]
    scorer = _VentStubScorer(base=40.0, wall=5.0, roof=4.0, vent=-2.0)

    # Act
    package: OptimisedPackage = optimise_package(
        groups=groups,
        scorer=scorer,
        baseline_epc=build_epc(),
        budget=None,
        target_sap=None,
        dependencies=[_ventilation_dependency(cost=900.0)],
    )

    # Assert — ventilation is in the package and its negative contribution lands
    # in the truthful total: 40 base + 5 wall − 2 ventilation = 43.
    assert _selected_types(package.selected) == {
        "cavity_wall_insulation",
        "mechanical_ventilation",
    }
    assert abs(package.score.sap_continuous - 43.0) <= 1e-9


def test_dependency_not_injected_without_a_trigger_measure() -> None:
    # Arrange — only loft is selected; the wall-triggered ventilation dependency
    # must not fire.
    groups: list[list[ScoredOption]] = [
        [_scored_overlay("loft_insulation", gain=4.0, cost=1000.0, overlay=_ROOF_OVERLAY)],
    ]
    scorer = _VentStubScorer(base=40.0, wall=5.0, roof=4.0, vent=-2.0)

    # Act
    package: OptimisedPackage = optimise_package(
        groups=groups,
        scorer=scorer,
        baseline_epc=build_epc(),
        budget=None,
        target_sap=None,
        dependencies=[_ventilation_dependency(cost=900.0)],
    )

    # Assert — no trigger, no ventilation; 40 base + 4 roof = 44.
    assert _selected_types(package.selected) == {"loft_insulation"}
    assert abs(package.score.sap_continuous - 44.0) <= 1e-9


def test_dependency_is_forced_even_when_it_pushes_over_budget() -> None:
    # Arrange — the budget covers the wall but not the forced ventilation;
    # ventilation is mandatory-when-triggered, so it is injected regardless.
    groups: list[list[ScoredOption]] = [
        [_scored_overlay("cavity_wall_insulation", gain=10.0, cost=1000.0, overlay=_WALL_OVERLAY)],
    ]
    scorer = _VentStubScorer(base=40.0, wall=5.0, roof=4.0, vent=-2.0)

    # Act — budget exactly covers the wall; ventilation (£900) overspends.
    package: OptimisedPackage = optimise_package(
        groups=groups,
        scorer=scorer,
        baseline_epc=build_epc(),
        budget=1000.0,
        target_sap=None,
        dependencies=[_ventilation_dependency(cost=900.0)],
    )

    # Assert — forced in despite the overspend.
    assert "mechanical_ventilation" in _selected_types(package.selected)


def test_injected_ventilation_penalty_drives_extra_repair() -> None:
    # Arrange — wall (+5) injects ventilation (−2): re-score 43 < target 46.
    # Repair adds the roof (true +4) to reach 47, paying for the ventilation
    # penalty out of the budget the dependency's cost has already eaten into.
    groups: list[list[ScoredOption]] = [
        [_scored_overlay("cavity_wall_insulation", gain=10.0, cost=1000.0, overlay=_WALL_OVERLAY)],
        [_scored_overlay("loft_insulation", gain=0.0, cost=1000.0, overlay=_ROOF_OVERLAY)],
    ]
    scorer = _VentStubScorer(base=40.0, wall=5.0, roof=4.0, vent=-2.0)

    # Act
    package: OptimisedPackage = optimise_package(
        groups=groups,
        scorer=scorer,
        baseline_epc=build_epc(),
        budget=5000.0,
        target_sap=46.0,
        dependencies=[_ventilation_dependency(cost=900.0)],
    )

    # Assert — repair pulled the roof in to clear the target net of ventilation:
    # 40 + 5 wall − 2 vent + 4 roof = 47.
    assert _selected_types(package.selected) == {
        "cavity_wall_insulation",
        "loft_insulation",
        "mechanical_ventilation",
    }
    assert abs(package.score.sap_continuous - 47.0) <= 1e-9