Model/scripts/fill_domna_addresses.py

"""Fill the DOMNA columns in the AddressProfilingResults spreadsheet.

Input: scripts/manipulation(2).xlsx, sheet "AddressProfilingResults", columns
    Organisation Reference | UPRN | DOMNA FOUND UPRN | DOMNA FOUND ADDRESS | Address | Postcode

Per-row rule ("if there's a UPRN in the UPRN column we're done"):

  * UPRN present AND Address present   -> nothing to do (already sorted).
  * UPRN present AND Address missing   -> reverse-lookup the address from the UPRN
                                          via the EPC API -> DOMNA FOUND ADDRESS.
  * UPRN missing  AND Address present  -> resolve a UPRN from address + postcode
                                          (EPC API, then Ordnance Survey) -> writes
                                          DOMNA FOUND UPRN + DOMNA FOUND ADDRESS.
  * not resolvable                     -> marked "NOT FOUND" and listed in the
                                          unresolved report.

Relaxed matching (this batch only — production AddressMatch is untouched): the
landlord writes flats as "3 GLADYS COURT" while EPC stores "Flat 3 Gladys
Court", which the production matcher hard-rejects. So per address we try several
query variants — the full string, just the first comma-segment, and a
"Flat <n> ..." form — and keep the best-scoring, unambiguous match. The unit
number must still match exactly (AddressMatch zeroes mismatched numbers), so a
wrong-unit match stays unlikely. Each fill carries its score + source so you can
spot-check (DOMNA SCORE / DOMNA SOURCE).

Rows that already have a DOMNA FOUND UPRN are skipped (idempotent / resumable).

    python -m scripts.fill_domna_addresses
    python -m scripts.fill_domna_addresses --limit 200   # smoke test first N

Keys come from backend/.env (OPEN_EPC_API_TOKEN, ORDNANCE_SURVEY_API_KEY). Run
from the worktree root (import trap).
"""

from __future__ import annotations

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

import pandas as pd

_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 get_epc_data_with_postcode  # noqa: E402
from backend.address2UPRN.scoring import all_uprns_match, rank_address_similarity  # noqa: E402
from backend.ordnanceSurvey.helpers import (  # noqa: E402
    lookup_os_places,
    os_places_results_to_dataframe,
)
from backend.utils.addressMatch import AddressMatch  # noqa: E402
from datatypes.epc.search import EpcSearchResult  # noqa: E402
from infrastructure.epc_client.epc_client_service import EpcClientService  # noqa: E402
from scripts.resolve_uprns_for_finaliser import clean_postcode, load_keys  # noqa: E402

SHEET = "AddressProfilingResults"
UPRN_COL = "UPRN"
ADDRESS_COL = "Address"
POSTCODE_COL = "Postcode"
REF_COL = "Organisation Reference"
FOUND_UPRN_COL = "DOMNA FOUND UPRN"
FOUND_ADDRESS_COL = "DOMNA FOUND ADDRESS"
SCORE_COL = "DOMNA SCORE"
SOURCE_COL = "DOMNA SOURCE"
NOT_FOUND = "NOT FOUND"

# EPC matches are tight (short addresses) so we hold the production 0.7 bar; OS
# addresses carry more trailing tokens, so a slightly lower bar is appropriate.
EPC_THRESHOLD = 0.7
OS_THRESHOLD = 0.6

_DEFAULT_IN = _REPO_ROOT / "scripts" / "manipulation(2).xlsx"
_DEFAULT_OUT = _REPO_ROOT / "scripts" / "manipulation_filled.xlsx"
_DEFAULT_UNRESOLVED = _REPO_ROOT / "scripts" / "manipulation_unresolved.csv"

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


def cell_str(value: object) -> str:
    """Coerce a spreadsheet cell to a trimmed string ("" for NaN/None)."""
    if value is None:
        return ""
    text = str(value).strip()
    return "" if text.lower() == "nan" else text


def parse_uprn_cell(value: object) -> Optional[int]:
    """Read a UPRN cell that pandas loaded as float64 back into an int."""
    text = cell_str(value)
    if not text:
        return None
    try:
        return int(float(text))
    except ValueError:
        return None


