app_simple_ui.py

"""
Subdivision Cluster / FIPS-Validation pipeline — Hugging Face Space (Gradio).

Faithful port of the Colab notebook
`LatestUndeterminedZipCityCensusTract70_Threshold_95_Clusters.ipynb`.

Colab-specific I/O replaced:
  * Drive input folder      -> file upload (multiple .xlsx/.xls)
  * google.colab.files      -> Gradio file outputs (written to a temp dir)
  * Drive safety-list path  -> get_safety_list_path() (repo file OR private dataset)

All processing logic (FIPS validation, RawFile snapshot, RapidFuzz clustering,
wide tables, majority-vote replacement, column reorder, OnlyBlanks warnings,
chunked CSV export, self-test) matches the notebook.
"""

import os
import re
import gc
import math
import time
import tempfile
import zipfile
from pathlib import Path
from collections import defaultdict

import numpy as np
import pandas as pd
import gradio as gr
from rapidfuzz import fuzz


# =====================================================================
# REFERENCE FILE LOADER  (the Subdivision Program Safety List)
# ---------------------------------------------------------------------
# 1. A file committed in the Space repo (override name via SAFETY_LIST_FILE).
#    Use ONLY if the Space is PRIVATE.
# 2. A private HF *dataset* repo, downloaded with the HF_TOKEN secret. Set:
#       SAFETY_LIST_REPO        e.g. "your-username/subdivision-safety-list"
#       SAFETY_LIST_REPO_FILE   filename inside that dataset (optional)
#       HF_TOKEN                read token, stored as a Space Secret
# =====================================================================
DEFAULT_SAFETY_LIST_NAME = "Subdivision_Program_Safety_List.xlsx"
EXCEL_ROW_LIMIT = 1_048_575   # matches the notebook

# Fixed pipeline settings (notebook values). Not exposed in the UI on purpose:
# users should not be able to deviate from the notebook's behavior.
STRICT_BLANK_CITIES = True    # blank cities that are not in the safety list
CLUSTER_THRESHOLD = 95        # RapidFuzz similarity for LegalSubDivision clusters
REPLACE_THRESHOLD = 70        # majority-vote % required to fill a blank
FILTER_DOMINANT_COUNTY = False  # process ALL counties (multi-county / full-state files)

# Master bundle format:
#   "excel" -> Master_Dataset.xlsx, every result on its own sheet (full output
#              split across Output_..._1/_2/... if it exceeds Excel's row cap)
#   "csv"   -> Master_Dataset.zip, every result bundled as a CSV (faster: no Excel
#              writing -- use this for very large / state-sized runs)
# The individual per-category files are CSV in both modes.
MASTER_OUTPUT = "excel"


def get_safety_list_path() -> str:
    local_name = os.environ.get("SAFETY_LIST_FILE", DEFAULT_SAFETY_LIST_NAME)
    if Path(local_name).exists():
        return local_name

    repo_id = os.environ.get("SAFETY_LIST_REPO")
    token = os.environ.get("HF_TOKEN")
    if repo_id and token:
        from huggingface_hub import hf_hub_download
        filename = os.environ.get("SAFETY_LIST_REPO_FILE", local_name)
        return hf_hub_download(
            repo_id=repo_id, filename=filename, repo_type="dataset", token=token
        )

    raise FileNotFoundError(
        "Safety list not found.\n"
        f"  - Commit '{local_name}' into this Space (keep the Space private), OR\n"
        "  - Set SAFETY_LIST_REPO + HF_TOKEN to pull it from a private dataset."
    )


_SAFETY_CACHE: dict | None = None


def build_whitelist_from_path(path):
    """Parse a safety-list file (xlsx or csv) into {pairs, cities, rows}."""
    if str(path).lower().endswith(".csv"):
        safety_df = pd.read_csv(path)
    else:
        safety_df = pd.read_excel(path)
    acceptable_cols = [c for c in safety_df.columns if str(c).startswith("Acceptable FIPS")]

    valid_pairs, valid_cities = set(), set()
    for _, r in safety_df.iterrows():
        c = _norm_city(r["City Name"])
        if c is None:
            continue
        valid_cities.add(c)
        for col in acceptable_cols:
            f = _norm_fips(r[col])
            if f:
                valid_pairs.add(c + "||" + f)

    return {"pairs": valid_pairs, "cities": valid_cities, "rows": len(safety_df)}


def load_whitelist():
    """Server-side safety list (repo file or private dataset). Cached per process."""
    global _SAFETY_CACHE
    if _SAFETY_CACHE is None:
        _SAFETY_CACHE = build_whitelist_from_path(get_safety_list_path())
    return _SAFETY_CACHE


# =====================================================================
# NORMALIZATION HELPERS  (identical to the notebook)
# =====================================================================
PLACEHOLDERS = {"UNDETERMINED CITY", "UNDETERMINED ZIP", "NULL", "NAN", "NONE", ""}


def _norm_city(x):
    if pd.isna(x):
        return None
    return str(x).strip().upper()


def _norm_fips(x):
    """5-digit FIPS string, or None. Handles 28071, 28071.0, '01067', '28071 - LAFAYETTE'."""
    if pd.isna(x):
        return None
    m = re.search(r"\d+", str(x))
    return str(int(m.group(0))).zfill(5) if m else None


