Model/applications/modelling_e2e/handler.py

"""SQS-triggered Lambda: fetch EPC (or predict) → run modelling → persist plan.

One SQS message = one batch of properties sharing a portfolio, scenario, and
(by caller convention) postcode. The handler reads ``property_ids``,
``portfolio_id``, ``scenario_id``, ``no_solar``, and ``dry_run`` from the
message body, fetches or predicts each property's EPC, runs the full modelling
pipeline (SAP10 → optimiser) via ``harness.console.run_modelling``, and
persists the resulting Plan via ``PostgresUnitOfWork`` in one atomic transaction
per property.

When no lodged EPC is found, EPC Prediction (Path 3, ADR-0031) synthesises one
from the postcode cohort. ``_cohort_cache`` is module-level so warm Lambda
containers re-processing the same postcode avoid redundant fetches.

All Measure Types are considered: pricing goes through
``catalogue_with_off_catalogue_overrides`` so the measures the live ``material``
catalogue cannot supply (``secondary_heating_removal``, the glazing and heating
gaps) are priced from the committed off-catalogue overlay instead of crashing.

DB engine is module-scoped so the connection pool is reused across warm
invocations (ADR-0012).
"""

from __future__ import annotations

import dataclasses
import io
import os
from collections.abc import Callable
from typing import Any, Optional, cast

import boto3
import pandas as pd  # pyright: ignore[reportMissingTypeStubs]
from sqlalchemy import Engine, text
from sqlmodel import Session

from datatypes.epc.domain.epc_property_data import (
    BuildingPartIdentifier,
    EpcPropertyData,
)
from domain.epc_prediction.comparable_properties import (
    ComparableProperty,
    select_comparables,
)
from domain.epc_prediction.epc_prediction import EpcPrediction
from domain.epc_prediction.prediction_target import (
    PredictionTarget,
    build_prediction_target,
)
from domain.geospatial.coordinates import Coordinates
from domain.geospatial.planning_restrictions import PlanningRestrictions
from domain.geospatial.spatial_reference import SpatialReference
from domain.property.property import Property, PropertyIdentity
from domain.property_baseline.calculator_rebaseliner import CalculatorRebaseliner
from domain.sap10_calculator.calculator import Sap10Calculator
from domain.tasks.tasks import Source
from harness.console import run_modelling
from orchestration.property_baseline_orchestrator import (
    PropertyBaselineOrchestrator,
)
from infrastructure.epc_client.epc_client_service import EpcClientService
from infrastructure.postcodes_io.postcodes_io_client import PostcodesIoClient
from infrastructure.postgres.config import PostgresConfig
from infrastructure.postgres.engine import make_engine
from infrastructure.solar.google_solar_api_client import (
    BuildingInsightsNotFoundError,
    GoogleSolarApiClient,
)
from applications.modelling_e2e.modelling_e2e_trigger_body import (
    ModellingE2ETriggerBody,
)
from repositories.comparable_properties.epc_comparable_properties_repository import (
    EpcComparablePropertiesRepository,
    SkippedCohortCert,
)
from repositories.geospatial.geospatial_s3_repository import (
    GeospatialS3Repository,
    ParquetReader,
)
from repositories.fuel_rates.fuel_rates_static_file_repository import (
    FuelRatesStaticFileRepository,
)
from repositories.postgres_unit_of_work import PostgresUnitOfWork
from repositories.product.composite_product_repository import (
    catalogue_with_off_catalogue_overrides,
)
from repositories.property.landlord_override_overlays import overlays_from
from repositories.property.override_backed_prediction_attributes_reader import (
    OverrideBackedPredictionAttributesReader,
)
from repositories.property.property_overrides_postgres_reader import (
    PropertyOverridesPostgresReader,
)
from repositories.scenario.scenario_postgres_repository import (
    ScenarioPostgresRepository,
)
from utilities.aws_lambda.task_handler import task_handler
from utilities.logger import setup_logger

_engine: Optional[Engine] = None
_cohort_cache: dict[str, list[ComparableProperty]] = {}
# Broadened (nearby-postcode) cohorts, keyed by (seed postcode, target property
# type): the early-stop walk depends on the type it is filling for, so two types
# in the same postcode must not share a cached result.
_nearby_cohort_cache: dict[tuple[str, str], list[ComparableProperty]] = {}

logger = setup_logger()