def address_variants(address: str) -> list[str]:
    """Query forms to try for one input address, best-discriminating first.

    Landlord flats read "3 GLADYS COURT, 260 REIGATE ROAD" but EPC stores
    "Flat 3 Gladys Court"; the full string scores low (extra tokens) and the
    bare "3 ..." trips the flat guard. So we also try the first comma-segment
    and a "Flat <segment>" form.
    """
    address = address.strip()
    first = address.split(",")[0].strip()
    variants = [address, first]
    if re.match(r"^\d", first):  # starts with a unit/house number
        variants.append("Flat " + first)
        variants.append("Flat " + address)
    seen: set[str] = set()
    out: list[str] = []
    for v in variants:
        key = v.lower()
        if v and key not in seen:
            seen.add(key)
            out.append(v)
    return out


def resolve_epc_relaxed(
    address: str,
    postcode_clean: str,
    epc_cache: dict[str, pd.DataFrame],
    threshold: float = EPC_THRESHOLD,
) -> Optional[Hit]:
    """Best unambiguous EPC match across the address variants (cached per postcode)."""
    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
    if epc_df.empty:
        return None

    best: Optional[Hit] = None
    for variant in address_variants(address):
        scored = rank_address_similarity(epc_df, user_address=variant)
        if scored.empty:
            continue
        score = float(scored.iloc[0]["lexiscore"])
        if best is not None and score <= best[2]:
            continue
        top_rank = scored[scored["lexirank"] == 1]
        # rank-1 rows must agree on one UPRN, else it's ambiguous — skip.
        if not all_uprns_match(top_rank, top_rank.iloc[0]["uprn"]):
            continue
        uprn = str(top_rank.iloc[0]["uprn"])
        if uprn in ("", "nan"):
            continue
        best = (uprn, str(scored.iloc[0]["address"]), score, "epc")

    return best if best is not None and best[2] >= threshold else None


