Model/scripts/expired_prediction_pairs_harness.py

"""Pre-2012 x SAP-10.2 pairs harness for Expired-Enhanced Prediction (ADR-0054).

For each postcode, find properties holding BOTH a pre-2012 cert in the historic
S3 backup AND a SAP-10.2 cert on the new gov API (RdSAP 10 went live June 2025;
only a same-spec lodged figure is a valid validation target — see Component
Accuracy, ADR-0030). Each pair is predicted twice from its leave-one-out
postcode cohort:

- PLAIN arm: property type + built form only (what a blind prediction sees);
- CONDITIONED arm: the historic cert's stable attributes conditioning cohort
  selection (ADR-0054).

Both arms are scored against the lodged SAP-10.2 cert with the SAME metric
suite as the prediction corpus (ADR-0030 Component Accuracy): per-component
classification hit-rates, mean-absolute numeric residuals, plus the secondary
calculator-floored SAP residual (calc(predicted) vs the lodged score). The
per-attribute breakdown is the whitelist evidence: an attribute whose
conditioned hit-rate is WORSE than plain is stale and gets demoted.

The expensive cohort fetch (search-by-postcode + per-cert fetch) happens only
for postcodes where a pair actually exists, so the script can sweep hundreds
of postcodes cheaply.

Usage:
    python scripts/expired_prediction_pairs_harness.py "B93 8SY" "LS6 1AA"
    python scripts/expired_prediction_pairs_harness.py --postcodes-file pcs.txt --out report.md

Env: OPEN_EPC_API_TOKEN, DATA_BUCKET; ambient AWS credentials for S3.
"""

from __future__ import annotations

import argparse
import os
import sys
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from datatypes.epc.domain.epc_property_data import EpcPropertyData  # noqa: E402
from datatypes.epc.domain.historic_epc import HistoricEpc  # noqa: E402
from domain.epc_prediction.comparable_properties import (  # noqa: E402
    ComparableProperty,
    select_comparables,
)
from domain.epc_prediction.epc_prediction import EpcPrediction  # noqa: E402
from domain.epc_prediction.historic_conditioning import (  # noqa: E402
    attributes_with_historic_fallback,
    conditioning_from_historic,
    target_with_conditioning,
)
from domain.epc_prediction.prediction_comparison import (  # noqa: E402
    PredictionComparison,
    compare_prediction,
)
from domain.epc_prediction.prediction_target import (  # noqa: E402
    PredictionTarget,
    build_prediction_target,
)
from domain.postcode import Postcode  # noqa: E402
from domain.property.property import PropertyIdentity  # noqa: E402

_PRE_2012 = "2012-01-01"
_VALIDATION_SAP_VERSION = 10.2

_RESIDUAL_COMPONENTS = (
    "floor_area_m2",
    "building_parts",
    "window_count",
    "total_window_area_m2",
    "door_count",
)


def latest_pre_2012_by_uprn(records: list[HistoricEpc]) -> dict[str, HistoricEpc]:
    """One historic cert per UPRN: the latest lodgement strictly before 2012
    (ISO date strings compare lexicographically). UPRN-less rows are dropped —
    without a UPRN there is nothing to pair."""
    by_uprn: dict[str, HistoricEpc] = {}
    for record in records:
        if not record.uprn or not record.lodgement_date:
            continue
        if record.lodgement_date >= _PRE_2012:
            continue
        current = by_uprn.get(record.uprn)
        if current is None or record.lodgement_date > current.lodgement_date:
            by_uprn[record.uprn] = record
    return by_uprn


@dataclass(frozen=True)
class PairScore:
    """One arm's score for one pair: the component comparison plus the
    calculator-floored SAP residual (calc(predicted) − lodged), None when the
    calculator could not score the predicted picture."""

    comparison: PredictionComparison
    sap_residual: Optional[float]


@dataclass(frozen=True)
class ArmScores:
    """One arm aggregated in the ComponentAccuracy shape (ADR-0030):
    classification maps component -> (hits, applicable-total); residuals maps a
    numeric component -> signed values; sap_residuals are calc − lodged."""

    classification: dict[str, tuple[int, int]]
    residuals: dict[str, list[float]]
    sap_residuals: list[float]


def aggregate(scores: list[PairScore]) -> ArmScores:
    counts: dict[str, list[int]] = defaultdict(lambda: [0, 0])
    residuals: dict[str, list[float]] = defaultdict(list)
    sap_residuals: list[float] = []
    for score in scores:
        comparison = score.comparison
        for component, hit in comparison.categorical_hits.items():
            if hit is None:
                continue
            counts[component][1] += 1
            if hit:
                counts[component][0] += 1
        residuals["floor_area_m2"].append(comparison.floor_area_residual)
        residuals["building_parts"].append(float(comparison.building_parts_residual))
        residuals["window_count"].append(float(comparison.window_count_residual))
        residuals["total_window_area_m2"].append(
            comparison.total_window_area_residual
        )
        residuals["door_count"].append(float(comparison.door_count_residual))
        if score.sap_residual is not None:
            sap_residuals.append(score.sap_residual)
    return ArmScores(
        classification={k: (v[0], v[1]) for k, v in counts.items()},
        residuals=dict(residuals),
        sap_residuals=sap_residuals,
    )


