Model/scripts/resolve_uprns_for_finaliser.py

"""Resolve a CSV of addresses to UPRNs, ready to feed the bulk-upload finaliser.

Takes a CSV with `Address 1/2/3` + `postcode` columns and, per row, resolves a
UPRN by trying — in order — the new EPC API (address2uprn), the historic EPC S3
dataset, then the Ordnance Survey Places API as a fallback. Whichever source
wins, the result is written into the SAME three columns the finaliser reads
(`bulk_upload_finaliser_orchestrator`):

    address2uprn_uprn        UPRN integer (empty when unresolved)
    address2uprn_address     the matched address
    address2uprn_lexiscore   the match score in [0, 1]

A `resolution_source` diagnostic column (epc / epc_historic / ordnance_survey /
none) is appended too — the finaliser ignores unknown columns. All original
columns are preserved in their original order, so the output CSV drops straight
into the finaliser.

    python -m scripts.resolve_uprns_for_finaliser input.csv -o resolved.csv

    # OS-only / EPC-only, custom postcode column, custom OS score threshold
    python -m scripts.resolve_uprns_for_finaliser in.csv -o out.csv --no-epc
    python -m scripts.resolve_uprns_for_finaliser in.csv -o out.csv --postcode-col Postcode --os-threshold 0.6

Keys are read from backend/.env: OPEN_EPC_API_TOKEN (EPC) and
ORDNANCE_SURVEY_API_KEY (OS Places). Run from the worktree root (import trap).

The module-level functions (`load_keys`, `read_rows`, `resolve_row`, `process`,
`write_rows`) are written to be driven line-by-line from a REPL as well as via
the CLI.
"""

from __future__ import annotations

import argparse
import csv
import os
import sys
from pathlib import Path
from typing import Optional

import pandas as pd
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 backend.address2UPRN.main import (  # noqa: E402
    get_epc_data_with_postcode,
    get_uprn_from_historic_epc,
    get_uprn_with_epc_df,
)
from backend.ordnanceSurvey.helpers import (  # noqa: E402
    lookup_os_places,
    os_places_results_to_dataframe,
)
from backend.utils.addressMatch import AddressMatch  # noqa: E402

# Columns the finaliser reads (bulk_upload_finaliser_orchestrator).
UPRN_COL = "address2uprn_uprn"
MATCHED_ADDRESS_COL = "address2uprn_address"
LEXISCORE_COL = "address2uprn_lexiscore"
SOURCE_COL = "resolution_source"
_RESULT_COLS = (UPRN_COL, MATCHED_ADDRESS_COL, LEXISCORE_COL, SOURCE_COL)

# A resolved hit: (uprn, matched_address, lexiscore, source).
Resolution = tuple[str, str, float, str]


def load_keys() -> tuple[Optional[str], Optional[str]]:
    """Load (epc_token, os_api_key) from backend/.env (and the process env)."""
    load_dotenv(_REPO_ROOT / "backend" / ".env")
    epc_token = os.environ.get("OPEN_EPC_API_TOKEN")
    os_api_key = os.environ.get("ORDNANCE_SURVEY_API_KEY")
    return epc_token, os_api_key


def read_rows(path: Path) -> tuple[list[dict[str, str]], list[str]]:
    """Read a CSV into (rows, fieldnames). Preserves column order."""
    with path.open(newline="", encoding="utf-8-sig") as fh:
        reader = csv.DictReader(fh)
        fieldnames = list(reader.fieldnames or [])
        rows = [dict(row) for row in reader]
    return rows, fieldnames


def clean_postcode(postcode: str) -> str:
    """Sanitise to the no-space upper form the EPC/OS lookups expect (e.g. E84SQ)."""
    return postcode.upper().replace(" ", "").strip()


def build_address(row: dict[str, str]) -> str:
    """Concatenate Address 1/2/3 the same way the address2uprn lambda does."""
    return " ".join(
        str(row.get(col, "") or "").strip() for col in ("Address 1", "Address 2", "Address 3")
    ).strip()