def resolve_os_relaxed(
    address: str,
    postcode_clean: str,
    os_api_key: str,
    os_cache: dict[str, pd.DataFrame],
    threshold: float = OS_THRESHOLD,
) -> Optional[Hit]:
    """Best OS Places match across the address variants (cached 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 and "data" in response:
            places_df = os_places_results_to_dataframe(response["data"])
        else:
            places_df = pd.DataFrame()
        os_cache[postcode_clean] = places_df
    if places_df.empty or "ADDRESS" not in places_df.columns:
        return None

    records: list[dict[str, object]] = places_df.to_dict(orient="records")
    best: Optional[Hit] = None
    for variant in address_variants(address):
        for rec in records:
            candidate = str(rec.get("ADDRESS", ""))
            score = AddressMatch.score(variant, candidate)
            if best is None or score > best[2]:
                best = (str(rec.get("UPRN", "")), candidate, score, "ordnance_survey")
    return best if best is not None and best[2] >= threshold else None


def _address_from_search(result: EpcSearchResult) -> str:
    parts = [
        result.address_line_1,
        result.address_line_2,
        result.address_line_3,
        result.address_line_4,
        result.post_town,
    ]
    return ", ".join(p.strip() for p in parts if p and p.strip())


def reverse_address_from_uprn(
    uprn: int,
    postcode_clean: str,
    service: EpcClientService,
    search_cache: dict[str, list[EpcSearchResult]],
) -> Optional[str]:
    """Find the EPC address for a known UPRN by searching its postcode (cached)."""
    results = search_cache.get(postcode_clean)
    if results is None:
        results = service.search_by_postcode(postcode_clean)
        search_cache[postcode_clean] = results
    for result in results:
        if result.uprn is not None and int(result.uprn) == uprn:
            return _address_from_search(result)
    return None


def fill(df: pd.DataFrame, *, os_api_key: Optional[str]) -> list[dict[str, str]]:
    """Fill the DOMNA columns in place. Returns the unresolved rows."""
    for col in (FOUND_UPRN_COL, FOUND_ADDRESS_COL, SCORE_COL, SOURCE_COL):
        if col not in df.columns:
            df[col] = ""
    df[FOUND_UPRN_COL] = df[FOUND_UPRN_COL].astype("object")
    df[FOUND_ADDRESS_COL] = df[FOUND_ADDRESS_COL].astype("object")

    token = os.environ.get("OPEN_EPC_API_TOKEN")
    service = EpcClientService(auth_token=token) if token else None
    epc_cache: dict[str, pd.DataFrame] = {}
    os_cache: dict[str, pd.DataFrame] = {}
    search_cache: dict[str, list[EpcSearchResult]] = {}

    unresolved: list[dict[str, str]] = []
    resolved_uprn = resolved_addr = skipped = 0
    total = len(df)

    for n, idx in enumerate(df.index, start=1):
        ref = cell_str(df.at[idx, REF_COL])
        given_uprn = parse_uprn_cell(df.at[idx, UPRN_COL])
        address = cell_str(df.at[idx, ADDRESS_COL])
        postcode_raw = cell_str(df.at[idx, POSTCODE_COL])
        postcode_clean = clean_postcode(postcode_raw)

        # Already sorted (UPRN + address) or already filled by a prior run.
        if given_uprn is not None and address:
            skipped += 1
            continue
        if cell_str(df.at[idx, FOUND_UPRN_COL]) and cell_str(df.at[idx, FOUND_UPRN_COL]) != NOT_FOUND:
            skipped += 1
            continue

        def mark_not_found(reason: str) -> None:
            df.at[idx, FOUND_UPRN_COL] = NOT_FOUND if given_uprn is None else ""
            df.at[idx, FOUND_ADDRESS_COL] = NOT_FOUND
            df.at[idx, SOURCE_COL] = "not_found"
            unresolved.append(
                {
                    "Organisation Reference": ref,
                    "reason": reason,
                    "Address": address,
                    "Postcode": postcode_raw,
                }
            )

        # Case B — UPRN present, address missing: reverse-lookup the address.
        if given_uprn is not None and not address:
            found: Optional[str] = None
            if service is not None and postcode_clean:
                try:
                    found = reverse_address_from_uprn(
                        given_uprn, postcode_clean, service, search_cache
                    )
                except Exception as exc:
                    print(f"  reverse failed {ref} {given_uprn}: {exc}")
            if found:
                df.at[idx, FOUND_ADDRESS_COL] = found
                df.at[idx, SOURCE_COL] = "epc_reverse"
                resolved_addr += 1
            else:
                mark_not_found("no address for UPRN")
            continue

        # Case A — no UPRN, has address: resolve a UPRN.
        if given_uprn is None and address:
            if not postcode_clean:
                mark_not_found("no postcode")
                continue
            hit: Optional[Hit] = None
            if token:
                try:
                    hit = resolve_epc_relaxed(address, postcode_clean, epc_cache)
                except Exception as exc:
                    print(f"  EPC failed {ref} {postcode_clean}: {exc}")
            if hit is None and os_api_key:
                try:
                    hit = resolve_os_relaxed(address, postcode_clean, os_api_key, os_cache)
                except Exception as exc:
                    print(f"  OS failed {ref} {postcode_clean}: {exc}")
            if hit is not None:
                uprn, matched, score, source = hit
                df.at[idx, FOUND_UPRN_COL] = uprn
                df.at[idx, FOUND_ADDRESS_COL] = matched
                df.at[idx, SCORE_COL] = round(score, 4)
                df.at[idx, SOURCE_COL] = source
                resolved_uprn += 1
            else:
                mark_not_found("no UPRN match")
            if n % 100 == 0:
                print(
                    f"[{n}/{total}] resolved={resolved_uprn} not_found={len(unresolved)}"
                )
            continue

        # Case C — neither a UPRN nor an address.
        mark_not_found("no UPRN and no address")

    print(
        f"\nResolved {resolved_uprn} UPRNs, {resolved_addr} addresses; "
        f"{skipped} already sorted/done; {len(unresolved)} not found."
    )
    return unresolved


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("--unresolved", type=Path, default=_DEFAULT_UNRESOLVED)
    parser.add_argument("--limit", type=int, default=None, help="process first N rows")
    return parser.parse_args()


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

    df = pd.read_excel(args.inp, sheet_name=SHEET)
    if args.limit is not None:
        df = df.head(args.limit).copy()
    print(f"Loaded {len(df)} rows from {args.inp} [{SHEET}]")

    unresolved = fill(df, os_api_key=os_api_key)

    df.to_excel(args.out, sheet_name=SHEET, index=False)
    print(f"Wrote filled sheet -> {args.out}")
    if unresolved:
        pd.DataFrame(unresolved).to_csv(args.unresolved, index=False)
        print(f"Wrote {len(unresolved)} unresolved rows -> {args.unresolved}")
    return 0


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