def _mean_abs(values: list[float]) -> str:
    return f"{sum(abs(v) for v in values) / len(values):.1f}" if values else "n/a"


def _rate_cell(hits: tuple[int, int]) -> str:
    hit, total = hits
    return f"{hit}/{total} ({hit / total:.0%})" if total else "n/a"


def format_report(plain: ArmScores, conditioned: ArmScores, pairs: int) -> str:
    """The two arms side by side, in the corpus gate's shape: classification
    hit-rates per component, then mean-abs residuals, then the secondary SAP
    residual."""
    components = sorted(set(plain.classification) | set(conditioned.classification))
    lines = [
        f"# Expired-Enhanced Prediction pairs report ({pairs} pairs)",
        "",
        "## Classification hit-rates (hits/applicable)",
        "",
        "| component | plain | conditioned |",
        "|---|---|---|",
    ]
    for component in components:
        p = _rate_cell(plain.classification.get(component, (0, 0)))
        c = _rate_cell(conditioned.classification.get(component, (0, 0)))
        lines.append(f"| {component} | {p} | {c} |")
    lines += [
        "",
        "## Mean absolute residuals (predicted − actual)",
        "",
        "| component | plain | conditioned |",
        "|---|---|---|",
    ]
    for component in _RESIDUAL_COMPONENTS:
        p = _mean_abs(plain.residuals.get(component, []))
        c = _mean_abs(conditioned.residuals.get(component, []))
        lines.append(f"| {component} | {p} | {c} |")
    lines += [
        "",
        "## SAP residual — secondary, calculator-floored (calc(predicted) − lodged)",
        "",
        "| metric | plain | conditioned |",
        "|---|---|---|",
        f"| mean abs | {_mean_abs(plain.sap_residuals)} "
        f"| {_mean_abs(conditioned.sap_residuals)} |",
        f"| scored | {len(plain.sap_residuals)} | {len(conditioned.sap_residuals)} |",
    ]
    return "\n".join(lines)


def _predict_arm(
    target: Optional[PredictionTarget],
    cohort: list[ComparableProperty],
    predictor: EpcPrediction,
) -> Optional[EpcPropertyData]:
    if target is None:
        return None
    comparables = select_comparables(target, cohort)
    if not comparables.members:
        return None
    return predictor.predict(target, comparables)


def _part0(epc: EpcPropertyData) -> Optional[object]:
    parts = epc.sap_building_parts
    return parts[0] if parts else None


def _actual_age_band(epc: EpcPropertyData) -> Optional[object]:
    part = _part0(epc)
    return getattr(part, "construction_age_band", None)


def _actual_wall(epc: EpcPropertyData) -> Optional[object]:
    part = _part0(epc)
    return getattr(part, "wall_construction", None)


def _actual_fuel(epc: EpcPropertyData) -> Optional[object]:
    details = epc.sap_heating.main_heating_details
    return details[0].main_fuel_type if details else None


def _tfa_within_band(actual_tfa: float, hist_tfa: Optional[float]) -> Optional[bool]:
    if hist_tfa is None:
        return None
    return abs(actual_tfa - hist_tfa) <= 0.05 * hist_tfa


@dataclass(frozen=True)
class LadderStep:
    """One conditioning filter's fate in the sequential relax ladder: how many
    of the incoming cohort matched, and whether it engaged (matches >= k)."""

    matches: int
    cohort_before: int
    engaged: bool


def simulate_conditioning_ladder(
    base: list[ComparableProperty],
    *,
    age_band: Optional[str],
    main_fuel: Optional[int],
    total_floor_area_m2: Optional[float],
    minimum_cohort: int = 5,
) -> dict[str, Optional[LadderStep]]:
    """Replay select_comparables' age->fuel->TFA conditioning sequence over the
    plain arm's selected cohort, recording per filter whether it ENGAGED
    (>= minimum_cohort matches survive) or RELAXED. None = attribute unresolved,
    filter never active."""
    steps: dict[str, Optional[LadderStep]] = {}
    cohort = list(base)

    def apply(name: str, active: bool, matches: list[ComparableProperty]) -> None:
        nonlocal cohort
        if not active:
            steps[name] = None
            return
        engaged = len(matches) >= minimum_cohort
        steps[name] = LadderStep(len(matches), len(cohort), engaged)
        if engaged:
            cohort = matches

    apply(
        "construction_age_band",
        age_band is not None,
        [c for c in cohort if _actual_age_band(c.epc) == age_band],
    )
    apply(
        "main_fuel",
        main_fuel is not None,
        [c for c in cohort if _actual_fuel(c.epc) == main_fuel],
    )
    apply(
        "total_floor_area",
        total_floor_area_m2 is not None,
        [
            c
            for c in cohort
            if total_floor_area_m2 is not None
            and abs(c.epc.total_floor_area_m2 - total_floor_area_m2)
            <= 0.05 * total_floor_area_m2
        ],
    )
    return steps