# ---------------------------------------------------------------------
# FIPS-COUNTY COLUMN DETECTION (by CONTENT, not header name)
# ---------------------------------------------------------------------
# The source header has changed before ("** REC FIPS-COUNTY" ->
# "** REC FIPS-COUNTY NAME"), so we identify the column by what's IN it:
# values shaped like "05045-FAULKNER" (FIPS digits, a dash, a county name).
# Header names are only used as a last-resort fallback.
FIPS_COUNTY_PATTERN = r"^\s*\d{1,5}\s*-\s*[A-Za-z]"
FALLBACK_FIPS_HEADERS = ["** REC FIPS-COUNTY NAME", "** REC FIPS-COUNTY", "County"]


def detect_fips_county_column(df, sample_size=1000, min_score=0.5):
    """Return (column_name, reason). Picks the column whose values best match
    the 'FIPS-COUNTYNAME' shape; falls back to known headers, then to any
    header mentioning FIPS."""
    best_col, best_score = None, 0.0
    for col in df.columns:
        s = df[col].dropna()
        if s.empty:
            continue
        sample = s.head(sample_size).astype(str)
        try:
            score = float(sample.str.match(FIPS_COUNTY_PATTERN).mean())
        except Exception:
            continue
        if score > best_score:
            best_col, best_score = col, score

    if best_col is not None and best_score >= min_score:
        return best_col, f"content match {best_score:.0%} of sampled rows"

    for name in FALLBACK_FIPS_HEADERS:
        if name in df.columns:
            return name, "fallback: known header name"

    for col in df.columns:
        if "FIPS" in str(col).upper():
            return col, "fallback: header mentions FIPS"

    return None, "not found"


# =====================================================================
# SELF-TEST  (startup health check — same logic as the notebook's cell)
# =====================================================================
def run_self_test() -> str:
    path = get_safety_list_path()
    st = pd.read_excel(path)
    fc = [c for c in st.columns if str(c).startswith("Acceptable FIPS")]
    pairs = set()
    for _, r in st.iterrows():
        c = _norm_city(r["City Name"])
        if c is None:
            continue
        for col in fc:
            f = _norm_fips(r[col])
            if f:
                pairs.add(c + "||" + f)

    good_city = st.iloc[0]["City Name"]
    good_fips = st.iloc[0]["Acceptable FIPS 1"]
    test = pd.DataFrame({
        "DDPPropertyCityFinal": [good_city, good_city],
        "** REC FIPS-COUNTY": [good_fips, 99999],   # row 2 = wrong FIPS on purpose
    })
    cn = test["DDPPropertyCityFinal"].map(_norm_city)
    fn = test["** REC FIPS-COUNTY"].map(_norm_fips)
    key = cn.fillna("") + "||" + fn.fillna("")
    ok = key.isin(pairs)
    assert bool(ok.iloc[0]) and not bool(ok.iloc[1]), \
        "SELF-TEST FAILED -- matching logic is not working as expected."
    return f"SELF-TEST PASSED: known-good ({good_city}) kept, wrong FIPS flagged."


# =====================================================================
# STEP 1 — FIPS SAFETY-LIST VALIDATION
# =====================================================================
def fips_validate(df: pd.DataFrame, blank_cities_not_in_list: bool, log, wl=None,
                  fips_col=None):
    if wl is None:
        wl = load_whitelist()
    valid_pairs, valid_cities = wl["pairs"], wl["cities"]

    if fips_col is None:
        fips_col, _reason = detect_fips_county_column(df)
    if fips_col is None:
        raise ValueError(
            "Could not find a FIPS-county column. Expected values like "
            "'05045-FAULKNER' in one of the columns."
        )
    FIPS_COL = fips_col
    CITY_COL, ZIP_COL, TRACT_COL = (
        "DDPPropertyCityFinal", "DDPPropertyZIPFinal", "DDPRevisedCensusTract",
    )

    city_norm = df[CITY_COL].map(_norm_city)
    fips_norm = df[FIPS_COL].map(_norm_fips)
    row_key = city_norm.fillna("") + "||" + fips_norm.fillna("")

    is_real_city = df[CITY_COL].notna() & (~city_norm.fillna("").isin(PLACEHOLDERS))
    city_in_list = city_norm.isin(valid_cities)
    pair_ok = row_key.isin(valid_pairs)

    if blank_cities_not_in_list:
        fails_mask = is_real_city & (~pair_ok)
    else:
        fails_mask = is_real_city & city_in_list & (~pair_ok)

    n_real = int(is_real_city.sum())
    n_inlist = int((is_real_city & city_in_list).sum())
    n_valid = int((is_real_city & pair_ok).sum())
    n_fail = int(fails_mask.sum())
    log.append(
        f"FIPS validation — list: {wl['rows']:,} rows | {len(valid_cities):,} cities | "
        f"{len(valid_pairs):,} pairs"
    )
    log.append(
        f"  real:{n_real:,} in-list:{n_inlist:,} valid-kept:{n_valid:,} blanked:{n_fail:,}"
    )
    if n_real > 0 and n_inlist == 0:
        log.append("  !! 0 cities matched the list — likely a format/spelling mismatch.")

    # Audit built BEFORE overwrite (insert position matches the notebook).
    key_col = "RecordKey" if "RecordKey" in df.columns else None
    base_cols = [c for c in [key_col, CITY_COL, FIPS_COL, ZIP_COL, TRACT_COL] if c]
    audit = df.loc[fails_mask, base_cols].copy()
    audit.insert(
        len(base_cols) if key_col is None else 2,
        "NormalizedFIPS", fips_norm[fails_mask].values,
    )
    audit["Reason"] = np.where(
        city_in_list[fails_mask].values, "FIPS does not match city", "City not in safety list"
    )
    audit = audit.rename(
        columns={CITY_COL: "OLD_City", ZIP_COL: "OLD_ZIP", TRACT_COL: "OLD_CensusTract"}
    )

    # Recent pandas reads all-numeric Excel columns as float64 and refuses to
    # accept a string into them; Colab's older pandas silently upcast to object.
    # Cast first so the placeholder write reproduces the notebook's result.
    df[CITY_COL] = df[CITY_COL].astype(object)
    df[ZIP_COL] = df[ZIP_COL].astype(object)
    df.loc[fails_mask, CITY_COL] = "UNDETERMINED CITY"
    df.loc[fails_mask, ZIP_COL] = "UNDETERMINED ZIP"
    df.loc[fails_mask, TRACT_COL] = np.nan
    return df, audit


