Model/scripts/run_modelling_e2e.py

"""Run Modelling end-to-end for specific Properties (by ``property_id``) and
print the recommendations for inspection.

The local DB's Properties have no linked, ingested EPC yet (Ingestion's source
clients are still stubbed — #1136), so this script does the ingestion step
inline: it reads each Property's UPRN from the DB, fetches the latest EPC
**live** from the gov EPC API by UPRN, resolves the UPRN's spatial reference
from S3, and fetches Google Solar — then runs the Modelling stage (every
Recommendation Generator → the Optimiser → a costed, attributed Plan). The same
local computation runs whether or not you store the result: by default it
persists **nothing** (the run is for inspecting recommendations); pass
`--persist` to write the inputs + the Plan to the DB. With `--persist`, a
lodged-EPC Property **also** gets its Baseline Performance row written (lodged
vs calculator-effective SAP + the bill block) via the same orchestrator the
first-run pipeline uses — run per Property so one bad cert does not abort the
batch. A predicted (EPC-less) Property has no lodged figures, so it gets a Plan
but no baseline row.

To keep the inspected recommendations identical to what gets stored, **both
modes price against the live ``material`` catalogue (read-only)** and model
against a real **Scenario** read from the DB — not the JSON sample catalogue.
Pass `--scenario-id` to target a real Scenario; its ``goal_value`` drives the
band and **its ``exclusions`` drive which measures the run considers** (the live
scenario table persists exclusions only, no inclusions). Without `--scenario-id`
the run synthesises an Increasing-EPC-to-``--goal`` Scenario with no exclusions.
`--measures` / `--exclude-measures` are optional overlays layered on top of the
Scenario's own exclusions.

``secondary_heating_removal`` is priced from the committed off-catalogue JSON
overlay (the live ``material.type`` enum cannot carry it), so no exclusion flag
is needed. ASHP is priced off the rate sheet (``material_id=None``), also fine.

Config: loads `backend/.env` for the DB creds (`DB_*`), the EPC API token
(`OPEN_EPC_API_TOKEN` — the Bearer token for the new gov API), the Google Solar
key (`GOOGLE_SOLAR_API_KEY`) and the S3
reference bucket (`DATA_BUCKET`) — the agent never sees the secrets. AWS creds
come from the ambient `~/.aws` profile. Run from the worktree root:

    # inspect only (no DB writes), Scenario 1266, measures from the Scenario:
    python -m scripts.run_modelling_e2e --scenario-id 1266 \
        --exclude-measures secondary_heating_removal 709634 709635 709636
    # same run, but persist EPC + spatial + solar + Plan (needs --portfolio-id):
    python -m scripts.run_modelling_e2e --scenario-id 1266 --portfolio-id 785 \
        --persist --exclude-measures secondary_heating_removal 709634 709635
    python -m scripts.run_modelling_e2e --no-solar 709634 709635   # skip Google leg

Per Property the spatial reference (S3 Open-UPRN parquet) gives the planning
protections (conservation/listed/heritage — gate the wall + solar measures) and
the coordinates that drive the Google Solar fetch (ADR-0026). Buildings S3
doesn't cover, or that Google has no solar coverage for, fall back to
unrestricted / no-solar and are still modelled. Pass `--no-solar` to skip the
Google leg.
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path
from typing import Any, Callable, Optional

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))  # worktree root first — avoid the import trap

from datatypes.epc.domain.epc_property_data import (  # noqa: E402
    BuildingPartIdentifier,
    EpcPropertyData,
)
from domain.property.property import Property, PropertyIdentity  # noqa: E402
from repositories.property.landlord_override_overlays import (  # noqa: E402
    overlays_from,
)
from repositories.property.property_overrides_postgres_reader import (  # noqa: E402
    PropertyOverridesPostgresReader,
)
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.prediction_target import (  # noqa: E402
    PredictionTarget,
    build_prediction_target,
)
from domain.geospatial.coordinates import Coordinates  # noqa: E402
from domain.geospatial.planning_restrictions import PlanningRestrictions  # noqa: E402
from domain.geospatial.spatial_reference import SpatialReference  # noqa: E402
from domain.modelling.considered_measures import (  # noqa: E402
    combine_considered_measures,
)
from domain.modelling.measure_type import MeasureType  # noqa: E402
from domain.modelling.plan import Plan, PlanMeasure  # noqa: E402
from domain.modelling.recommendation import Recommendation  # noqa: E402
from domain.modelling.scenario import Scenario  # noqa: E402
from harness.console import candidate_recommendations, run_modelling  # noqa: E402
from harness.plan_table import format_plan_table  # noqa: E402
from infrastructure.epc_client.epc_client_service import EpcClientService  # noqa: E402
from infrastructure.postcodes_io.postcodes_io_client import (  # noqa: E402
    PostcodesIoClient,
)
from infrastructure.solar.google_solar_api_client import (  # noqa: E402
    BuildingInsightsNotFoundError,
    GoogleSolarApiClient,
)
from repositories.comparable_properties.epc_comparable_properties_repository import (  # noqa: E402
    EpcComparablePropertiesRepository,
)
from repositories.geospatial.geospatial_s3_repository import (  # noqa: E402
    GeospatialS3Repository,
)
from repositories.product.composite_product_repository import (  # noqa: E402
    catalogue_with_off_catalogue_overrides,
)
from repositories.product.product_repository import ProductRepository  # noqa: E402
from repositories.property.override_backed_prediction_attributes_reader import (  # noqa: E402
    OverrideBackedPredictionAttributesReader,
)
from repositories.postgres_unit_of_work import PostgresUnitOfWork  # noqa: E402
from orchestration.property_baseline_orchestrator import (  # noqa: E402
    PropertyBaselineOrchestrator,
)
from domain.property_baseline.calculator_rebaseliner import (  # noqa: E402
    CalculatorRebaseliner,
)
from domain.sap10_calculator.calculator import Sap10Calculator  # noqa: E402
from repositories.fuel_rates.fuel_rates_static_file_repository import (  # noqa: E402
    FuelRatesStaticFileRepository,
)
from repositories.scenario.scenario_postgres_repository import (  # noqa: E402
    ScenarioPostgresRepository,
)
from scripts.e2e_common import (  # noqa: E402
    ENV_PATH,
    build_engine,
    load_env,
    s3_parquet_reader,
)
from sqlalchemy import Engine, text  # noqa: E402
from sqlmodel import Session  # noqa: E402

_MARKDOWN_PATH = Path("modelling_e2e.md")
_CSV_PATH = Path("modelling_e2e.csv")
_CANDIDATES_CSV_PATH = Path("modelling_e2e_candidates.csv")


def _spatial_for(repo: GeospatialS3Repository, uprn: int) -> Optional[SpatialReference]:
    """The UPRN's spatial reference (coordinates + planning protections), or
    None when S3 doesn't cover it — a missing reference must not abort the run,
    so a lookup error degrades to None (unrestricted, no solar)."""
    try:
        return repo.spatial_for(uprn)
    except Exception as error:  # noqa: BLE001 — S3/parquet hiccup is non-fatal
        print(
            f"  spatial lookup failed for uprn {uprn}: {type(error).__name__}: {error}"
        )
        return None


def _solar_insights_for(
    solar_client: GoogleSolarApiClient, spatial: Optional[SpatialReference]
) -> Optional[dict[str, Any]]:
    """The raw Google Solar `buildingInsights` for the reference's coordinates,
    or None when there are no coordinates / Google has no coverage there."""
    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  # no Google solar coverage at this point — model without it
    # A transient Solar failure (timeout/reset) is NOT swallowed: it propagates so
    # the property is marked ERROR and the wrapper's retry sweep re-runs it later
    # when Solar recovers. We must not silently model a coverage-having property
    # without its solar leg.


def _uprns_for(engine: Engine, property_ids: list[int]) -> dict[int, Optional[int]]:
    """Read each Property's UPRN from the DB (read-only)."""
    with engine.connect() as conn:
        rows = conn.execute(
            text("SELECT id, uprn FROM property WHERE id = ANY(:ids)"),
            {"ids": property_ids},
        ).fetchall()
    return {int(pid): (int(uprn) if uprn is not None else None) for pid, uprn in rows}


