Model/repositories/comparable_properties/epc_comparable_properties_repository.py

"""EPC-API + geospatial adapter for the ComparableProperties port (ADR-0031).

Assembles a postcode's candidate cohort: the EPC search lists the postcode's
lodged certs, each is fetched + mapped to `EpcPropertyData`, and the certs' UPRNs
are resolved to coordinates in one batched geospatial read (closely-numbered
UPRNs share a partition). Register metadata the cert itself doesn't carry
(address, registration date) is threaded off the search row.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import date
from typing import Callable, Optional, Protocol

from datatypes.epc.domain.epc_property_data import EpcPropertyData
from datatypes.epc.search.epc_search_result import EpcSearchResult
from domain.epc_prediction.comparable_properties import ComparableProperty
from domain.geospatial.coordinates import Coordinates
from repositories.comparable_properties.comparable_properties_repository import (
    ComparablePropertiesRepository,
)

# The same default floor `select_comparables` uses: keep walking nearby postcodes
# until this many candidates match, so the broadened cohort is big enough for the
# downstream relax ladder rather than stopping at the first stray match.
_DEFAULT_MINIMUM_COHORT = 5


class CohortEpcClient(Protocol):
    """The slice of the EPC-API client the cohort fetch needs (e.g.
    `EpcClientService`)."""

    def search_by_postcode(self, postcode: str) -> list[EpcSearchResult]: ...

    def get_by_certificate_number(self, cert_num: str) -> EpcPropertyData: ...


class CohortGeospatial(Protocol):
    """The geospatial slice the cohort fetch needs — batch UPRN→coordinate."""

    def coordinates_for_uprns(
        self, uprns: list[int]
    ) -> dict[int, Coordinates]: ...


class NearbyPostcodes(Protocol):
    """Resolves the real unit postcodes physically near a seed postcode (e.g.
    `PostcodesIoClient`). The gov EPC API cannot search by radius, so this is how
    the cohort reaches beyond the target's own postcode (ADR-0034)."""

    def nearby(
        self, postcode: str, coordinates: Optional[Coordinates] = None
    ) -> list[str]: ...


logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SkippedCohortCert:
    """A postcode-cohort cert the mapper could not consume, so it was excluded
    from the cohort rather than sinking the whole prediction. The certificate
    number + error are captured so the gap surfaces (subtask outputs + logs) and
    can be closed by extending the mapper later (ADR-0031)."""

    certificate_number: str
    error: str