# =====================================================================
# STEP 2 — RAPIDFUZZ CLUSTERING  (identical to the notebook)
# =====================================================================
def group_similar_strings_optimized(strings, threshold=95, prefix_len=3, max_len_diff=3):
    cleaned, seen = [], set()
    for s in strings:
        if pd.isna(s):
            continue
        s = str(s).strip()
        if s and s not in seen:
            cleaned.append(s)
            seen.add(s)

    # Canonical order so the SAME set of names always clusters the same way,
    # regardless of the order rows arrived in (single-county vs whole-state,
    # app vs app_simple_ui, or any file-upload order). The greedy grouping below
    # is order-sensitive, so we sort first to make it deterministic.
    cleaned.sort(key=lambda x: (x.upper(), x))

    visited, groups = set(), []
    buckets = defaultdict(list)
    for s in cleaned:
        buckets[s.upper()[:prefix_len]].append(s)

    for key in sorted(buckets.keys()):
        bucket = buckets[key]
        for i, base in enumerate(bucket):
            if base in visited:
                continue
            group = [base]
            visited.add(base)
            base_upper = base.upper()
            base_len = len(base_upper)
            for j in range(i + 1, len(bucket)):
                compare = bucket[j]
                if compare in visited:
                    continue
                comp_upper = compare.upper()
                if abs(base_len - len(comp_upper)) > max_len_diff:
                    continue
                if fuzz.ratio(base_upper, comp_upper) >= threshold:
                    group.append(compare)
                    visited.add(compare)
            groups.append(group)
    return groups


def build_clusters(city_df, zip_df, census_tracts_df, cluster_threshold):
    all_legal = pd.concat(
        [
            city_df[["County", "LegalSubDivision"]],
            zip_df[["County", "LegalSubDivision"]],
            census_tracts_df[["County", "LegalSubDivision"]],
        ],
        ignore_index=True,
    ).dropna(subset=["County", "LegalSubDivision"])
    all_legal["County"] = all_legal["County"].astype(str).str.strip()
    all_legal["LegalSubDivision"] = all_legal["LegalSubDivision"].astype(str).str.strip()
    all_legal = all_legal[(all_legal["County"] != "") & (all_legal["LegalSubDivision"] != "")]
    all_legal = all_legal[["County", "LegalSubDivision"]].drop_duplicates()

    groups = []   # (county, [members]) so each cluster stays tied to its county
    for county in sorted(all_legal["County"].unique()):
        strings = all_legal[all_legal["County"] == county]["LegalSubDivision"].tolist()
        for grp in group_similar_strings_optimized(
                strings, threshold=cluster_threshold, prefix_len=3, max_len_diff=3):
            groups.append((county, grp))

    rows = []
    for cid, (county, group) in enumerate(groups, start=1):
        for sub in group:
            rows.append({"Cluster_ID": cid, "County": county, "LegalSubDivision": sub})
    return pd.DataFrame(rows, columns=["Cluster_ID", "County", "LegalSubDivision"])


