Model/orchestration/modelling_orchestrator.py

from __future__ import annotations

from collections.abc import Callable
from typing import Final, Optional

from datatypes.epc.domain.epc import Epc
from datatypes.epc.domain.epc_property_data import EpcPropertyData
from domain.billing.bill import Bill, EnergyBreakdown
from domain.billing.bill_derivation import BillDerivation
from domain.modelling.generators.floor_recommendation import recommend_floor_insulation
from domain.modelling.optimisation.measure_dependency import ventilation_dependency
from domain.modelling.optimisation.optimiser import (
    MeasureDependency,
    OptimisedPackage,
    ScoredOption,
    optimise_package,
)
from domain.modelling.scoring.package_scorer import PackageScorer, Score
from domain.modelling.plan import Plan, PlanMeasure
from domain.modelling.recommendation import MeasureOption, Recommendation
from domain.modelling.generators.roof_recommendation import recommend_loft_insulation
from domain.modelling.scenario import Scenario
from domain.modelling.scoring.scoring import (
    MeasureImpact,
    cascade_scores,
    independent_option_impacts,
    marginals_from_scores,
)
from domain.modelling.generators.wall_recommendation import recommend_cavity_wall
from domain.sap10_calculator.calculator import SapCalculator
from repositories.fuel_rates.fuel_rates_repository import FuelRatesRepository
from repositories.product.product_repository import ProductRepository
from repositories.unit_of_work import UnitOfWork

# The PortfolioGoal value that targets a SAP band (cf.
# backend.app.db.models.portfolio.PortfolioGoal.INCREASING_EPC). Other goals
# (Energy Savings, Reducing CO2 emissions) don't yet set a SAP repair target —
# the optimiser just maximises SAP gain within budget for them (later slice).
_INCREASING_EPC_GOAL: Final[str] = "Increasing EPC"

# Best-practice install sequence for the role-3 attribution cascade (ADR-0016):
# walls → roof → ventilation → floor, per the legacy `Recommendations` class.
# Ventilation sits after the fabric that triggers it so its (negative) marginal
# is attributed against the insulated envelope.
_BEST_PRACTICE_ORDER: Final[tuple[str, ...]] = (
    "cavity_wall_insulation",
    "external_wall_insulation",
    "internal_wall_insulation",
    "loft_insulation",
    "mechanical_ventilation",
    "suspended_floor_insulation",
    "solid_floor_insulation",
)


class ModellingOrchestrator:
    """Stage 3 — scores each baselined Property against its Scenarios into Plans
    and persists them (CONTEXT.md: Modelling; ADR-0011 / ADR-0012 / ADR-0016 /
    ADR-0017).

    Runs the whole batch in **one** Unit of Work and commits once. For each
    (Property × Scenario) it reads the Property's Effective EPC and the Scenario
    through repos, generates the candidate Recommendations (wall / roof /
    floor), scores each Option independently (role 1), runs the grouped-knapsack
    Optimiser + whole-package re-score + greedy repair toward the Scenario's SAP
    target (role 2, ADR-0016), attributes each selected measure via the
    best-practice marginal cascade (role 3), and persists a **Plan** with its
    **Plan Measures**. Single-phase — multi-phase is deferred (ADR-0005).

    Reads only through repos and threads only IDs (`property_ids`,
    `scenario_ids`, `portfolio_id`) — never an in-memory hand-off from Baseline
    (ADR-0011). The injected `SapCalculator` is the scoring-engine seam.
    """

    def __init__(
        self,
        *,
        unit_of_work: Callable[[], UnitOfWork],
        calculator: SapCalculator,
        fuel_rates: FuelRatesRepository,
    ) -> None:
        self._unit_of_work = unit_of_work
        self._calculator = calculator
        self._fuel_rates = fuel_rates

    def run(
        self, property_ids: list[int], scenario_ids: list[int], portfolio_id: int
    ) -> None:
        scorer = PackageScorer(self._calculator)
        # Resolve Fuel Rates once and reuse the BillDerivation across the batch,
        # so every baseline/post bill is priced at the same snapshot (ADR-0014).
        bill_derivation = BillDerivation(self._fuel_rates.get_current())
        with self._unit_of_work() as uow:
            properties = uow.property.get_many(property_ids)
            scenarios: list[Scenario] = uow.scenario.get_many(scenario_ids)
            for property_id, prop in zip(property_ids, properties, strict=True):
                effective_epc: EpcPropertyData = prop.effective_epc
                for scenario in scenarios:
                    plan = self._plan_for(
                        scorer,
                        bill_derivation,
                        effective_epc,
                        uow.product,
                        scenario,
                        current_market_value=prop.current_market_value,
                    )
                    uow.plan.save(
                        plan,
                        property_id=property_id,
                        scenario_id=scenario.id,
                        portfolio_id=portfolio_id,
                        is_default=scenario.is_default,
                    )
            uow.commit()

    def _plan_for(
        self,
        scorer: PackageScorer,
        bill_derivation: BillDerivation,
        effective_epc: EpcPropertyData,
        products: ProductRepository,
        scenario: Scenario,
        *,
        current_market_value: Optional[float],
    ) -> Plan:
        """Generate → score → optimise → re-score/repair → attribute → bill →
        assemble the Plan for one Property + Scenario."""
        groups: list[list[ScoredOption]] = _scored_candidate_groups(
            scorer, effective_epc, products
        )
        # Forced Measure Dependencies (ventilation) are excluded from the pool
        # but injected into the package before the re-score (ADR-0016).
        dependencies: list[MeasureDependency] = _measure_dependencies(
            effective_epc, products
        )
        package: OptimisedPackage = optimise_package(
            groups=groups,
            scorer=scorer,
            baseline_epc=effective_epc,
            budget=scenario.budget,
            target_sap=_target_sap(scenario),
            dependencies=dependencies,
        )

        # Role-3 attribution: re-apply the *selected* set in best-practice order
        # so each measure's marginal telescopes to the truthful package total.
        ordered: list[MeasureOption] = sorted(
            (scored.option for scored in package.selected), key=_best_practice_key
        )
        # Score the baseline + every cumulative prefix once (cascade[0] is the
        # baseline, cascade[-1] the whole package), then reuse those Scores for
        # both the marginal attribution and the per-measure bill cascade.
        cascade: list[Score] = cascade_scores(
            scorer, effective_epc, [option.overlay for option in ordered]
        )
        impacts: list[MeasureImpact] = marginals_from_scores(cascade)
        # Bill every prefix at one Fuel Rates snapshot; consecutive Bill deltas
        # are each measure's marginal energy/cost saving — negative for
        # ventilation — telescoping exactly to the Plan totals (ADR-0014). The
        # Plan's baseline/post Bills are the cascade endpoints, so the
        # per-measure savings and the headline savings share one source.
        bills: list[Bill] = [_bill_for(bill_derivation, score) for score in cascade]
        measures: tuple[PlanMeasure, ...] = tuple(
            _plan_measure(option, impact, before, after)
            for option, impact, before, after in zip(
                ordered, impacts, bills[:-1], bills[1:], strict=True
            )
        )
        return Plan(
            measures=measures,
            baseline=cascade[0],
            post_retrofit=package.score,
            baseline_bill=bills[0],
            post_bill=bills[-1],
            current_market_value=current_market_value,
        )


