Model/scripts/hyde_epc_schema_versions.py

"""Tally the EPC schema versions across the hyde list (manipulation_filled UPRNs).

For every resolved UPRN we look up its EPC certificate's ``schemaType`` (e.g.
``RdSAP-Schema-21.0.1``, ``RdSAP-Schema-17.1``, ``SAP-Schema-16.2``). The
gov EPC ``/api/domestic/search`` endpoint returns ``schemaType`` per row, so one
search-per-postcode covers every UPRN in that postcode — far cheaper than a
certificate fetch per UPRN. The latest cert (max registrationDate) wins per UPRN.

Outputs: a per-schema-version tally with one example UPRN each, plus a CSV
mapping every UPRN -> schema version.

    python -m scripts.hyde_epc_schema_versions
    python -m scripts.hyde_epc_schema_versions --workers 8 --out scripts/hyde_schema_versions.csv

Reads OPEN_EPC_API_TOKEN from backend/.env. Run from the worktree root.
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from collections import Counter, defaultdict
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Optional

import httpx
from dotenv import load_dotenv

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

from scripts.fill_domna_addresses import clean_postcode  # noqa: E402
from scripts.finalise_to_property_table import load_rows  # noqa: E402

_BASE = "https://api.get-energy-performance-data.communities.gov.uk"
_SEARCH = f"{_BASE}/api/domestic/search"
NOT_IN_EPC = "NOT_IN_EPC"

_DEFAULT_IN = _REPO_ROOT / "scripts" / "manipulation_filled.xlsx"
_DEFAULT_OUT = _REPO_ROOT / "scripts" / "hyde_schema_versions.csv"


def search_postcode(
    client: httpx.Client, postcode: str, headers: dict[str, str]
) -> list[dict[str, Any]]:
    """Return the search rows for a postcode, retrying on rate-limit (429)."""
    for attempt in range(5):
        resp = client.get(_SEARCH, params={"postcode": postcode}, headers=headers, timeout=30)
        if resp.status_code == 429:
            retry_after = float(resp.headers.get("Retry-After", "2"))
            time.sleep(min(retry_after, 10) * (attempt + 1))
            continue
        # 400 = malformed postcode (data-entry typo), 404 = no certs — skip both.
        if resp.status_code in (400, 404):
            return []
        resp.raise_for_status()
        return resp.json().get("data", [])
    return []


def build_uprn_schema_map(
    postcodes: list[str], token: str, workers: int
) -> dict[int, tuple[str, str]]:
    """Map UPRN -> (schemaType, registrationDate) for the latest cert per UPRN.

    One search per postcode (concurrent); later we look our UPRNs up in here.
    """
    headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"}
    by_uprn: dict[int, tuple[str, str]] = {}
    done = 0
    total = len(postcodes)

    def fetch(pc: str) -> list[dict[str, Any]]:
        with httpx.Client() as client:
            return search_postcode(client, pc, headers)

    with ThreadPoolExecutor(max_workers=workers) as pool:
        for rows in pool.map(fetch, postcodes):
            for row in rows:
                uprn = row.get("uprn")
                schema = row.get("schemaType")
                reg = row.get("registrationDate") or ""
                if uprn is None or not schema:
                    continue
                prev = by_uprn.get(int(uprn))
                # Keep the latest-registered cert's schema for this UPRN.
                if prev is None or reg > prev[1]:
                    by_uprn[int(uprn)] = (str(schema), str(reg))
            done += 1
            if done % 250 == 0:
                print(f"  searched {done}/{total} postcodes, {len(by_uprn)} uprns seen")
    return by_uprn


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--in", dest="inp", type=Path, default=_DEFAULT_IN)
    parser.add_argument("--out", type=Path, default=_DEFAULT_OUT)
    parser.add_argument("--workers", type=int, default=8)
    return parser.parse_args()


def main() -> int:
    args = _parse_args()
    load_dotenv(_REPO_ROOT / "backend" / ".env")
    token = os.environ.get("OPEN_EPC_API_TOKEN")
    if not token:
        print("OPEN_EPC_API_TOKEN not set (backend/.env)")
        return 2

    _, rows = load_rows(args.inp, include_unmatched=False)
    pairs: list[tuple[int, str, str]] = []  # (uprn, postcode_clean, address)
    for r in rows:
        uprn = r["address2uprn_uprn"]
        if uprn:
            pairs.append((int(uprn), clean_postcode(r["postcode"]), r["address2uprn_address"]))
    postcodes = sorted({pc for _, pc, _ in pairs if pc})
    print(f"{len(pairs)} UPRNs across {len(postcodes)} unique postcodes")

    by_uprn = build_uprn_schema_map(postcodes, token, args.workers)
    print(f"EPC search returned schema for {len(by_uprn)} distinct UPRNs")

    # Resolve each hyde UPRN to its schema version.
    tally: Counter[str] = Counter()
    example: dict[str, tuple[int, str]] = {}
    out_lines: list[tuple[int, str, str, str]] = []  # uprn, schema, postcode, address
    seen: set[int] = set()
    for uprn, pc, address in pairs:
        if uprn in seen:
            continue
        seen.add(uprn)
        schema = by_uprn.get(uprn, (NOT_IN_EPC, ""))[0]
        tally[schema] += 1
        example.setdefault(schema, (uprn, address))
        out_lines.append((uprn, schema, pc, address))

    # Write the full per-UPRN mapping.
    import csv

    with args.out.open("w", newline="", encoding="utf-8") as fh:
        w = csv.writer(fh)
        w.writerow(["uprn", "schema_version", "postcode", "matched_address"])
        w.writerows(out_lines)

    print(f"\nSchema versions across {len(seen)} distinct UPRNs:\n")
    print(f"  {'schema version':<26} {'count':>7}   example UPRN")
    print(f"  {'-'*26} {'-'*7}   {'-'*12}")
    for schema, count in tally.most_common():
        ex_uprn, ex_addr = example[schema]
        print(f"  {schema:<26} {count:>7}   {ex_uprn}  ({ex_addr})")
    print(f"\nFull mapping -> {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())