# =====================================================================
# STEP 3 — WIDE TABLES + MAJORITY-VOTE "REPLACE WITH"
# =====================================================================
def build_cluster_wide_table(source_df, cluster_df, value_col, prefix):
    # Map by (County, LegalSubDivision) -- NOT LegalSubDivision alone. The same
    # subdivision name occurs in many counties, each with its own cluster; keying
    # on the name only would collapse them and pool values across counties.
    lk = cluster_df[["County", "LegalSubDivision", "Cluster_ID"]].drop_duplicates(
        ["County", "LegalSubDivision"])
    cluster_map = {(c, s): cid for c, s, cid in
                   zip(lk["County"].astype(str).str.strip(),
                       lk["LegalSubDivision"].astype(str).str.strip(), lk["Cluster_ID"])}

    w = source_df[["County", "LegalSubDivision", value_col]].copy()
    w["County"] = w["County"].astype(str).str.strip()
    w["LegalSubDivision"] = w["LegalSubDivision"].astype(str).str.strip()
    w[value_col] = w[value_col].replace("", np.nan)
    w["Cluster_ID"] = [cluster_map.get((c, s))
                       for c, s in zip(w["County"], w["LegalSubDivision"])]
    w = w.dropna(subset=["Cluster_ID", value_col])
    if w.empty:
        return pd.DataFrame(columns=["Cluster_ID"])
    w["Cluster_ID"] = w["Cluster_ID"].astype(int)
    w[value_col] = w[value_col].astype(str).str.strip()
    w = w[w[value_col] != ""]

    counts = w.groupby(["Cluster_ID", value_col], observed=True).size().unstack(fill_value=0)
    totals = counts.sum(axis=1)
    pct = counts.div(totals.replace(0, np.nan), axis=0).mul(100).fillna(0)
    counts.columns = [f"{prefix}{c}_Count" for c in counts.columns]
    pct.columns = [f"{prefix}{c}_Percentage" for c in pct.columns]
    wide = pd.concat([counts, pct], axis=1).reset_index()

    # Match the notebook: Cluster_ID first, all other columns sorted.
    other = sorted([c for c in wide.columns if c != "Cluster_ID"])
    return wide[["Cluster_ID"] + other]


def build_replace_with(wide_df, prefix_to_remove, threshold):
    pct_cols = [c for c in wide_df.columns if c.endswith("_Percentage")]
    if not pct_cols:
        res = wide_df[["Cluster_ID"]].copy()
        res["Replace with"] = None
        return res.drop_duplicates(subset=["Cluster_ID"])
    labels = np.array(
        [c.replace(prefix_to_remove, "").replace("_Percentage", "") for c in pct_cols],
        dtype=object,
    )
    values = wide_df[pct_cols].to_numpy(copy=False)
    out = []
    for row in values:
        matched = labels[row >= threshold]
        out.append(", ".join(matched) if len(matched) else None)
    res = wide_df[["Cluster_ID"]].copy()
    res["Replace with"] = out
    return res.drop_duplicates(subset=["Cluster_ID"])


def clean_census_tract_suffix(value):
    if pd.isna(value):
        return np.nan
    try:
        f = float(value)
        return str(int(f)) if f == int(f) else str(f)
    except ValueError:
        return str(value)


def clean_zip(value):
    """Normalize ZIPs read from Excel as numbers: drop the float '.0' artifact and
    zero-pad to 5 digits (recovering leading zeros lost on numeric import, e.g.
    7221.0 -> 07221). Handles comma-joined majority-vote results. ZIP+4 left as-is."""
    if pd.isna(value):
        return np.nan
    out = []
    for part in str(value).split(","):
        p = part.strip()
        if p.endswith(".0"):
            p = p[:-2]
        digits = re.sub(r"\D", "", p)
        if not digits:
            continue
        out.append(digits.zfill(5) if len(digits) <= 5 else digits)
    return ", ".join(out) if out else np.nan


_BOOL_MAP = {
    True: "TRUE", False: "FALSE", 1: "TRUE", 0: "FALSE",
    "1": "TRUE", "0": "FALSE", "true": "TRUE", "false": "FALSE",
    "t": "TRUE", "f": "FALSE", "yes": "TRUE", "no": "FALSE",
}


def boolify(series):
    """Render a 0/1 (or True/False) column as literal 'TRUE'/'FALSE' text.
    Recognized values are mapped; anything unrecognized is left as-is; NaN kept."""
    def conv(v):
        if pd.isna(v):
            return v
        if v in _BOOL_MAP:
            return _BOOL_MAP[v]
        return _BOOL_MAP.get(str(v).strip().lower(), v)
    return series.map(conv)


# =====================================================================
# CHUNKED CSV WRITER  (faithful port of save_csv_in_parts)
# =====================================================================
def save_csv_in_parts(df, base_filename, out_dir, row_limit=1_000_000):
    base_name = base_filename.replace(".csv", "")
    created = []
    total = len(df)
    num_parts = math.ceil(total / row_limit)
    for i in range(num_parts):
        start = i * row_limit
        end = min((i + 1) * row_limit, total)
        part = f"{base_name}_part_{i + 1}.csv" if num_parts > 1 else base_filename
        p = out_dir / part
        df.iloc[start:end].to_csv(p, index=False)
        created.append(str(p))
        gc.collect()
    return created