def format_diagnosis(rows: list[dict[str, object]]) -> str:
    """Aggregate the per-pair telemetry into the two diagnosis tables: did each
    conditioning filter ever ENGAGE, and does the historic value still AGREE
    with the newly lodged one (the staleness measurement)."""
    if not rows:
        return ""
    filters = ("construction_age_band", "main_fuel", "total_floor_area")
    lines = [
        "",
        "## Diagnosis — filter engagement (conditioned arm)",
        "",
        "| filter | resolved | engaged | relaxed (too few matches) |",
        "|---|---|---|---|",
    ]
    for name in filters:
        resolved = engaged = 0
        for row in rows:
            step = row.get(f"ladder_{name}")
            if step is None:
                continue
            resolved += 1
            if isinstance(step, LadderStep) and step.engaged:
                engaged += 1
        lines.append(
            f"| {name} | {resolved}/{len(rows)} | {engaged}/{resolved or 1} "
            f"| {resolved - engaged}/{resolved or 1} |"
        )
    attrs = (
        "property_type",
        "built_form",
        "wall_construction",
        "construction_age_band",
        "main_fuel",
        "tfa_within_5pct",
    )
    lines += [
        "",
        "## Diagnosis — historic vs newly-lodged agreement (staleness)",
        "",
        "| attribute | historic resolved | agrees with new cert |",
        "|---|---|---|",
    ]
    for name in attrs:
        resolved = agrees = 0
        for row in rows:
            value = row.get(f"agrees_{name}")
            if value is None:
                continue
            resolved += 1
            if value:
                agrees += 1
        pct = f" ({agrees / resolved:.0%})" if resolved else ""
        lines.append(f"| {name} | {resolved}/{len(rows)} | {agrees}/{resolved or 1}{pct} |")
    sizes: list[int] = [
        size
        for row in rows
        if isinstance((size := row.get("plain_cohort_size")), int)
    ]
    if sizes:
        lines += [
            "",
            f"Mean plain-arm cohort size: {sum(sizes) / len(sizes):.1f} "
            f"(min {min(sizes)}, max {max(sizes)}); relax threshold k=5.",
        ]
    return "\n".join(lines)