def _postcodes_for(engine: Engine, property_ids: list[int]) -> dict[int, str]:
    """Read each Property's postcode from the DB (read-only). Needed to find the
    EPC-Prediction cohort (the postcode's other lodged certs) and to seed the
    PredictionTarget when a Property has no EPC."""
    with engine.connect() as conn:
        rows = conn.execute(
            text("SELECT id, postcode FROM property WHERE id = ANY(:ids)"),
            {"ids": property_ids},
        ).fetchall()
    return {int(pid): (postcode or "") for pid, postcode in rows}


def _dump_overrides(engine: Engine, property_ids: list[int]) -> None:
    """Print each target Property's ``property_overrides`` rows (read-only), so the
    Landlord Overrides folded into the Effective EPC are visible before modelling."""
    with engine.connect() as conn:
        rows = conn.execute(
            text(
                "SELECT property_id, building_part, override_component, override_value "
                "FROM property_overrides WHERE property_id = ANY(:ids) "
                "ORDER BY property_id, building_part, override_component"
            ),
            {"ids": property_ids},
        ).fetchall()
    if not rows:
        print("landlord overrides: none for the target propertie(s)\n")
        return
    print("landlord overrides (folded into the Effective EPC):")
    for property_id, building_part, component, value in rows:
        print(f"  property {property_id} · part {building_part} · {component} = {value}")
    print()


def _main_wall_summary(epc: EpcPropertyData) -> str:
    """The MAIN building part's wall codes — what the calculator scores for the
    wall U-value. Used to show whether a Landlord Override moved them."""
    for part in epc.sap_building_parts:
        if part.identifier is BuildingPartIdentifier.MAIN:
            return (
                f"wall_construction={part.wall_construction} "
                f"wall_insulation_type={part.wall_insulation_type}"
            )
    return "no MAIN building part"