def write_master_workbook(path, sheets_in_order, row_limit=1_048_575):
    """Write a multi-sheet .xlsx from [(sheet_name, df), ...]. Any sheet larger than
    Excel's 1,048,576-row cap is split into <name>_1, <name>_2, ... so nothing is
    dropped. Uses xlsxwriter in constant_memory mode, driven row-by-row (low RAM --
    important for the full-output sheet on state-sized files). We write the rows
    ourselves because the pandas + constant_memory path corrupts rows that contain
    NaN cells. Falls back to openpyxl if xlsxwriter isn't installed."""
    try:
        import xlsxwriter
    except Exception:
        xlsxwriter = None

    if xlsxwriter is not None:
        wb = xlsxwriter.Workbook(str(path), {"constant_memory": True, "nan_inf_to_errors": True})

        def _write(name, dfx):
            ws = wb.add_worksheet(str(name)[:31])
            ws.write_row(0, 0, [str(c) for c in dfx.columns])
            clean = dfx.where(pd.notna(dfx), None)   # NaN -> blank cell
            for r, rowvals in enumerate(clean.itertuples(index=False, name=None), start=1):
                ws.write_row(r, 0, rowvals)

        for name, dfx in sheets_in_order:
            if dfx is None:
                continue
            n = len(dfx)
            if n <= row_limit:
                _write(name, dfx)
            else:
                for i in range(math.ceil(n / row_limit)):
                    _write(f"{name}_{i + 1}", dfx.iloc[i * row_limit:(i + 1) * row_limit])
                    gc.collect()
        wb.close()
    else:
        with pd.ExcelWriter(path, engine="openpyxl") as writer:
            for name, dfx in sheets_in_order:
                if dfx is None:
                    continue
                n = len(dfx)
                if n <= row_limit:
                    dfx.to_excel(writer, sheet_name=str(name)[:31], index=False)
                else:
                    for i in range(math.ceil(n / row_limit)):
                        dfx.iloc[i * row_limit:(i + 1) * row_limit].to_excel(
                            writer, sheet_name=f"{name}_{i + 1}"[:31], index=False)
    return str(path)