def resolve_epc(
    address: str, postcode_clean: str, epc_cache: dict[str, pd.DataFrame]
) -> Optional[Resolution]:
    """Resolve via the new EPC API (cached per postcode), then historic EPC S3.

    `epc_cache` is mutated to memoise one EPC API call per postcode — pass the
    same dict across rows so a postcode is only fetched once.
    """
    epc_df = epc_cache.get(postcode_clean)
    if epc_df is None:
        epc_df = get_epc_data_with_postcode(postcode=postcode_clean)
        epc_cache[postcode_clean] = epc_df

    result = get_uprn_with_epc_df(
        user_inputed_address=address, epc_df=epc_df, verbose=True
    )
    if isinstance(result, tuple):
        uprn, matched, score = result
        return str(uprn), str(matched), float(score), "epc"

    historic = get_uprn_from_historic_epc(
        user_inputed_address=address, postcode=postcode_clean
    )
    if historic is not None:
        uprn, matched, score = historic
        return str(uprn), str(matched), float(score), "epc_historic"

    return None


def resolve_os(
    address: str,
    postcode_clean: str,
    os_api_key: str,
    os_cache: dict[str, pd.DataFrame],
    threshold: float,
) -> Optional[Resolution]:
    """Resolve via the OS Places API: best-scoring address above `threshold`.

    `os_cache` memoises one OS Places call per postcode.
    """
    places_df = os_cache.get(postcode_clean)
    if places_df is None:
        response = lookup_os_places(postcode_clean, os_api_key)
        if response.get("status") != 200 or "data" not in response:
            places_df = pd.DataFrame()
        else:
            places_df = os_places_results_to_dataframe(response["data"])
        os_cache[postcode_clean] = places_df

    if places_df.empty or "ADDRESS" not in places_df.columns:
        return None

    # Iterate plain records — avoids pandas' partially-unknown indexing types.
    records: list[dict[str, object]] = places_df.to_dict(orient="records")
    best: Optional[Resolution] = None
    for rec in records:
        candidate = str(rec.get("ADDRESS", ""))
        score = AddressMatch.score(address, candidate)
        if score >= threshold and (best is None or score > best[2]):
            best = (str(rec.get("UPRN", "")), candidate, score, "ordnance_survey")
    return best


def resolve_row(
    row: dict[str, str],
    *,
    epc_token: Optional[str],
    os_api_key: Optional[str],
    epc_cache: dict[str, pd.DataFrame],
    os_cache: dict[str, pd.DataFrame],
    postcode_col: str,
    use_epc: bool,
    use_os: bool,
    os_threshold: float,
    validate_postcode: bool,
) -> dict[str, str]:
    """Resolve one row in place and return it with the finaliser columns filled.

    Tries EPC (new + historic) first, then OS Places. On no match the three
    result columns are written empty and `resolution_source` is "none".
    """
    address = build_address(row)
    postcode_clean = clean_postcode(str(row.get(postcode_col, "") or ""))

    def write(res: Optional[Resolution]) -> dict[str, str]:
        if res is None:
            row[UPRN_COL] = ""
            row[MATCHED_ADDRESS_COL] = ""
            row[LEXISCORE_COL] = ""
            row[SOURCE_COL] = "none"
        else:
            uprn, matched, score, source = res
            row[UPRN_COL] = uprn
            row[MATCHED_ADDRESS_COL] = matched
            row[LEXISCORE_COL] = str(score)
            row[SOURCE_COL] = source
        return row

    if not address or not postcode_clean:
        return write(None)

    if validate_postcode and not AddressMatch.is_valid_postcode(postcode_clean):
        return write(None)

    if use_epc and epc_token:
        try:
            res = resolve_epc(address, postcode_clean, epc_cache)
            if res is not None:
                return write(res)
        except Exception as exc:  # keep going on a per-row API/lookup failure
            print(f"  EPC lookup failed for {address!r} / {postcode_clean}: {exc}")

    if use_os and os_api_key:
        try:
            res = resolve_os(address, postcode_clean, os_api_key, os_cache, os_threshold)
            if res is not None:
                return write(res)
        except Exception as exc:
            print(f"  OS lookup failed for {address!r} / {postcode_clean}: {exc}")

    return write(None)


