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.optimiser import (
    OptimisedPackage,
    ScoredOption,
    optimise,
    optimise_package,
)
from domain.modelling.package_scorer import Score
from domain.modelling.recommendation import Cost, MeasureOption
from domain.modelling.simulation import BuildingPartOverlay, EpcSimulation
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