def _scenario_for(session: Session, scenario_id: int) -> Scenario:
    """Read the Scenario the run targets (read-only). An Increasing-EPC Scenario
    must carry a ``goal_value`` (band) — the old null-band rows were a fixed bug
    and crash the Optimiser's target — so reject one that does not."""
    scenario: Scenario = ScenarioPostgresRepository(session).get_many([scenario_id])[0]
    if scenario.goal == "Increasing EPC" and not scenario.goal_value:
        raise ValueError(
            f"scenario {scenario_id} has no goal_value (band); pick a recent one"
        )
    return scenario


def _parse_measures(raw: Optional[str]) -> Optional[frozenset[MeasureType]]:
    """Parse `--measures a,b,c` into a `considered_measures` allowlist, or None
    (consider every modelled measure) when unset. Raises on an unknown type."""
    if raw is None:
        return None
    return frozenset(
        MeasureType(token.strip()) for token in raw.split(",") if token.strip()
    )


def _resolve_considered(
    allowlist: Optional[frozenset[MeasureType]],
    excluded: Optional[frozenset[MeasureType]],
) -> Optional[frozenset[MeasureType]]:
    """Combine the `--measures` allowlist with the `--exclude-measures` set. With
    no exclusions the allowlist is returned unchanged (None = every measure).
    With exclusions the result is (the allowlist, or every measure) minus the
    excluded types — so `--exclude-measures secondary_heating_removal` considers
    every measure except that one, without enumerating the rest."""
    if not excluded:
        return allowlist
    base = allowlist if allowlist is not None else frozenset(MeasureType)
    return base - excluded


def _context_summary(
    spatial: Optional[SpatialReference], solar_insights: Optional[dict[str, Any]]
) -> str:
    """A one-line note on what the geospatial leg contributed: which planning
    protections gated the measures, and whether Google Solar potential fired."""
    if spatial is None:
        restrictions_note = "no spatial reference"
    else:
        flags = [
            name
            for name, on in (
                ("conservation", spatial.restrictions.in_conservation_area),
                ("listed", spatial.restrictions.is_listed),
                ("heritage", spatial.restrictions.is_heritage),
            )
            if on
        ]
        restrictions_note = ", ".join(flags) if flags else "unrestricted"
    solar_note = "solar ✓" if solar_insights is not None else "no solar"
    return f"{restrictions_note}; {solar_note}"


def _measure_summary(measure: PlanMeasure) -> str:
    return (
        f"  - {measure.measure_type}: "
        f"+{measure.impact.sap_points:.2f} SAP · £{measure.cost.total:,.0f} "
        f"— {measure.description}"
    )


def _candidate_lines(
    recommendations: list[Recommendation], selected: set[MeasureType]
) -> list[str]:
    """Render every candidate Option (the full menu the Generators produced,
    not just the Plan the Optimiser selected) with its per-Option cost, flagging
    the Options that made it into the Plan — so measures the Optimiser passed
    over (e.g. an ASHP it found too costly for the target band) are visible."""
    lines: list[str] = []
    for recommendation in recommendations:
        for option in recommendation.options:
            cost = option.cost
            cost_note = (
                f"£{cost.total:,.0f} (+{cost.contingency_rate * 100:.0f}% cont.)"
                if cost is not None
                else "no cost"
            )
            flag = " ✓ SELECTED" if option.measure_type in selected else ""
            lines.append(
                f"    [{recommendation.surface}] {option.measure_type} · "
                f"{cost_note}{flag} — {option.description}"
            )
    return lines


def _candidate_csv_rows(
    property_id: int,
    uprn: Optional[int],
    recommendations: list[Recommendation],
    selected: set[MeasureType],
) -> list[str]:
    """One CSV row per candidate Option: the full measure menu with cost,
    contingency, and whether the Optimiser selected it."""
    rows: list[str] = []
    for recommendation in recommendations:
        for option in recommendation.options:
            cost = option.cost
            total = f"{cost.total:.2f}" if cost is not None else ""
            contingency = f"{cost.contingency_rate:.4f}" if cost is not None else ""
            chosen = "yes" if option.measure_type in selected else "no"
            description = option.description.replace(",", ";")
            rows.append(
                f"{property_id},{uprn or ''},{recommendation.surface},"
                f"{option.measure_type},{total},{contingency},{chosen},{description}"
            )
    return rows