def process(
    rows: list[dict[str, str]],
    *,
    epc_token: Optional[str],
    os_api_key: Optional[str],
    postcode_col: str = "postcode",
    use_epc: bool = True,
    use_os: bool = True,
    os_threshold: float = 0.5,
    validate_postcode: bool = True,
) -> list[dict[str, str]]:
    """Resolve every row, printing a per-row line so REPL/CLI progress is visible."""
    epc_cache: dict[str, pd.DataFrame] = {}
    os_cache: dict[str, pd.DataFrame] = {}
    for i, row in enumerate(rows, start=1):
        resolve_row(
            row,
            epc_token=epc_token,
            os_api_key=os_api_key,
            epc_cache=epc_cache,
            os_cache=os_cache,
            postcode_col=postcode_col,
            use_epc=use_epc,
            use_os=use_os,
            os_threshold=os_threshold,
            validate_postcode=validate_postcode,
        )
        print(
            f"[{i}/{len(rows)}] {build_address(row)!r} -> "
            f"{row[UPRN_COL] or '(no match)'} ({row[SOURCE_COL]})"
        )
    return rows


def write_rows(rows: list[dict[str, str]], path: Path, fieldnames: list[str]) -> None:
    """Write rows to CSV, preserving input columns and appending the result columns."""
    out_fields = list(fieldnames)
    for col in _RESULT_COLS:
        if col not in out_fields:
            out_fields.append(col)
    with path.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=out_fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path, help="input CSV (Address 1/2/3 + postcode)")
    parser.add_argument(
        "-o", "--out", type=Path, required=True, help="output CSV for the finaliser"
    )
    parser.add_argument("--postcode-col", default="postcode", help="postcode column name")
    parser.add_argument("--no-epc", action="store_true", help="skip EPC resolution")
    parser.add_argument("--no-os", action="store_true", help="skip Ordnance Survey fallback")
    parser.add_argument(
        "--os-threshold", type=float, default=0.5, help="min OS match score (default 0.5)"
    )
    parser.add_argument(
        "--no-validate-postcode",
        action="store_true",
        help="skip the postcodes.io validity check (one HTTP call per postcode)",
    )
    parser.add_argument("--limit", type=int, default=None, help="process only the first N rows")
    return parser.parse_args()


def main() -> int:
    args = _parse_args()
    epc_token, os_api_key = load_keys()

    use_epc = not args.no_epc
    use_os = not args.no_os
    if use_epc and not epc_token:
        print("OPEN_EPC_API_TOKEN not set (backend/.env) — EPC resolution disabled")
        use_epc = False
    if use_os and not os_api_key:
        print("ORDNANCE_SURVEY_API_KEY not set (backend/.env) — OS fallback disabled")
        use_os = False
    if not use_epc and not use_os:
        print("No resolver enabled (missing keys or both --no-* flags). Nothing to do.")
        return 2

    rows, fieldnames = read_rows(args.input)
    if args.limit is not None:
        rows = rows[: args.limit]
    print(f"Loaded {len(rows)} rows from {args.input}")

    process(
        rows,
        epc_token=epc_token,
        os_api_key=os_api_key,
        postcode_col=args.postcode_col,
        use_epc=use_epc,
        use_os=use_os,
        os_threshold=args.os_threshold,
        validate_postcode=not args.no_validate_postcode,
    )

    write_rows(rows, args.out, fieldnames)
    matched = sum(1 for r in rows if r.get(UPRN_COL))
    print(f"\nResolved {matched}/{len(rows)} rows. Wrote {args.out}")
    return 0


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