# =====================================================================
# MAIN PIPELINE
# =====================================================================
def run_pipeline(uploaded_files, reference_file, blank_cities_not_in_list,
                 filter_dominant_county, cluster_threshold, replace_threshold):
    t_start = time.time()
    log = []
    if not uploaded_files:
        return "Please upload at least one Excel file.", []

    out_dir = Path(tempfile.mkdtemp(prefix="subdiv_"))
    outputs = []

    # ---- Resolve the reference (safety) list -------------------------
    # Default: server-side list (repo file or private dataset). If the user
    # uploads an override on the sidebar, use it for THIS run only.
    try:
        if reference_file is not None:
            wl = build_whitelist_from_path(reference_file.name)
            log.append(f"Using UPLOADED reference list: {Path(reference_file.name).name}")
        else:
            wl = load_whitelist()
            log.append("Using server-side safety list.")
    except FileNotFoundError as e:
        return "\n".join(log) + f"\n\nREFERENCE FILE ERROR:\n{e}", []
    except Exception as e:
        return "\n".join(log) + f"\n\nCould not parse reference list: {e}", []

    # ---- Load & merge uploads (Excel or CSV) -------------------------
    dfs = []
    for f in uploaded_files:
        name = f.name
        try:
            if str(name).lower().endswith(".csv"):
                dfs.append(pd.read_csv(name))
            else:
                dfs.append(pd.read_excel(name))
            log.append(f"Loaded {Path(name).name}")
        except Exception as e:
            log.append(f"ERROR loading {Path(name).name}: {e}")
    if not dfs:
        return "\n".join(log) + "\nNo files could be read.", []
    df = pd.concat(dfs, ignore_index=True)
    log.append(f"Merged shape: {df.shape}")

    # ---- Locate the FIPS-county column by CONTENT --------------------
    fips_col, fips_reason = detect_fips_county_column(df)
    if fips_col is None:
        return ("\n".join(log) +
                "\n\nCOLUMN ERROR: could not find a FIPS-county column. Expected "
                "values like '05045-FAULKNER' in one of the columns."), []
    log.append(f"FIPS-county column: '{fips_col}' ({fips_reason})")

    # ---- FIPS validation (uses the reference list) -------------------
    try:
        df, audit = fips_validate(df, blank_cities_not_in_list, log, wl=wl, fips_col=fips_col)
    except FileNotFoundError as e:
        return "\n".join(log) + f"\n\nREFERENCE FILE ERROR:\n{e}", []

    # ---- Verification files (audit + RawFile_FIPSValidated), CSV -----
    if len(audit) > 0:
        ap = out_dir / "FIPS_Validation_Audit_ChangedRecords.csv"
        audit.to_csv(ap, index=False)
        outputs.append(str(ap))

    rp = out_dir / "RawFile_FIPSValidated.csv"
    df.to_csv(rp, index=False)
    outputs.append(str(rp))

    # ---- UNDETERMINED/NULL -> NaN ------------------------------------
    for col in ["DDPPropertyCityFinal", "DDPPropertyZIPFinal",
                "DDPRevisedCensusTract", "LegalSubDivision"]:
        if col in df.columns:
            df[col] = df[col].replace(["UNDETERMINED CITY", "UNDETERMINED ZIP", "NULL"], np.nan)

    # ---- Rename detected FIPS col -> County --------------------------
    if fips_col != "County":
        if "County" in df.columns:
            df = df.drop(columns=["County"])
        df = df.rename(columns={fips_col: "County"})

    # ---- Filter to single most-dominant county -----------------------
    # DISABLED: we now process big multi-county files (up to a full state) in
    # one run, so we keep every county. Clustering already runs per-county.
    # To re-enable single-county mode, uncomment the block below.
    # if filter_dominant_county:
    #     dominant = df["County"].value_counts().idxmax()
    #     df = df[df["County"] == dominant].copy()
    #     log.append(f"Filtered to dominant county: {dominant} -> {df.shape}")

    # ---- Source frames -----------------------------------------------
    city_df = df.loc[df["DDPPropertyCityFinal"].notna()].copy()
    zip_df = df.loc[df["DDPPropertyZIPFinal"].notna()].copy()
    census_tracts_df = df.loc[df["DDPRevisedCensusTract"].notna()].copy()
    original = df.loc[df["LegalSubDivision"].notna()].copy()

    # ---- Cluster ------------------------------------------------------
    cluster_df = build_clusters(city_df, zip_df, census_tracts_df, cluster_threshold)
    log.append(f"Clusters: {cluster_df['Cluster_ID'].nunique():,}")

    # ---- Replace-with (majority vote) --------------------------------
    zip_w = build_replace_with(
        build_cluster_wide_table(zip_df, cluster_df, "DDPPropertyZIPFinal", "ZIP_"),
        "ZIP_", replace_threshold).rename(columns={"Replace with": "Replace with Zip Codes"})
    city_w = build_replace_with(
        build_cluster_wide_table(city_df, cluster_df, "DDPPropertyCityFinal", "City_"),
        "City_", replace_threshold).rename(columns={"Replace with": "Replace with Cities"})
    cen_w = build_replace_with(
        build_cluster_wide_table(census_tracts_df, cluster_df, "DDPRevisedCensusTract", "CensusTract_"),
        "CensusTract_", replace_threshold).rename(columns={"Replace with": "Replace with Census Tracts"})

    # ---- Map back onto original --------------------------------------
    # Same fix: assign each row its county's cluster, keyed on (County, name).
    lk = cluster_df[["County", "LegalSubDivision", "Cluster_ID"]].drop_duplicates(
        ["County", "LegalSubDivision"])
    cmap = {(c, s): cid for c, s, cid in
            zip(lk["County"].astype(str).str.strip(),
                lk["LegalSubDivision"].astype(str).str.strip(), lk["Cluster_ID"])}
    original_final = original.copy()
    _oc = original_final["County"].astype(str).str.strip()
    _os = original_final["LegalSubDivision"].astype(str).str.strip()
    original_final["Cluster_ID"] = [cmap.get((c, s)) for c, s in zip(_oc, _os)]
    original_final["Replace with Zip Codes"] = original_final["Cluster_ID"].map(
        dict(zip(zip_w["Cluster_ID"], zip_w["Replace with Zip Codes"])))
    original_final["Replace with Cities"] = original_final["Cluster_ID"].map(
        dict(zip(city_w["Cluster_ID"], city_w["Replace with Cities"])))
    original_final["Replace with Census Tracts"] = original_final["Cluster_ID"].map(
        dict(zip(cen_w["Cluster_ID"], cen_w["Replace with Census Tracts"])))

    # Normalize ZIPs (drop float '.0', restore 5-digit zero-padding).
    original_final["Replace with Zip Codes"] = (
        original_final["Replace with Zip Codes"].apply(clean_zip))

    # ---- Reorder columns (desired_leading_columns) -------------------
    desired_leading = [
        "RecordKey", "ComesFromMtg2", "IsFlagged", "Cluster_ID", "LegalSubDivision",
        "County", "Replace with Zip Codes", "Replace with Cities", "Replace with Census Tracts",
        "DDPPropertyZIPFinal", "DDPPropertyCityFinal", "DDPRevisedCensusTract",
        "LotBlockSubdivisionTract", "DDPRevisedLatitude", "DDPRevisedLongitude",
    ]
    existing = original_final.columns.tolist()
    lead = [c for c in desired_leading if c in existing]
    other = [c for c in existing if c not in lead]
    original_final = original_final[lead + other]

    original_final_output = original_final.copy()   # full output snapshot

    # ---- Split into final cities / zips / census tracts --------------
    keep = ["RecordKey", "ComesFromMtg2", "IsFlagged",
            "Replace with Zip Codes", "Replace with Cities", "Replace with Census Tracts",
            "DDPPropertyZIPFinal", "DDPPropertyCityFinal", "DDPRevisedCensusTract"]
    keep = [c for c in keep if c in original_final.columns]
    of = original_final[keep].rename(columns={
        "DDPPropertyZIPFinal": "DDPPropertyZIPFinalToRemove",
        "DDPPropertyCityFinal": "DDPPropertyCityFinalToRemove",
        "DDPRevisedCensusTract": "DDPRevisedCensusTractToRemove",
        "Replace with Zip Codes": "DDPPropertyZIPFinal",
        "Replace with Cities": "DDPPropertyCityFinal",
        "Replace with Census Tracts": "DDPRevisedCensusTract",
    })
    of["DDPRevisedCensusTract"] = of["DDPRevisedCensusTract"].apply(clean_census_tract_suffix)

    idcols = [c for c in ["RecordKey", "ComesFromMtg2", "IsFlagged"] if c in of.columns]
    final_cities = of[idcols + ["DDPPropertyCityFinalToRemove", "DDPPropertyCityFinal"]].copy()
    final_zips = of[idcols + ["DDPPropertyZIPFinalToRemove", "DDPPropertyZIPFinal"]].copy()
    final_ct = of[idcols + ["DDPRevisedCensusTractToRemove", "DDPRevisedCensusTract"]].copy()

    # keep only rows whose ORIGINAL value was blank (the ones we filled)
    final_cities = final_cities[final_cities["DDPPropertyCityFinalToRemove"].isna()].copy()
    final_zips = final_zips[final_zips["DDPPropertyZIPFinalToRemove"].isna()].copy()
    final_ct = final_ct[final_ct["DDPRevisedCensusTractToRemove"].isna()].copy()

    final_cities = final_cities.drop(columns=["DDPPropertyCityFinalToRemove"])
    final_zips = final_zips.drop(columns=["DDPPropertyZIPFinalToRemove"])
    final_ct = final_ct.drop(columns=["DDPRevisedCensusTractToRemove"])

    # ---- OnlyBlanks warnings (column entirely null) ------------------
    def only_blanks_warning(d, col, prefix):
        if col in d.columns and d[col].isnull().all():
            p = out_dir / f"OnlyBlanks_{prefix}.csv"
            d.to_csv(p, index=False)
            outputs.append(str(p))
            log.append(f"WARNING: {prefix} has only blanks in '{col}' -> {p.name}")

    only_blanks_warning(final_cities, "DDPPropertyCityFinal", "final_cities")
    only_blanks_warning(final_zips, "DDPPropertyZIPFinal", "final_zips")
    only_blanks_warning(final_ct, "DDPRevisedCensusTract", "final_census_tracts")

    # ---- Drop NaNs so only program assignments remain ----------------
    final_cities = final_cities[final_cities["DDPPropertyCityFinal"].notna()].copy()
    final_zips = final_zips[final_zips["DDPPropertyZIPFinal"].notna()].copy()
    final_ct = final_ct[final_ct["DDPRevisedCensusTract"].notna()].copy()

    # ---- Boolean columns as TRUE/FALSE (platform can't read 0/1) -----
    for d in (final_cities, final_zips, final_ct):
        for bc in ("ComesFromMtg2", "IsFlagged"):
            if bc in d.columns:
                d[bc] = boolify(d[bc])

    # ---- Write the per-category + cluster outputs (CSV) --------------
    for d, name in [(final_cities, "final_cities.csv"),
                    (final_zips, "final_zips.csv"),
                    (final_ct, "final_census_tracts.csv"),
                    (cluster_df, "cluster_df.csv")]:
        p = out_dir / name
        d.to_csv(p, index=False)
        outputs.append(str(p))

    # ---- Full output CSV (chunked at 1,000,000 rows) -----------------
    # NB: no '%' in the filename -- '%' is the URL-encoding character and makes
    # the browser download fail with "unknown server error".
    csv_name = f"OutputFile_{replace_threshold}Threshold_{cluster_threshold}ClustersRapidFuzz.csv"
    outputs.extend(save_csv_in_parts(original_final_output, csv_name, out_dir, row_limit=1_000_000))

    # ---- Summary sheet: ONE ROW PER COUNTY ---------------------------
    # Totals (records / subdivisions) come from the FULL post-validation frame
    # (df -- every row, incl. those with no LegalSubDivision). The blank/fill/
    # cluster metrics come from the clustered subset (output_df), which is the
    # only set the pipeline can actually fill. Other output files are unchanged.
    output_df = original_final_output     # rows WITH a LegalSubDivision + fills
    full_df = df                          # ALL post-validation rows, has County
    city_df, zips_df, census_df = final_cities, final_zips, final_ct

    def _fips_name(label):
        raw = str(label)
        parts = raw.split("-")
        fips = parts[0].strip()
        name = parts[1].strip().title() if len(parts) > 1 else raw.strip().title()
        return fips, name

    RW_CITY, RW_ZIP, RW_TRACT = ("Replace with Cities", "Replace with Zip Codes",
                                 "Replace with Census Tracts")
    ORIG_CITY, ORIG_ZIP, ORIG_TRACT = ("DDPPropertyCityFinal", "DDPPropertyZIPFinal",
                                       "DDPRevisedCensusTract")

    def _blank(g, col):
        return g[col].isna() if col in g.columns else pd.Series(False, index=g.index)

    def _fill(g, col):
        return g[col].notna() if col in g.columns else pd.Series(False, index=g.index)

    def _out_subset(label):
        if "County" not in output_df.columns:
            return output_df.iloc[0:0]
        m = output_df["County"].isna() if pd.isna(label) else (output_df["County"] == label)
        return output_df[m]

    summary_rows = []
    if "County" in full_df.columns and len(full_df):
        for county_label, gfull in full_df.groupby("County", dropna=False):
            fips, name = _fips_name(county_label)
            g = _out_subset(county_label)   # clustered rows for this county
            city_blank, zip_blank, tract_blank = (_blank(g, ORIG_CITY),
                                                  _blank(g, ORIG_ZIP), _blank(g, ORIG_TRACT))
            city_fill, zip_fill, tract_fill = (_fill(g, RW_CITY),
                                               _fill(g, RW_ZIP), _fill(g, RW_TRACT))
            summary_rows.append({
                "County": name,
                "FIPS": fips,
                # true totals from the full frame
                "Total Records in County": len(gfull),
                # how many of those records carry a LegalSubDivision (non-unique)
                "Total Subdivisions": int(gfull["LegalSubDivision"].count())
                                      if "LegalSubDivision" in gfull.columns else 0,
                "Total Clusters Made": int(g["Cluster_ID"].nunique())
                                       if "Cluster_ID" in g.columns else 0,
                "Total Blank Cities": int(city_blank.sum()),
                "Total Cities Assignment": int((city_blank & city_fill).sum()),
                "Total Blank Zips": int(zip_blank.sum()),
                "Total Zips Assignment": int((zip_blank & zip_fill).sum()),
                "Total Blank Census": int(tract_blank.sum()),
                "Total Census Assignment": int((tract_blank & tract_fill).sum()),
            })

    summary_cols = ["County", "FIPS", "Total Records in County", "Total Subdivisions",
                    "Total Clusters Made", "Total Blank Cities", "Total Cities Assignment",
                    "Total Blank Zips", "Total Zips Assignment", "Total Blank Census",
                    "Total Census Assignment"]
    summary_df = pd.DataFrame(summary_rows, columns=summary_cols)
    log.append(f"Summary: {len(summary_df)} county row(s).")

    summary_path = out_dir / "Summary.csv"
    summary_df.to_csv(summary_path, index=False)
    outputs.append(str(summary_path))

    # ---- Master bundle: Excel workbook OR zip of CSVs (see MASTER_OUTPUT) -
    if MASTER_OUTPUT == "csv":
        # Bundle the already-written CSVs into one zip. No Excel writing -> fast.
        master_path = out_dir / "Master_Dataset.zip"
        member_names = ["FIPS_Validation_Audit_ChangedRecords.csv", "cluster_df.csv",
                        "final_cities.csv", "final_zips.csv", "final_census_tracts.csv",
                        "Summary.csv"]
        members = [out_dir / n for n in member_names if (out_dir / n).exists()]
        members += sorted(out_dir.glob("OutputFile_*.csv"))   # full output (chunked)
        with zipfile.ZipFile(master_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for p in members:
                zf.write(p, arcname=p.name)
    else:
        # Every result on its own sheet; the full output splits across
        # Output_70_95_RapidFuzz_1, _2, ... if it exceeds Excel's row cap.
        master_path = out_dir / "Master_Dataset.xlsx"
        write_master_workbook(master_path, [
            ("Output_70_95_RapidFuzz", output_df),
            ("ValidationAuditResults", audit),
            ("cluster_df", cluster_df),
            ("final_cities", final_cities),
            ("final_zips", final_zips),
            ("final_census_tracts", final_ct),
            ("Summary", summary_df),
        ])
    outputs.append(str(master_path))

    elapsed_min = (time.time() - t_start) / 60
    log.append(f"Done. {len(outputs)} files ready.")
    log.append(f"Total processing time: {elapsed_min:.1f} minutes.")
    gc.collect()
    return "\n".join(log), outputs


# =====================================================================
# STARTUP HEALTH CHECK
# =====================================================================
def _startup_health_check():
    if os.environ.get("SELF_TEST_ON_STARTUP", "1") != "1":
        return
    try:
        print("[startup] " + run_self_test())
    except FileNotFoundError as e:
        # Reference file not configured yet — let the UI load so the user can fix it.
        print(f"[startup] Self-test skipped — safety list unavailable:\n{e}")
    # AssertionError intentionally propagates: a failed self-test fails the boot.


_startup_health_check()


# =====================================================================
# GRADIO UI  (simple vertical layout — previous design)
# =====================================================================
with gr.Blocks() as demo:
    gr.Markdown(
        "## Subdivision Cluster + FIPS Validation\n"
        "Upload the Excel files to process. Records are validated against the "
        "Subdivision Program Safety List, then `LegalSubDivision` values are "
        "fuzzy-clustered and city/ZIP/census-tract are filled by majority vote.\n\n"
        "*Runs with the notebook's fixed settings.*"
    )
    files_in = gr.File(file_count="multiple", file_types=[".xlsx", ".xls", ".csv"],
                       label="Files to process (XLSX, XLS, or CSV)")
    run_btn = gr.Button("Run pipeline", variant="primary")
    status = gr.Textbox(label="Log", lines=14)
    files_out = gr.File(label="Output files", file_count="multiple")

    def _run_simple(files):
        # No reference upload in this layout -> always the server-side list.
        # All settings are fixed (see constants at top). No user-facing options.
        return run_pipeline(files, None, STRICT_BLANK_CITIES, FILTER_DOMINANT_COUNTY,
                            CLUSTER_THRESHOLD, REPLACE_THRESHOLD)

    run_btn.click(
        _run_simple,
        inputs=[files_in],
        outputs=[status, files_out],
    )

if __name__ == "__main__":
    demo.launch()