def _persist(
    engine: Engine,
    *,
    property_id: int,
    uprn: int,
    portfolio_id: int,
    scenario: Scenario,
    epc: Optional[EpcPropertyData],
    spatial: Optional[SpatialReference],
    solar_insights: Optional[dict[str, Any]],
    plan: Plan,
) -> None:
    """Write the run's inputs (EPC + spatial + solar) and the computed Plan to
    the DB in one Unit of Work, then commit. ``PlanPostgresRepository`` replaces
    any existing Plan for ``(property_id, scenario.id)`` (idempotent re-run). A
    predicted Property has no lodged EPC to store (``epc is None``), so only the
    spatial/solar inputs and the Plan are persisted for it."""
    with PostgresUnitOfWork(lambda: Session(engine)) as uow:
        if epc is not None:
            uow.epc.save(epc, property_id=property_id, portfolio_id=portfolio_id)
        if spatial is not None:
            uow.spatial.save(uprn, spatial)
        # The live `solar` table is keyed by UPRN and needs the fetch's
        # coordinates; insights are only present when those coordinates were
        # (see `_solar_insights_for`), so `spatial.coordinates` is non-None here.
        if solar_insights is not None:
            assert 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,
        )
        # Mark the Property as run under the new process (old engine's
        # `has_recommendations` marker + a bumped `updated_at`); the modelling
        # compute above runs on in-memory fakes, so this DB UoW must set it.
        uow.property.mark_modelled(
            property_id, has_recommendations=bool(plan.measures)
        )
        uow.commit()


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 the Property is
    ineligible (``property_type`` unresolvable) or no comparable neighbours exist.

    The cohort is found by POSTCODE, so a wrong postcode on the property row
    yields the wrong neighbours — a prediction is only as good as the postcode it
    is given. When the own postcode holds no same-type comparables, the cohort is
    broadened to the real unit postcodes physically nearest it (``broaden``)."""
    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  # property_type unresolvable — gated out of prediction
    comparables = select_comparables(target, cohort_for(target.postcode))
    if not comparables.members:
        # Sparse own postcode — reach out to the nearest real postcodes.
        comparables = select_comparables(target, broaden(target))
    if not comparables.members:
        return None  # no comparable neighbours nearby either
    predicted = predictor.predict(target, comparables)
    # The calculator needs a MAIN building part; a cohort whose template carries
    # none (e.g. a malformed flat record) yields an unscoreable picture, so reject
    # it as not-predictable rather than letting the calculator StopIteration.
    if not any(
        part.identifier is BuildingPartIdentifier.MAIN
        for part in predicted.sap_building_parts
    ):
        return None
    return predicted


def _run_from_db(
    args: argparse.Namespace,
    *,
    engine: Engine,
    products: ProductRepository,
    scenario: Optional[Scenario],
    considered: Optional[frozenset[MeasureType]],
    baseline_orchestrator: Optional[PropertyBaselineOrchestrator],
    md_path: Path,
    csv_path: Path,
    candidates_path: Path,
    target: str,
    measures_note: str,
) -> None:
    """Re-model from already-persisted inputs — **zero gov-API calls**.

    Reads each Property's Effective EPC (lodged-or-predicted, overrides folded),
    planning protections and solar straight from the DB (a prior ``--persist``
    ingestion must have stored them), runs the same modelling, and — with
    ``--persist`` — re-writes the Plan and, for lodged-EPC Properties, the
    Baseline. A predicted Property has no lodged figures, so it gets no baseline
    row (same rule as the live path). One bad property is logged and skipped.
    """
    md_lines: list[str] = [f"# Modelling recommendations ({target}, {measures_note})\n"]
    csv_rows: list[str] = [
        "property_id,uprn,api_sap,baseline_sap,sap_delta,post_sap,measures,"
        "measure_types,cost_of_works"
    ]
    candidate_csv_rows: list[str] = [
        "property_id,uprn,surface,measure_type,cost_total,contingency_rate,"
        "selected,description"
    ]
    total = len(args.property_ids)
    run_start = time.monotonic()
    errors = 0
    for index, property_id in enumerate(args.property_ids, start=1):
        elapsed = time.monotonic() - run_start
        eta = (elapsed / (index - 1)) * (total - index + 1) if index > 1 else 0.0
        print(
            f"[{index}/{total}] · {errors} err · elapsed {elapsed / 60:.1f}m "
            f"· ETA {eta / 60:.1f}m · property {property_id} (from DB)",
            flush=True,
        )
        try:
            with PostgresUnitOfWork(lambda: Session(engine)) as uow:
                prop = uow.property.get(property_id)
                effective_epc: EpcPropertyData = prop.effective_epc
                restrictions: PlanningRestrictions = prop.planning_restrictions
                uprn: Optional[int] = prop.identity.uprn
                epc: Optional[EpcPropertyData] = prop.epc
                solar_insights: Optional[dict[str, Any]] = (
                    uow.solar.get(uprn) if uprn is not None else None
                )
            predicted = epc is None
            plan: Plan = run_modelling(
                effective_epc,
                goal_band=args.goal,
                planning_restrictions=restrictions,
                solar_insights=solar_insights,
                considered_measures=considered,
                products=products,
                scenario=scenario,
                print_table=False,
            )
            candidates: list[Recommendation] = candidate_recommendations(
                epc if epc is not None else effective_epc,
                planning_restrictions=restrictions,
                solar_insights=solar_insights,
                considered_measures=considered,
                products=products,
            )
            if args.persist:
                assert scenario is not None  # guaranteed by the --persist guard
                with PostgresUnitOfWork(lambda: Session(engine)) as uow:
                    uow.plan.save(
                        plan,
                        property_id=property_id,
                        scenario_id=scenario.id,
                        portfolio_id=args.portfolio_id,
                        is_default=scenario.is_default,
                    )
                    uow.property.mark_modelled(
                        property_id, has_recommendations=bool(plan.measures)
                    )
                    uow.commit()
                # Lodged EPC also gets its Baseline Performance re-established from
                # the persisted EPC; predicted Properties have no lodged figures.
                if epc is not None:
                    assert baseline_orchestrator is not None
                    baseline_orchestrator.run([property_id])
        except Exception as error:  # noqa: BLE001 — one bad property must not stop the run
            errors += 1
            line = f"property {property_id}: ERROR — {type(error).__name__}: {error}"
            print(line + "\n")
            md_lines.append(f"## Property {property_id}\n\n`{line}`\n")
            csv_rows.append(f"{property_id},,,,,,,ERROR,")
            continue

        measure_types = [m.measure_type for m in plan.measures]
        selected: set[MeasureType] = {m.measure_type for m in plan.measures}
        flags = [
            name
            for name, on in (
                ("conservation", restrictions.in_conservation_area),
                ("listed", restrictions.is_listed),
                ("heritage", restrictions.is_heritage),
            )
            if on
        ]
        context = (
            f"{', '.join(flags) if flags else 'unrestricted'}; "
            f"{'solar ✓' if solar_insights is not None else 'no solar'}"
        )
        source_tag = " · ⚠ PREDICTED (no lodged EPC)" if predicted else ""
        candidate_lines = _candidate_lines(candidates, selected)
        print(
            f"=== Property {property_id} (uprn {uprn}) === "
            f"SAP {plan.baseline.sap_continuous:.1f} -> {plan.post_sap_continuous:.1f} "
            f"· {len(plan.measures)} measure(s) · £{plan.cost_of_works:,.0f} "
            f"· {context}{source_tag}"
        )
        print(format_plan_table(plan))
        md_lines.append(f"## Property {property_id} (uprn {uprn}){source_tag}\n")
        md_lines.append(
            f"SAP {plan.baseline.sap_continuous:.1f} → {plan.post_sap_continuous:.1f} "
            f"· {len(plan.measures)} measure(s) · cost £{plan.cost_of_works:,.0f} "
            f"· {context}\n"
        )
        md_lines.append("**Selected Plan**\n")
        md_lines.extend(_measure_summary(m) for m in plan.measures)
        md_lines.append("")
        md_lines.append("**All candidate measures (cost per measure)**\n")
        md_lines.extend(candidate_lines)
        md_lines.append("")
        api_sap: Optional[int] = epc.energy_rating_current if epc is not None else None
        calc_sap: float = plan.baseline.sap_continuous
        api_cell = "" if api_sap is None else str(api_sap)
        delta_cell = "" if api_sap is None else f"{calc_sap - api_sap:.2f}"
        csv_rows.append(
            f"{property_id},{uprn},{api_cell},{calc_sap:.2f},{delta_cell},"
            f"{plan.post_sap_continuous:.2f},{len(plan.measures)},"
            f"{'|'.join(measure_types)},{plan.cost_of_works:.0f}"
        )
        candidate_csv_rows.extend(
            _candidate_csv_rows(property_id, uprn, candidates, selected)
        )

    md_path.write_text("\n".join(md_lines) + "\n", encoding="utf-8")
    csv_path.write_text("\n".join(csv_rows) + "\n", encoding="utf-8")
    candidates_path.write_text("\n".join(candidate_csv_rows) + "\n", encoding="utf-8")
    print(f"wrote {md_path.resolve()}")
    print(f"wrote {csv_path.resolve()}")
    print(f"wrote {candidates_path.resolve()}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "property_ids", type=int, nargs="+", help="Property ids to model"
    )
    parser.add_argument(
        "--goal", default="C", help="target band when no --scenario-id (default C)"
    )
    parser.add_argument(
        "--scenario-id", type=int, default=None, help="model against this DB Scenario"
    )
    parser.add_argument(
        "--measures",
        default=None,
        help="optional override: comma-separated measure types to consider. The "
        "Scenario's exclusions already drive this; the flag narrows it further.",
    )
    parser.add_argument(
        "--exclude-measures",
        default=None,
        help="optional override: comma-separated measure types to exclude on top "
        "of the Scenario's own exclusions (e.g. secondary_heating_removal, which "
        "the live catalogue does not yet stock)",
    )
    parser.add_argument(
        "--portfolio-id",
        type=int,
        default=None,
        help="portfolio id (required for --persist)",
    )
    parser.add_argument(
        "--persist",
        action="store_true",
        help="WRITE the inputs + Plan to the DB (default: inspect only, no writes)",
        default=False,
    )
    parser.add_argument(
        "--no-solar",
        action="store_true",
        help="skip the live Google Solar fetch (no Solar PV Options)",
    )
    parser.add_argument(
        "--out-prefix",
        default=None,
        help="write outputs to <prefix>.md / <prefix>.csv / <prefix>_candidates.csv "
        "(parent dirs created) instead of ./modelling_e2e.*; lets batched runs "
        "keep separate, durable output files",
    )
    parser.add_argument(
        "--from-db",
        action="store_true",
        default=False,
        help="re-model from already-persisted inputs: read each Property's "
        "Effective EPC + planning protections + solar from the DB and skip the "
        "live EPC/spatial/solar fetch entirely (zero gov-API calls). Requires a "
        "prior --persist ingestion run; with --persist it re-writes the Plan "
        "(and Baseline for lodged-EPC Properties) without re-fetching.",
    )
    args = parser.parse_args()

    if args.persist and (args.scenario_id is None or args.portfolio_id is None):
        parser.error("--persist requires --scenario-id and --portfolio-id")

    if args.out_prefix:
        _base = Path(args.out_prefix)
        _base.parent.mkdir(parents=True, exist_ok=True)
        md_path = _base.with_suffix(".md")
        csv_path = _base.with_suffix(".csv")
        candidates_path = _base.parent / f"{_base.name}_candidates.csv"
    else:
        md_path, csv_path, candidates_path = (
            _MARKDOWN_PATH,
            _CSV_PATH,
            _CANDIDATES_CSV_PATH,
        )

    load_env(ENV_PATH)
    # The new gov EPC API (Bearer) authenticates with OPEN_EPC_API_TOKEN — the
    # name is misleading; EPC_AUTH_TOKEN is dead (403). Verified against the
    # /api/domestic/search endpoint.
    epc_client = EpcClientService(os.environ["OPEN_EPC_API_TOKEN"])
    geospatial = GeospatialS3Repository(s3_parquet_reader(os.environ["DATA_BUCKET"]))
    solar_client = GoogleSolarApiClient(os.environ["GOOGLE_SOLAR_API_KEY"])
    engine = build_engine()
    cli_considered = _resolve_considered(
        _parse_measures(args.measures), _parse_measures(args.exclude_measures)
    )
    uprns = _uprns_for(engine, args.property_ids)
    postcodes = _postcodes_for(engine, args.property_ids)
    # Landlord Overrides are read from property_overrides and folded onto the lodged
    # EPC to form the Effective EPC the calculator scores (ADR-0032).
    overrides_reader = PropertyOverridesPostgresReader(lambda: Session(engine))
    _dump_overrides(engine, args.property_ids)
    # EPC Prediction (Path 3, ADR-0031): when a Property has no lodged EPC, an
    # EpcPropertyData is synthesised from its postcode cohort. The cohort comes
    # from the live EPC API (search-by-postcode + per-cert fetch), memoised per
    # postcode so co-located missing Properties don't refetch the same cohort.
    prediction_attributes = OverrideBackedPredictionAttributesReader(overrides_reader)
    comparables_repo = EpcComparablePropertiesRepository(
        epc_client, geospatial, nearby_postcodes=PostcodesIoClient()
    )
    predictor = EpcPrediction()
    _cohort_cache: dict[str, list[ComparableProperty]] = {}
    _nearby_cohort_cache: dict[tuple[str, str], list[ComparableProperty]] = {}

    def cohort_for(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]:
        # Broadened cohort for a gated-out target: the nearest real postcodes,
        # walked until enough same-type comparables surface (ADR-0031). Memoised
        # per (postcode, property_type).
        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]
    # One read-only session for the live `material` catalogue, reused across the
    # batch so both store and no-store runs price against the same DB rows.
    catalogue_session = Session(engine)
    products = catalogue_with_off_catalogue_overrides(catalogue_session)
    # When persisting, a lodged-EPC Property also gets a Baseline Performance row
    # via the production orchestrator (lodged-vs-calculator SAP, the bill block) —
    # the same establish-and-persist the first-run pipeline runs, here per Property
    # so one bad cert doesn't abort the batch. Predicted Properties have no lodged
    # figures, so they are skipped below.
    baseline_orchestrator: Optional[PropertyBaselineOrchestrator] = None
    if args.persist:
        baseline_orchestrator = PropertyBaselineOrchestrator(
            unit_of_work=lambda: PostgresUnitOfWork(lambda: Session(engine)),
            rebaseliner=CalculatorRebaseliner(Sap10Calculator()),
            fuel_rates=FuelRatesStaticFileRepository(),
        )
    scenario: Optional[Scenario] = (
        _scenario_for(catalogue_session, args.scenario_id)
        if args.scenario_id is not None
        else None
    )
    # The Scenario's own exclusions drive which measures the run considers; the
    # --measures/--exclude-measures flags are an optional override layered on top.
    considered = combine_considered_measures(
        scenario.considered_measures() if scenario is not None else None,
        cli_considered,
    )

    target = (
        f"scenario {scenario.id} (band {scenario.goal_value})"
        if scenario is not None
        else f"synthesised Increasing-EPC band {args.goal}"
    )
    measures_note = ",".join(sorted(considered)) if considered else "all measures"
    mode = "PERSISTING to DB" if args.persist else "no DB writes"
    source = "persisted DB inputs" if args.from_db else "live EPC/solar"
    print(
        f"modelling {len(args.property_ids)} propertie(s) · {target} · {measures_note} · "
        f"{mode} (DB material catalogue, {source})...\n"
    )

    if args.from_db:
        # Read inputs from the DB and skip every live fetcher (no gov-API calls).
        # Self-contained loop + file writing; the live path below is left as-is.
        _run_from_db(
            args,
            engine=engine,
            products=products,
            scenario=scenario,
            considered=considered,
            baseline_orchestrator=baseline_orchestrator,
            md_path=md_path,
            csv_path=csv_path,
            candidates_path=candidates_path,
            target=target,
            measures_note=measures_note,
        )
        catalogue_session.close()
        return

    md_lines: list[str] = [f"# Modelling recommendations ({target}, {measures_note})\n"]
    csv_rows: list[str] = [
        "property_id,uprn,api_sap,baseline_sap,sap_delta,post_sap,measures,"
        "measure_types,cost_of_works"
    ]
    candidate_csv_rows: list[str] = [
        "property_id,uprn,surface,measure_type,cost_total,contingency_rate,"
        "selected,description"
    ]

    total = len(args.property_ids)
    run_start = time.monotonic()
    errors = 0
    for index, property_id in enumerate(args.property_ids, start=1):
        elapsed = time.monotonic() - run_start
        rate = elapsed / (index - 1) if index > 1 else 0.0
        eta = rate * (total - index + 1)
        bar_done = int(28 * (index - 1) / total)
        bar = "#" * bar_done + "-" * (28 - bar_done)
        print(
            f"[{bar}] {index}/{total} ({100 * (index - 1) / total:.1f}%) "
            f"· {errors} err · elapsed {elapsed / 60:.1f}m · ETA {eta / 60:.1f}m "
            f"· property {property_id}",
            flush=True,
        )
        uprn = uprns.get(property_id)
        try:
            if uprn is None:
                raise ValueError("no UPRN on the property row")
            postcode = postcodes.get(property_id, "")
            # Resolve the spatial reference once: its planning protections gate
            # measures, and its coordinates both drive solar AND distance-weight
            # the EPC-Prediction cohort, so resolve before the EPC branch.
            spatial: Optional[SpatialReference] = _spatial_for(geospatial, uprn)
            restrictions: PlanningRestrictions = (
                spatial.restrictions if spatial is not None else PlanningRestrictions()
            )
            coordinates: Optional[Coordinates] = (
                spatial.coordinates if spatial is not None else None
            )
            overrides = overlays_from(overrides_reader.overrides_for(property_id))
            epc: Optional[EpcPropertyData] = epc_client.get_by_uprn(uprn)
            predicted = False
            if epc is not None:
                # Lodged EPC: fold any Landlord Overrides onto it; with none, the
                # Effective EPC is the lodged EPC unchanged (ADR-0032).
                overlaid_property = Property(
                    identity=PropertyIdentity(
                        portfolio_id=args.portfolio_id or 0,
                        postcode=postcode,
                        address="",
                        uprn=uprn,
                    ),
                    epc=epc,
                    landlord_overrides=overrides,
                )
                effective_epc: EpcPropertyData = overlaid_property.effective_epc
                lodged_wall = _main_wall_summary(epc)
                effective_wall = _main_wall_summary(effective_epc)
                if lodged_wall != effective_wall:
                    print(
                        f"  overlay moved the main wall: lodged [{lodged_wall}] "
                        f"-> effective [{effective_wall}]"
                    )
                else:
                    print(f"  overlay no-op on main wall: [{lodged_wall}]")
            else:
                # No lodged EPC: synthesise one from the postcode cohort
                # (EPC Prediction Path 3, ADR-0031).
                predicted_epc = _predict_epc(
                    property_id=property_id,
                    uprn=uprn,
                    postcode=postcode,
                    portfolio_id=args.portfolio_id or 0,
                    attributes_reader=prediction_attributes,
                    coordinates=coordinates,
                    cohort_for=cohort_for,
                    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}')"
                    )
                # Property.effective_epc folds any Landlord Overrides onto the
                # synthesised EPC (cohort fills the unknown fields, the landlord's
                # known facts correct them) — same overlay the lodged path applies.
                effective_epc = Property(
                    identity=PropertyIdentity(
                        portfolio_id=args.portfolio_id or 0,
                        postcode=postcode,
                        address="",
                        uprn=uprn,
                    ),
                    epc=None,
                    predicted_epc=predicted_epc,
                    landlord_overrides=overrides,
                ).effective_epc
                predicted = True
                synth_wall = _main_wall_summary(predicted_epc)
                effective_wall = _main_wall_summary(effective_epc)
                if synth_wall != effective_wall:
                    print(
                        f"  no lodged EPC -> synthesised from '{postcode}' cohort; "
                        f"overlay moved wall [{synth_wall}] -> [{effective_wall}]"
                    )
                else:
                    print(
                        f"  no lodged EPC -> synthesised from '{postcode}' cohort "
                        f"(overlay no-op on wall) [{synth_wall}]"
                    )
            solar_insights: Optional[dict[str, Any]] = (
                None if args.no_solar else _solar_insights_for(solar_client, spatial)
            )
            plan: Plan = run_modelling(
                effective_epc,
                goal_band=args.goal,
                planning_restrictions=restrictions,
                solar_insights=solar_insights,
                considered_measures=considered,
                products=products,
                scenario=scenario,
                print_table=False,
            )
            # The full candidate menu (every Generator Option + its cost), so
            # measures the Optimiser did not select are still visible. A predicted
            # Property has no lodged cert, so the synthesised Effective EPC is used.
            candidates: list[Recommendation] = candidate_recommendations(
                epc if epc is not None else effective_epc,
                planning_restrictions=restrictions,
                solar_insights=solar_insights,
                considered_measures=considered,
                products=products,
            )
            if args.persist:
                assert scenario is not None  # guaranteed by the --persist guard
                _persist(
                    engine,
                    property_id=property_id,
                    uprn=uprn,
                    portfolio_id=args.portfolio_id,
                    scenario=scenario,
                    epc=epc,
                    spatial=spatial,
                    solar_insights=solar_insights,
                    plan=plan,
                )
                # A lodged EPC also gets its Baseline Performance persisted
                # (reads the EPC just saved above). Predicted Properties have no
                # lodged figures to baseline, so they are skipped.
                if epc is not None:
                    assert baseline_orchestrator is not None
                    baseline_orchestrator.run([property_id])
        except (
            Exception
        ) as error:  # noqa: BLE001 — one bad property must not stop the run
            # A failed catalogue query (e.g. a `material.type` enum mismatch)
            # aborts the shared session's transaction; without a rollback every
            # subsequent property reports `InFailedSqlTransaction` and masks its
            # own real error. Reset so each property surfaces what's wrong.
            catalogue_session.rollback()
            errors += 1
            line = f"property {property_id} (uprn {uprn}): ERROR — {type(error).__name__}: {error}"
            print(line + "\n")
            md_lines.append(f"## Property {property_id}\n\n`{line}`\n")
            csv_rows.append(f"{property_id},{uprn or ''},,,,,,ERROR,")
            continue

        measure_types = [m.measure_type for m in plan.measures]
        selected: set[MeasureType] = {m.measure_type for m in plan.measures}
        context = _context_summary(spatial, solar_insights)
        # Flag EPC-Prediction properties so a synthesised SAP is never mistaken
        # for one scored off a lodged cert.
        source_tag = " · ⚠ PREDICTED (no lodged EPC)" if predicted else ""
        candidate_lines = _candidate_lines(candidates, selected)
        header = (
            f"=== Property {property_id} (uprn {uprn}) === "
            f"SAP {plan.baseline.sap_continuous:.1f} -> {plan.post_sap_continuous:.1f} "
            f"· {len(plan.measures)} measure(s) · £{plan.cost_of_works:,.0f} · {context}"
            f"{source_tag}"
        )
        print(header)
        print(format_plan_table(plan))
        print(f"  candidate measures considered ({len(candidate_lines)} option(s)):")
        for candidate_line in candidate_lines:
            print(candidate_line)
        print()

        md_lines.append(f"## Property {property_id} (uprn {uprn}){source_tag}\n")
        md_lines.append(
            f"SAP {plan.baseline.sap_continuous:.1f} → {plan.post_sap_continuous:.1f} "
            f"· {len(plan.measures)} measure(s) · cost £{plan.cost_of_works:,.0f} "
            f"· {context}\n"
        )
        md_lines.append("**Selected Plan**\n")
        md_lines.extend(_measure_summary(m) for m in plan.measures)
        md_lines.append("")
        md_lines.append("**All candidate measures (cost per measure)**\n")
        md_lines.extend(candidate_lines)
        md_lines.append("")
        # api_sap is the lodged/register SAP (off the cert); a predicted Property
        # has none, so it and the delta are left blank. baseline_sap is the
        # calculator's score on the Effective EPC — the two whose divergence the
        # run is for reviewing (mirrors lodged vs effective in the baseline table).
        api_sap: Optional[int] = epc.energy_rating_current if epc is not None else None
        calc_sap: float = plan.baseline.sap_continuous
        api_sap_cell = "" if api_sap is None else str(api_sap)
        sap_delta_cell = "" if api_sap is None else f"{calc_sap - api_sap:.2f}"
        csv_rows.append(
            f"{property_id},{uprn},{api_sap_cell},{calc_sap:.2f},{sap_delta_cell},"
            f"{plan.post_sap_continuous:.2f},{len(plan.measures)},"
            f"{'|'.join(measure_types)},{plan.cost_of_works:.0f}"
        )
        candidate_csv_rows.extend(
            _candidate_csv_rows(property_id, uprn, candidates, selected)
        )

    catalogue_session.close()
    md_path.write_text("\n".join(md_lines) + "\n", encoding="utf-8")
    csv_path.write_text("\n".join(csv_rows) + "\n", encoding="utf-8")
    candidates_path.write_text(
        "\n".join(candidate_csv_rows) + "\n", encoding="utf-8"
    )
    print(f"wrote {md_path.resolve()}")
    print(f"wrote {csv_path.resolve()}")
    print(f"wrote {candidates_path.resolve()}")


if __name__ == "__main__":
    main()