def _bill_for(bill_derivation: BillDerivation, score: Score) -> Bill:
    """Derive the annual Bill for a scored end-state, pricing the delivered
    energy off the Score's SapResult. The real PackageScorer always attaches the
    SapResult; a missing one is a wiring error, so raise rather than bill at a
    default (ADR-0014)."""
    if score.sap_result is None:
        raise ValueError(
            "cannot derive a bill: the Score carries no SapResult to price"
        )
    return bill_derivation.derive(EnergyBreakdown.from_sap_result(score.sap_result))


def _candidate_recommendations(
    effective_epc: EpcPropertyData, products: ProductRepository
) -> list[Recommendation]:
    """Run every fabric Recommendation Generator; keep the ones that apply."""
    generators = (
        recommend_cavity_wall,
        recommend_loft_insulation,
        recommend_floor_insulation,
    )
    found = (generator(effective_epc, products) for generator in generators)
    return [recommendation for recommendation in found if recommendation is not None]


def _measure_dependencies(
    effective_epc: EpcPropertyData, products: ProductRepository
) -> list[MeasureDependency]:
    """The forced Measure Dependencies for this Property — currently just
    ventilation, suppressed when the dwelling is already mechanically
    ventilated (ADR-0016)."""
    dependency: Optional[MeasureDependency] = ventilation_dependency(
        effective_epc, products
    )
    return [dependency] if dependency is not None else []


def _scored_candidate_groups(
    scorer: PackageScorer,
    effective_epc: EpcPropertyData,
    products: ProductRepository,
) -> list[list[ScoredOption]]:
    """One group per Recommendation: each Option scored independently against
    the baseline (role-1 warm-start signal, ADR-0016)."""
    groups: list[list[ScoredOption]] = []
    for recommendation in _candidate_recommendations(effective_epc, products):
        options = list(recommendation.options)
        impacts: list[MeasureImpact] = independent_option_impacts(
            scorer, effective_epc, options
        )
        groups.append(
            [
                ScoredOption(option=option, sap_gain=impact.sap_points)
                for option, impact in zip(options, impacts, strict=True)
            ]
        )
    return groups


def _target_sap(scenario: Scenario) -> Optional[float]:
    """The SAP rating the Optimiser repairs toward — the floor of the goal
    band for an INCREASING_EPC goal, else None (no SAP target)."""
    if scenario.goal != _INCREASING_EPC_GOAL:
        return None
    return float(Epc(scenario.goal_value).sap_lower_bound())


def _best_practice_key(option: MeasureOption) -> int:
    try:
        return _BEST_PRACTICE_ORDER.index(option.measure_type)
    except ValueError:
        return len(_BEST_PRACTICE_ORDER)


def _plan_measure(
    option: MeasureOption, impact: MeasureImpact, before: Bill, after: Bill
) -> PlanMeasure:
    """Assemble a Plan Measure, attributing this measure's marginal bill saving
    as the delta between the running package Bill before and after it (delivered
    kWh and £). Signed so positive is a saving; ventilation is negative."""
    if option.cost is None:
        raise ValueError(
            f"measure option {option.measure_type!r} has no cost; cannot persist"
        )
    return PlanMeasure(
        measure_type=option.measure_type,
        description=option.description,
        cost=option.cost,
        impact=impact,
        kwh_savings=before.total_consumption_kwh - after.total_consumption_kwh,
        energy_cost_savings=before.total_gbp - after.total_gbp,
        material_id=option.material_id,
    )