def _get_engine() -> Engine:
    global _engine
    if _engine is None:
        config = PostgresConfig.from_env(dict(os.environ))
        # Reduced pool for Lambda: 32 concurrent containers × 3 connections = 96 max,
        # vs the default 3+5=8 which would reach 256+ and exhaust RDS max_connections.
        # pool_size=2 covers the simultaneous read_session + UoW session per invocation.
        _engine = make_engine(dataclasses.replace(config, pool_size=2, max_overflow=1))
    return _engine


def _s3_parquet_reader() -> ParquetReader:
    bucket = os.environ["DATA_BUCKET"]

    def read(key: str) -> pd.DataFrame:
        s3: Any = cast(
            Any, boto3.client("s3")
        )  # pyright: ignore[reportUnknownMemberType]
        raw = cast(bytes, s3.get_object(Bucket=bucket, Key=key)["Body"].read())
        return pd.read_parquet(io.BytesIO(raw))  # type: ignore[return-value]

    return read


def _spatial_for(
    geospatial: GeospatialS3Repository, uprn: int
) -> Optional[SpatialReference]:
    try:
        return geospatial.spatial_for(uprn)
    except Exception:  # noqa: BLE001
        return None


def _solar_insights_for(
    solar_client: GoogleSolarApiClient, spatial: Optional[SpatialReference]
) -> Optional[dict[str, Any]]:
    if spatial is None or spatial.coordinates is None:
        return None
    try:
        return solar_client.get_building_insights(
            spatial.coordinates.longitude, spatial.coordinates.latitude
        )
    except BuildingInsightsNotFoundError:
        return None


def _dedupe_skipped(
    skipped: list[SkippedCohortCert],
) -> list[SkippedCohortCert]:
    """First occurrence of each skipped cert number (the same cert can appear in
    more than one postcode cohort across a batch)."""
    seen: set[str] = set()
    unique: list[SkippedCohortCert] = []
    for cert in skipped:
        if cert.certificate_number not in seen:
            seen.add(cert.certificate_number)
            unique.append(cert)
    return unique


def _predict_epc(
    *,
    property_id: int,
    uprn: int,
    postcode: str,
    portfolio_id: int,
    attributes_reader: OverrideBackedPredictionAttributesReader,
    coordinates: Optional[Coordinates],
    cohort_for: Callable[[str], list[ComparableProperty]],
    broaden: Callable[[PredictionTarget], list[ComparableProperty]],
    predictor: EpcPrediction,
) -> Optional[EpcPropertyData]:
    """Synthesise an EpcPropertyData for an EPC-less property from its postcode
    cohort (EPC Prediction Path 3, ADR-0031), or None when ineligible.

    When the property's own postcode holds no same-type comparables (a sparse
    postcode — e.g. the only flat among houses), the cohort is broadened to the
    real unit postcodes physically nearest it (``broaden``) before giving up.

    Returns None when property_type is unresolvable (hard cohort filter cannot
    fire) or when even the broadened cohort is empty after filtering.
    """
    attributes = attributes_reader.attributes_for(property_id)
    identity = PropertyIdentity(
        portfolio_id=portfolio_id, postcode=postcode, address="", uprn=uprn
    )
    target = build_prediction_target(identity, coordinates, attributes)
    if target is None:
        return None
    comparables = select_comparables(target, cohort_for(target.postcode))
    if not comparables.members:
        comparables = select_comparables(target, broaden(target))
    if not comparables.members:
        return None
    predicted = predictor.predict(target, comparables)
    if not any(
        part.identifier is BuildingPartIdentifier.MAIN
        for part in predicted.sap_building_parts
    ):
        return None
    return predicted