class EpcComparablePropertiesRepository(ComparablePropertiesRepository):
    def __init__(
        self,
        epc_client: CohortEpcClient,
        geospatial: CohortGeospatial,
        nearby_postcodes: Optional[NearbyPostcodes] = None,
        widen_nearby_postcodes: Optional[NearbyPostcodes] = None,
    ) -> None:
        self._epc_client = epc_client
        self._geospatial = geospatial
        self._nearby_postcodes = nearby_postcodes
        # A second, wider-reaching `NearbyPostcodes` source, consulted only when
        # the normal-radius walk fails to reach `minimum` same-type matches — one
        # extra step outward before genuinely giving up (ADR-0034 amendment). A
        # single configured step today; the same pattern chains to further steps
        # later without changing this class's contract.
        self._widen_nearby_postcodes = widen_nearby_postcodes
        # Cohort certs skipped because they are not yet mappable. Accumulates
        # across every postcode the instance serves; the caller reads it after
        # the run to report the mapper gaps (see modelling_e2e handler).
        self.skipped: list[SkippedCohortCert] = []

    def candidates_for(self, postcode: str) -> list[ComparableProperty]:
        results: list[EpcSearchResult] = self._epc_client.search_by_postcode(
            postcode
        )
        uprns: list[int] = [r.uprn for r in results if r.uprn is not None]
        coordinates: dict[int, Coordinates] = self._geospatial.coordinates_for_uprns(
            uprns
        )
        cohort: list[ComparableProperty] = []
        for result in results:
            comparable = self._comparable_or_skip(result, coordinates)
            if comparable is not None:
                cohort.append(comparable)
        return cohort

    def candidates_near(
        self,
        postcode: str,
        coordinates: Optional[Coordinates] = None,
        *,
        enough: Optional[Callable[[ComparableProperty], bool]] = None,
        minimum: int = _DEFAULT_MINIMUM_COHORT,
    ) -> list[ComparableProperty]:
        """The broadened cohort: candidates drawn from the real unit postcodes
        nearest ``postcode`` (ADR-0034), for when the target's own postcode holds
        no same-type comparables. Postcodes are visited nearest first and each
        candidate is deduped by certificate number across them.

        ``enough`` lets the caller stop the walk early — once ``minimum``
        candidates satisfy it (e.g. they match the target's property type) the
        remaining, further-away postcodes are not fetched, so a dense area
        resolves in one or two searches instead of the whole radius. Without a
        configured ``NearbyPostcodes`` source this degrades to the seed postcode
        alone.

        If the normal-radius walk still falls short of ``minimum`` matches and a
        ``widen_nearby_postcodes`` source is configured, the walk is retried
        once against that wider source before giving up — a single extra step
        outward, not an unbounded expansion."""
        postcodes = self._nearby(self._nearby_postcodes, postcode, coordinates)
        candidates, matches = self._walk(postcodes, enough, minimum)

        if (
            enough is not None
            and matches < minimum
            and self._widen_nearby_postcodes is not None
        ):
            wider_postcodes = self._nearby(
                self._widen_nearby_postcodes, postcode, coordinates
            )
            candidates, _ = self._walk(wider_postcodes, enough, minimum)

        return candidates

    def _nearby(
        self,
        source: Optional[NearbyPostcodes],
        postcode: str,
        coordinates: Optional[Coordinates],
    ) -> list[str]:
        return source.nearby(postcode, coordinates) if source is not None else [postcode]

    def _walk(
        self,
        postcodes: list[str],
        enough: Optional[Callable[[ComparableProperty], bool]],
        minimum: int,
    ) -> tuple[list[ComparableProperty], int]:
        candidates: list[ComparableProperty] = []
        seen_certs: set[str] = set()
        matches = 0
        for nearby_postcode in postcodes:
            for candidate in self.candidates_for(nearby_postcode):
                if candidate.certificate_number in seen_certs:
                    continue
                seen_certs.add(candidate.certificate_number)
                candidates.append(candidate)
                if enough is not None and enough(candidate):
                    matches += 1
            if enough is not None and matches >= minimum:
                break
        return candidates, matches

    # Mapper-shape errors: the cert's lodged data does not fit the schema/mapper.
    # `ValueError` — missing required field / unmapped code (`UnmappedApiCode`);
    # `AttributeError` — a field lodged in the wrong shape (e.g. `photovoltaic_supply`
    # as a list where a `PhotovoltaicSupply` object is expected → "'list' object has
    # no attribute 'none_or_no_details'"); `KeyError`/`TypeError` — analogous
    # structural mismatches. Transient API failures (`EpcApiError` &c.) subclass
    # `Exception` directly, NOT these, so they still propagate and fail the run loud.
    _UNMAPPABLE_CERT_ERRORS = (ValueError, AttributeError, KeyError, TypeError)

    def _comparable_or_skip(
        self, result: EpcSearchResult, coordinates: dict[int, Coordinates]
    ) -> Optional[ComparableProperty]:
        """Map one cohort cert, or record + skip it when the mapper raises a
        mapper-shape error (see ``_UNMAPPABLE_CERT_ERRORS``). One unmappable cert
        must NOT abort the whole cohort; transient API errors still propagate."""
        try:
            epc: EpcPropertyData = self._epc_client.get_by_certificate_number(
                result.certificate_number
            )
        except self._UNMAPPABLE_CERT_ERRORS as exc:
            self.skipped.append(
                SkippedCohortCert(
                    certificate_number=result.certificate_number,
                    error=f"{type(exc).__name__}: {exc}",
                )
            )
            logger.warning(
                "skipping unmappable cohort cert %s: %s",
                result.certificate_number,
                exc,
            )
            return None
        resolved: Optional[Coordinates] = (
            coordinates.get(result.uprn) if result.uprn is not None else None
        )
        return ComparableProperty(
            epc=epc,
            certificate_number=result.certificate_number,
            address=result.address_line_1,
            registration_date=_parse_date(result.registration_date),
            coordinates=resolved,
        )


def _parse_date(value: str) -> Optional[date]:
    """The register's ISO registration date, or None when unparseable (the
    predictor falls back to an unweighted recency)."""
    try:
        return date.fromisoformat(value[:10])
    except ValueError:
        return None