def run(  # pragma: no cover - live IO composition
    postcodes: list[str], telemetry_path: Optional[Path] = None
) -> str:
    from domain.sap10_calculator.calculator import Sap10Calculator
    from infrastructure.epc_client.epc_client_service import EpcClientService
    from repositories.comparable_properties.epc_comparable_properties_repository import (
        EpcComparablePropertiesRepository,
    )
    from repositories.geospatial.geospatial_s3_repository import (
        GeospatialS3Repository,
    )
    from repositories.historic_epc.historic_epc_s3_repository import (
        HistoricEpcS3Repository,
    )
    from scripts.e2e_common import load_env, s3_parquet_reader

    load_env()
    epc_client = EpcClientService(os.environ["OPEN_EPC_API_TOKEN"])
    geospatial = GeospatialS3Repository(s3_parquet_reader(os.environ["DATA_BUCKET"]))
    comparables_repo = EpcComparablePropertiesRepository(epc_client, geospatial)
    historic_repo = HistoricEpcS3Repository.with_default_s3_client()
    predictor = EpcPrediction()
    calculator = Sap10Calculator()

    def sap_residual(
        predicted: Optional[EpcPropertyData], actual: EpcPropertyData
    ) -> Optional[float]:
        if predicted is None or actual.energy_rating_current is None:
            return None
        try:
            calculated: float = calculator.calculate(predicted).sap_score_continuous
        except Exception as error:  # the calculator strict-raises on gaps
            print(f"  calculator raised: {error}", file=sys.stderr)
            return None
        return calculated - float(actual.energy_rating_current)

    plain_scores: list[PairScore] = []
    conditioned_scores: list[PairScore] = []
    telemetry: list[dict[str, object]] = []
    pairs = 0
    for index, raw_postcode in enumerate(postcodes):
        postcode = str(Postcode(raw_postcode))
        historic = latest_pre_2012_by_uprn(
            historic_repo.get_for_postcode(Postcode(raw_postcode))
        )
        # Pair-check first: only a postcode with a SAP-10.2 relodgement of a
        # pre-2012 UPRN pays for the cohort fetch.
        paired: list[tuple[str, HistoricEpc, EpcPropertyData]] = []
        for uprn, record in historic.items():
            actual = epc_client.get_by_uprn(int(uprn))
            if actual is not None and actual.sap_version == _VALIDATION_SAP_VERSION:
                paired.append((uprn, record, actual))
        print(
            f"[{index + 1}/{len(postcodes)}] {postcode}: "
            f"{len(historic)} pre-2012 UPRNs, {len(paired)} pairs",
            file=sys.stderr,
            flush=True,
        )
        if not paired:
            continue
        cohort = comparables_repo.candidates_for(postcode)
        for uprn, record, actual in paired:
            pairs += 1
            loo_cohort = [c for c in cohort if c.epc.uprn != int(uprn)]
            identity = PropertyIdentity(
                portfolio_id=0, postcode=postcode, address=record.address, uprn=int(uprn)
            )
            conditioning = conditioning_from_historic(record)
            attributes = attributes_with_historic_fallback(None, conditioning)
            plain_target = build_prediction_target(identity, None, attributes)
            conditioned_target = (
                target_with_conditioning(plain_target, conditioning)
                if plain_target is not None
                else None
            )
            plain = _predict_arm(plain_target, loo_cohort, predictor)
            conditioned = _predict_arm(conditioned_target, loo_cohort, predictor)
            if plain_target is not None:
                base = list(select_comparables(plain_target, loo_cohort).members)
                ladder = simulate_conditioning_ladder(
                    base,
                    age_band=conditioning.construction_age_band,
                    main_fuel=conditioning.main_fuel,
                    total_floor_area_m2=conditioning.total_floor_area_m2,
                )
                telemetry.append(
                    {
                        "postcode": postcode,
                        "uprn": uprn,
                        "plain_cohort_size": len(base),
                        **{f"ladder_{k}": v for k, v in ladder.items()},
                        "agrees_property_type": (
                            None
                            if conditioning.property_type is None
                            else conditioning.property_type == actual.property_type
                        ),
                        "agrees_built_form": (
                            None
                            if conditioning.built_form is None
                            else conditioning.built_form == actual.built_form
                        ),
                        "agrees_wall_construction": (
                            None
                            if conditioning.wall_construction is None
                            else conditioning.wall_construction == _actual_wall(actual)
                        ),
                        "agrees_construction_age_band": (
                            None
                            if conditioning.construction_age_band is None
                            else conditioning.construction_age_band
                            == _actual_age_band(actual)
                        ),
                        "agrees_main_fuel": (
                            None
                            if conditioning.main_fuel is None
                            else conditioning.main_fuel == _actual_fuel(actual)
                        ),
                        "agrees_tfa_within_5pct": _tfa_within_band(
                            actual.total_floor_area_m2,
                            conditioning.total_floor_area_m2,
                        ),
                    }
                )
            if plain is not None:
                plain_scores.append(
                    PairScore(compare_prediction(plain, actual), sap_residual(plain, actual))
                )
            if conditioned is not None:
                conditioned_scores.append(
                    PairScore(
                        compare_prediction(conditioned, actual),
                        sap_residual(conditioned, actual),
                    )
                )

    if telemetry_path is not None:
        import dataclasses as _dc
        import json as _json

        with telemetry_path.open("w") as handle:
            for row in telemetry:
                serialisable = {
                    k: (_dc.asdict(v) if isinstance(v, LadderStep) else v)
                    for k, v in row.items()
                }
                handle.write(_json.dumps(serialisable) + "\n")
    report = format_report(
        aggregate(plain_scores), aggregate(conditioned_scores), pairs
    )
    return report + format_diagnosis(telemetry)


def main() -> None:  # pragma: no cover - CLI entry
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("postcodes", nargs="*", help="postcodes to scan")
    parser.add_argument("--postcodes-file", type=Path, default=None)
    parser.add_argument("--out", type=Path, default=None, help="write the report here")
    parser.add_argument(
        "--telemetry", type=Path, default=None, help="write per-pair JSONL here"
    )
    args = parser.parse_args()
    postcodes: list[str] = list(args.postcodes)
    if args.postcodes_file is not None:
        postcodes += [
            line.strip()
            for line in args.postcodes_file.read_text().splitlines()
            if line.strip()
        ]
    if not postcodes:
        parser.error("no postcodes given")
    report = run(postcodes, telemetry_path=args.telemetry)
    if args.out is not None:
        args.out.write_text(report + "\n")
    print(report)


if __name__ == "__main__":
    main()