Model/scripts/backfill_boiler_efficiency_band.py

"""Backfill the ``boiler_efficiency_band`` override for already-ingested boilers.

The SEDBUK band (ADR-0068) is losslessly present in every ``main_heating_system``
override's ``original_spreadsheet_description`` (``Boiler: C rated Combi``) but was
dropped at classification. This one-off populates the new ``boiler_efficiency_band``
override rows for the existing stock by re-parsing that text with the SAME guard the
live classifier uses (``boiler_efficiency_band_guard``), so the backfill and the
forward path cannot drift.

One ``boiler_efficiency_band`` row per ``main_heating_system`` row, keyed to the
same ``(property_id, building_part)`` — the parseable band (``A``–``G``) where the
description carries one, else an explicit ``Unknown``. This mirrors the finaliser
(which records Unknown rather than skipping), so backfilled historical data matches
go-forward data. The modelling gate (gas/oil boilers only, ADR-0068) is applied at
overlay time, not here — so a band parsed off an electric ``Boiler: A rated NA`` is
stored for fidelity but ignored by the calculator, and ``Unknown`` (incl. on a
non-boiler system) is inert to the calculator, exactly as the live path stores them.

GATED: the ``boiler_efficiency_band`` value of the FE-owned ``override_component``
pgEnum must exist first (the Drizzle migration in the assessment-model repo). Until
then this writes nothing in prod — the Class-A/B deferred pattern.

DRY-RUN BY DEFAULT: prints the counts it would write and writes nothing. Pass
``--apply`` to execute inside a transaction. Idempotent — a row already carrying the
target band is left untouched, so re-running is a no-op.
"""

from __future__ import annotations

import argparse
from collections.abc import Iterable
from dataclasses import dataclass

from sqlalchemy import Connection, text

from domain.epc.property_overrides.boiler_efficiency_band import BoilerEfficiencyBand
from domain.epc.property_overrides.boiler_efficiency_band_guard import (
    boiler_efficiency_band_guard,
)
from scripts.e2e_common import build_engine, load_env

_OVERRIDE_COMPONENT = "boiler_efficiency_band"


@dataclass(frozen=True)
class MainHeatingRow:
    """The fields of a ``main_heating_system`` ``property_overrides`` row the
    backfill reads."""

    property_id: int
    portfolio_id: int
    building_part: int
    original_spreadsheet_description: str


@dataclass(frozen=True)
class BandBackfillRow:
    """A ``boiler_efficiency_band`` override row to upsert for a boiler."""

    property_id: int
    portfolio_id: int
    building_part: int
    override_value: str
    original_spreadsheet_description: str


def band_backfill_rows(rows: Iterable[MainHeatingRow]) -> list[BandBackfillRow]:
    """The ``boiler_efficiency_band`` row to write for **every**
    ``main_heating_system`` row, keyed to the same property + building part — the
    parseable SEDBUK band (``A``–``G``) where the description carries one, else an
    explicit ``Unknown``. This mirrors the finaliser (which records Unknown rather
    than skipping), so backfilled historical data matches go-forward data;
    ``Unknown`` is inert to modelling (no efficiency anchor) and fine on a
    non-boiler system.

    Pure and DB-free, so the parse/decide logic is unit-tested without a
    database; the connection wrapper below just feeds it rows and upserts."""
    backfill: list[BandBackfillRow] = []
    for row in rows:
        band = boiler_efficiency_band_guard(row.original_spreadsheet_description)
        value = (
            band.value
            if band is not None and band is not BoilerEfficiencyBand.UNKNOWN
            else BoilerEfficiencyBand.UNKNOWN.value
        )
        backfill.append(
            BandBackfillRow(
                property_id=row.property_id,
                portfolio_id=row.portfolio_id,
                building_part=row.building_part,
                override_value=value,
                original_spreadsheet_description=row.original_spreadsheet_description,
            )
        )
    return backfill


_SELECT_MAIN_HEATING = text(
    """
    SELECT property_id, portfolio_id, building_part, original_spreadsheet_description
    FROM property_overrides
    WHERE override_component = 'main_heating_system'
    """
)
# Idempotent upsert on the (property, component, part) unique constraint — a row
# already carrying the target band is left untouched (updated_at unchanged).
_UPSERT_BAND = text(
    """
    INSERT INTO property_overrides
        (id, property_id, portfolio_id, building_part, override_component,
         override_value, original_spreadsheet_description, created_at, updated_at)
    VALUES
        (gen_random_uuid(), :property_id, :portfolio_id, :building_part,
         'boiler_efficiency_band', :override_value, :original_spreadsheet_description,
         now(), now())
    ON CONFLICT (property_id, override_component, building_part)
    DO UPDATE SET override_value = EXCLUDED.override_value,
                  original_spreadsheet_description =
                      EXCLUDED.original_spreadsheet_description,
                  updated_at = now()
    WHERE property_overrides.override_value <> EXCLUDED.override_value
    """
)


def backfill(conn: Connection, *, apply: bool) -> int:
    """Backfill the ``boiler_efficiency_band`` rows off the existing
    ``main_heating_system`` descriptions.

    In ``--apply`` returns the number of rows **actually written** — the upsert's
    ``WHERE override_value <> EXCLUDED.override_value`` no-ops rows already at the
    target band, so an idempotent re-run reports 0, not the candidate count. In
    dry-run returns the number of **candidate** rows (it cannot know how many
    differ without writing)."""
    rows = [
        MainHeatingRow(
            property_id=r.property_id,
            portfolio_id=r.portfolio_id,
            building_part=r.building_part,
            original_spreadsheet_description=r.original_spreadsheet_description or "",
        )
        for r in conn.execute(_SELECT_MAIN_HEATING)
    ]
    to_write = band_backfill_rows(rows)
    if not apply:
        return len(to_write)
    written = 0
    for band_row in to_write:
        result = conn.execute(
            _UPSERT_BAND,
            {
                "property_id": band_row.property_id,
                "portfolio_id": band_row.portfolio_id,
                "building_part": band_row.building_part,
                "override_value": band_row.override_value,
                "original_spreadsheet_description": (
                    band_row.original_spreadsheet_description
                ),
            },
        )
        written += result.rowcount or 0
    return written


def main() -> None:
    load_env()
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--apply",
        action="store_true",
        help="execute the writes (default: dry-run, writes nothing)",
    )
    args = parser.parse_args()

    engine = build_engine()
    with engine.begin() as conn:
        conn.execute(text("SET statement_timeout = 120000"))
        total = backfill(conn, apply=args.apply)

    noun = "row(s) written" if args.apply else "candidate row(s)"
    print(
        f"{total} boiler_efficiency_band {noun} from the existing "
        "main_heating_system descriptions (Boiler: <A-G> rated ...). The modelling "
        "gate (gas/oil boilers only) is applied at overlay time (ADR-0068)."
    )
    if args.apply:
        print(
            "(Rows already at their target band are no-ops, so a re-run reports 0.)"
        )
    else:
        print("\nDRY-RUN — nothing written. Re-run with --apply to execute.")


if __name__ == "__main__":
    main()