@task_handler(task_source="modelling_e2e", source=Source.PROPERTY)
def handler(body: dict[str, Any], context: Any) -> Optional[dict[str, Any]]:
    trigger = ModellingE2ETriggerBody.model_validate(body)
    property_ids = trigger.property_ids
    portfolio_id = trigger.portfolio_id
    scenario_id = trigger.scenario_id
    no_solar = trigger.no_solar
    dry_run = trigger.dry_run

    logger.info(
        f"start property_ids={property_ids} portfolio={portfolio_id} "
        f"scenario={scenario_id} no_solar={no_solar} dry_run={dry_run}"
    )

    engine = _get_engine()
    epc_client = EpcClientService(os.environ["OPEN_EPC_API_TOKEN"])
    geospatial = GeospatialS3Repository(_s3_parquet_reader())
    solar_client = GoogleSolarApiClient(os.environ["GOOGLE_SOLAR_API_KEY"])

    with engine.connect() as conn:
        uprn_rows = conn.execute(
            text("SELECT id, uprn FROM property WHERE id = ANY(:ids)"),
            {"ids": property_ids},
        ).fetchall()
        postcode_rows = conn.execute(
            text("SELECT id, postcode FROM property WHERE id = ANY(:ids)"),
            {"ids": property_ids},
        ).fetchall()

    uprns: dict[int, int] = {int(row[0]): int(row[1]) for row in uprn_rows}
    postcodes: dict[int, str] = {int(row[0]): (row[1] or "") for row in postcode_rows}

    overrides_reader = PropertyOverridesPostgresReader(lambda: Session(engine))
    prediction_attrs_reader = OverrideBackedPredictionAttributesReader(overrides_reader)
    comparables_repo = EpcComparablePropertiesRepository(
        epc_client, geospatial, nearby_postcodes=PostcodesIoClient()
    )
    predictor = EpcPrediction()

    def _get_cohort(postcode: str) -> list[ComparableProperty]:
        if postcode not in _cohort_cache:
            _cohort_cache[postcode] = (
                comparables_repo.candidates_for(postcode) if postcode else []
            )
        return _cohort_cache[postcode]

    def _broaden(target: PredictionTarget) -> list[ComparableProperty]:
        """The nearby-postcode cohort for a gated-out target — the real unit
        postcodes nearest it, walked until enough same-type comparables surface
        (ADR-0034). Memoised per (postcode, property_type) so co-located
        same-type misses share one walk."""
        key = (target.postcode, target.property_type)
        if key not in _nearby_cohort_cache:
            _nearby_cohort_cache[key] = (
                comparables_repo.candidates_near(
                    target.postcode,
                    target.coordinates,
                    enough=lambda c: c.epc.property_type == target.property_type,
                )
                if target.postcode
                else []
            )
        return _nearby_cohort_cache[key]

    # Re-establishes each lodged Property's Baseline Performance from the just-
    # persisted EPC (one UoW per property, committed after the Plan's). Predicted
    # Properties have no lodged figures, so they get no baseline (mirrors the e2e
    # runner and the ara_first_run Baseline stage).
    baseline_orchestrator = PropertyBaselineOrchestrator(
        unit_of_work=lambda: PostgresUnitOfWork(lambda: Session(engine)),
        rebaseliner=CalculatorRebaseliner(Sap10Calculator()),
        fuel_rates=FuelRatesStaticFileRepository(),
    )

    read_session = Session(engine)
    try:
        scenario = ScenarioPostgresRepository(read_session).get_many([scenario_id])[0]
        products = catalogue_with_off_catalogue_overrides(read_session)

        errors: list[int] = []

        for property_id in property_ids:
            try:
                uprn = uprns[property_id]
                postcode = postcodes.get(property_id, "")
                logger.info(f"property={property_id} uprn={uprn} postcode={postcode!r}")

                spatial = _spatial_for(geospatial, uprn)
                restrictions = (
                    spatial.restrictions
                    if spatial is not None
                    else PlanningRestrictions()
                )
                coordinates: Optional[Coordinates] = (
                    spatial.coordinates if spatial is not None else None
                )

                epc: Optional[EpcPropertyData] = epc_client.get_by_uprn(uprn)
                overrides = overlays_from(overrides_reader.overrides_for(property_id))
                predicted_epc: Optional[EpcPropertyData] = None

                if epc is not None:
                    logger.info(f"property={property_id} lodged EPC found")
                    effective_epc = Property(
                        identity=PropertyIdentity(
                            portfolio_id=portfolio_id,
                            postcode=postcode,
                            address="",
                            uprn=uprn,
                        ),
                        epc=epc,
                        landlord_overrides=overrides,
                    ).effective_epc
                else:
                    logger.info(
                        f"property={property_id} no lodged EPC — attempting prediction"
                    )
                    predicted_epc = _predict_epc(
                        property_id=property_id,
                        uprn=uprn,
                        postcode=postcode,
                        portfolio_id=portfolio_id,
                        attributes_reader=prediction_attrs_reader,
                        coordinates=coordinates,
                        cohort_for=_get_cohort,
                        broaden=_broaden,
                        predictor=predictor,
                    )
                    if predicted_epc is None:
                        raise ValueError(
                            f"no EPC for UPRN {uprn} and not predictable "
                            f"(unresolved property_type, or no same-type "
                            f"comparables in or near '{postcode}')"
                        )
                    effective_epc = Property(
                        identity=PropertyIdentity(
                            portfolio_id=portfolio_id,
                            postcode=postcode,
                            address="",
                            uprn=uprn,
                        ),
                        epc=None,
                        predicted_epc=predicted_epc,
                        landlord_overrides=overrides,
                    ).effective_epc

                solar_insights: Optional[dict[str, Any]] = (
                    None if no_solar else _solar_insights_for(solar_client, spatial)
                )

                # All Measure Types are considered: the off-catalogue overlay
                # (catalogue_with_off_catalogue_overrides) prices the measures the
                # live material catalogue cannot supply, so none need excluding.
                plan = run_modelling(
                    effective_epc,
                    planning_restrictions=restrictions,
                    solar_insights=solar_insights,
                    considered_measures=None,
                    products=products,
                    scenario=scenario,
                    print_table=False,
                )
                logger.info(
                    f"property={property_id} modelling complete "
                    f"measures={len(plan.measures)}"
                )

                if dry_run:
                    measure_types = (
                        ", ".join(m.measure_type for m in plan.measures) or "none"
                    )
                    logger.info(
                        f"[dry_run] property={property_id} "
                        f"measures=[{measure_types}] — skipping DB write"
                    )
                    continue

                with PostgresUnitOfWork(lambda: Session(engine)) as uow:
                    if epc is not None:
                        uow.epc.save(
                            epc, property_id=property_id, portfolio_id=portfolio_id
                        )
                    elif predicted_epc is not None:
                        # Persist the synthesised EPC in the predicted slot (ADR-0031),
                        # so the Baseline stage can re-hydrate it and downstream sees
                        # the picture the Plan was modelled from.
                        uow.epc.save(
                            predicted_epc,
                            property_id=property_id,
                            portfolio_id=portfolio_id,
                            source="predicted",
                        )
                    if spatial is not None:
                        uow.spatial.save(uprn, spatial)
                    if (
                        solar_insights is not None
                        and spatial is not None
                        and spatial.coordinates is not None
                    ):
                        uow.solar.save(
                            uprn,
                            longitude=spatial.coordinates.longitude,
                            latitude=spatial.coordinates.latitude,
                            insights=solar_insights,
                        )
                    uow.plan.save(
                        plan,
                        property_id=property_id,
                        scenario_id=scenario_id,
                        portfolio_id=portfolio_id,
                        is_default=scenario.is_default,
                    )
                    uow.property.mark_modelled(
                        property_id, has_recommendations=bool(plan.measures)
                    )
                    uow.commit()
                logger.info(f"property={property_id} plan saved")

                # Baseline Performance is re-established from the persisted EPC
                # (lodged or predicted), so it runs after the Plan UoW commits. By
                # here the property always has a persisted EPC — a property that
                # could be neither fetched nor predicted raised earlier.
                baseline_orchestrator.run([property_id])
                logger.info(f"property={property_id} baseline saved")

            except Exception as error:  # noqa: BLE001
                logger.error(
                    f"property={property_id}: {type(error).__name__}: {error}",
                    exc_info=True,
                )
                errors.append(property_id)

        # Cohort certs the mapper could not consume were skipped (not aborted on)
        # so prediction could proceed; surface them — with cert numbers — in the
        # subtask outputs so the mapper gaps can be closed later.
        skipped_certs: list[dict[str, str]] = [
            {"certificate_number": s.certificate_number, "error": s.error}
            for s in _dedupe_skipped(comparables_repo.skipped)
        ]
        if skipped_certs:
            logger.info(
                f"skipped {len(skipped_certs)} unmappable cohort cert(s): "
                f"{[s['certificate_number'] for s in skipped_certs]}"
            )

        if errors:
            message = f"failed property_ids: {errors}"
            if skipped_certs:
                message += f"; skipped_unmappable_cohort_certs: {skipped_certs}"
            raise RuntimeError(message)

        return {"skipped_unmappable_cohort_certs": skipped_certs} if skipped_certs else None
    finally:
        read_session.close()