Model/scripts/reclassify_main_heating.py

"""One-time re-classification of main-heating overrides funnelled into a garbage
drawer for want of a correct archetype (#1376).

Covers both dumping grounds the ``main_heating_guard`` now rescues:
- **HHRSH** (ADR-0044): "High heat retention storage heaters" → "Electric storage
  heaters, old" (SAP 401), the worst old type — 111 rows.
- **Gas CPSU** (ADR-0045): solid-fuel room heaters (open fire / + back boiler /
  closed + boiler) and electric "NA"-type boilers → "Gas CPSU" (120, mains gas).

The live classifier now applies ``main_heating_guard`` deterministically (so new
intakes are correct); this fixes the rows written before it. The **same guard**
decides the correction here, so the backfill and the live path cannot drift.

Updates the TEXT ``property_overrides.override_value`` (what the modelling reads —
the immediate fix) unconditionally. The ``landlord_main_heating_system_overrides``
``value`` classifier cache is a ``main_heating_system`` **pgEnum**; the new
archetypes are FE-owned values, so their cache writes are **deferred** until the
Drizzle migration adds them (the Class-A/B pattern).

DRY-RUN BY DEFAULT: prints the counts it would change and writes nothing. Pass
``--apply`` to execute inside a transaction. Idempotent — only rows whose stored
value differs from the target member are touched.
"""

from __future__ import annotations

import argparse
from collections.abc import Iterable
from typing import Optional

from sqlalchemy import Connection, text

from domain.epc.property_overrides.main_heating_guard import main_heating_guard
from domain.epc.property_overrides.main_heating_system_type import (
    MainHeatingSystemType,
)
from scripts.e2e_common import build_engine, load_env


def electric_cpsu_for(description: str, main_fuel: Optional[str]) -> Optional[str]:
    """``"Electric CPSU"`` for a CPSU-boiler override whose property's ``main_fuel``
    override is electricity, else ``None`` (ADR-0045). ``"Boiler: A/C rated CPSU"``
    is identical bar the band — the electric/gas split is *fuel*-determined, so this
    keys on the fuel the guard cannot see; a gas / unknown fuel keeps Gas CPSU
    (120), and a non-CPSU description is left to the guard."""
    if "cpsu" in description.lower() and main_fuel == "electricity":
        return MainHeatingSystemType.ELECTRIC_CPSU.value
    return None


def main_heating_corrections(
    stored: Iterable[tuple[str, str]],
) -> dict[str, str]:
    """``(description, stored override_value)`` → the faithful archetype value, for
    the descriptions the main-heating guard resolves (HHRSH, solid-fuel room
    heaters, electric boiler) whose stored value differs. Descriptions the guard
    defers and rows already on the target are omitted, so re-running against
    corrected data is a no-op."""
    corrections: dict[str, str] = {}
    for description, value in stored:
        member = main_heating_guard(description)
        if member is not None and value != member.value:
            corrections[description] = member.value
    return corrections


_DISTINCT = text(
    """
    SELECT DISTINCT lower(original_spreadsheet_description) AS description,
                    override_value AS value
    FROM property_overrides
    WHERE override_component = 'main_heating_system'
    """
)
_OVERRIDES_UPDATE = text(
    """
    UPDATE property_overrides
       SET override_value = :new_value
     WHERE override_component = 'main_heating_system'
       AND lower(original_spreadsheet_description) = :description
       AND override_value <> :new_value
    """
)
_OVERRIDES_COUNT = text(
    """
    SELECT count(*) FROM property_overrides
     WHERE override_component = 'main_heating_system'
       AND lower(original_spreadsheet_description) = :description
       AND override_value <> :new_value
    """
)
_CACHE_UPDATE = text(
    """
    UPDATE landlord_main_heating_system_overrides
       SET value = :new_value, updated_at = now()
     WHERE lower(description) = :description
       AND value::text <> :new_value
    """
)
_ENUM_VALUES = text(
    "SELECT e.enumlabel FROM pg_enum e JOIN pg_type t ON t.oid = e.enumtypid "
    "WHERE t.typname = 'main_heating_system'"
)
# CPSU boilers still on Gas CPSU, each with its property's main_fuel override —
# the electric/gas split the description can't carry (ADR-0045).
_CPSU_ROWS = text(
    """
    SELECT mh.property_id AS property_id,
           lower(mh.original_spreadsheet_description) AS description,
           mf.override_value AS main_fuel
    FROM property_overrides mh
    LEFT JOIN property_overrides mf
      ON mf.property_id = mh.property_id AND mf.override_component = 'main_fuel'
    WHERE mh.override_component = 'main_heating_system'
      AND mh.override_value = 'Gas CPSU'
      AND lower(mh.original_spreadsheet_description) LIKE '%cpsu%'
    """
)
_CPSU_PROPERTY_UPDATE = text(
    """
    UPDATE property_overrides
       SET override_value = :new_value
     WHERE property_id = :property_id
       AND override_component = 'main_heating_system'
       AND override_value = 'Gas CPSU'
    """
)


def reclassify(conn: Connection, *, apply: bool) -> tuple[int, set[str]]:
    """Re-map main-heating overrides off their garbage-drawer archetypes onto the
    faithful ones the guard resolves. Returns the number of ``property_overrides``
    rows found and the set of target values the live ``main_heating_system`` pgEnum
    does not yet carry (cache-deferred until the FE migration)."""
    stored = [(r.description, r.value) for r in conn.execute(_DISTINCT)]
    enum_values = {r[0] for r in conn.execute(_ENUM_VALUES)}
    total = 0
    deferred: set[str] = set()
    for description, new_value in main_heating_corrections(stored).items():
        params = {"description": description, "new_value": new_value}
        total += conn.execute(_OVERRIDES_COUNT, params).scalar() or 0
        in_enum = new_value in enum_values
        if not in_enum:
            deferred.add(new_value)
        if apply:
            conn.execute(_OVERRIDES_UPDATE, params)
            if in_enum:
                conn.execute(_CACHE_UPDATE, params)

    # Fuel-aware phase: an electric CPSU (fuel-determined, per-property) — update
    # only the TEXT layer. The description-keyed classifier cache cannot represent
    # a fuel-dependent value, and "Electric CPSU" is a deferred pgEnum value anyway.
    for row in conn.execute(_CPSU_ROWS):
        new_value = electric_cpsu_for(row.description, row.main_fuel)
        if new_value is None:
            continue
        total += 1
        deferred.add(new_value)
        if apply:
            conn.execute(
                _CPSU_PROPERTY_UPDATE,
                {"property_id": row.property_id, "new_value": new_value},
            )
    return total, deferred


def main() -> None:
    load_env()
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--apply",
        action="store_true",
        help="execute the updates (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, deferred = reclassify(conn, apply=args.apply)

    verb = "re-classified" if args.apply else "would re-classify"
    print(
        f"{verb} {total} main-heating override row(s) off their garbage-drawer "
        "archetypes (old storage / Gas CPSU) onto faithful HHRSH / solid-fuel / "
        "electric-boiler archetypes — property_overrides / TEXT, what the "
        "modelling reads."
    )
    if deferred:
        print(
            f"\n{len(deferred)} target value(s) NOT yet in the main_heating_system "
            "pgEnum — their classifier-cache rows are deferred until the FE-repo "
            "enum migration adds these members (property_overrides was still "
            "updated, which is what the modelling reads):"
        )
        for value in sorted(deferred):
            print(f"    {value!r}")
    if not args.apply:
        print("\nDRY-RUN — nothing written. Re-run with --apply to execute.")


if __name__ == "__main__":
    main()