"""Sweep min-building-footprint thresholds across all setback distances and
write a single matrix to out/sweep_matrix.txt. Loads the data once; reuses
the buffered residential exclusion and per-parcel buildable polygons across
all thresholds. Does NOT touch the canonical eligible_*.csv files.
"""
from __future__ import annotations
import sys
import time
from pathlib import Path

import geopandas as gpd
from shapely import make_valid, unary_union
from shapely.geometry import MultiPolygon, Polygon

ROOT = Path(__file__).resolve().parent.parent
CACHE = ROOT / "cache"
OUT = ROOT / "out"

# Mirror the constants and helpers from analyze.py.
DC_INTERNAL_SETBACK_FT = 50
SETBACKS_FT = [200, 500, 1000]
BUILDING_THRESHOLDS = [5_000, 20_000, 50_000, 100_000, 250_000, 500_000]

RESIDENTIAL_DISTRICT_ZONECLASSES = {
    "R-1", "R-2", "R-2A", "R-2B", "R-2C", "R-3", "R-4", "R-5", "R-6",
    "RP-2", "RP-3", "RP-4", "RP-5", "RP-6",
    "RO-1", "RO-2", "A-1", "RM-3", "RM-4", "RM-5",
}
RESIDENTIAL_LANDUSE_CODES = {"11", "12", "13", "14"}
RESIDENTIAL_PLACETYPES = {
    "SF RESIDENCE", "APARTMENT", "APARTMENTS", "MOBILE HOME",
    "CONDOMINIUM", "TOWNHOUSE", "DUPLEX", "PUBLIC HOUSING", "DORMITORY",
}


def fix(g):
    if g is None or g.is_empty:
        return g
    return g if g.is_valid else make_valid(g)


def max_poly(g) -> float:
    if g is None or g.is_empty:
        return 0.0
    if isinstance(g, Polygon):
        return g.area
    if isinstance(g, MultiPolygon):
        return max((p.area for p in g.geoms), default=0.0)
    try:
        return max((max_poly(p) for p in g.geoms), default=0.0)
    except AttributeError:
        return 0.0


def landuse_res(s) -> bool:
    if s is None or (isinstance(s, float) and s != s):
        return False
    prefix = str(s).split(" - ", 1)[0]
    return bool({c.strip() for c in prefix.split(":")} & RESIDENTIAL_LANDUSE_CODES)


def main() -> int:
    t0 = time.time()
    print("loading layers...", flush=True)
    zoning = gpd.read_parquet(CACHE / "zoning.parquet")
    parcels = gpd.read_parquet(CACHE / "parcels.parquet")
    addresses = gpd.read_parquet(CACHE / "addresses.parquet")
    zoning["geometry"] = zoning.geometry.apply(fix)
    parcels["geometry"] = parcels.geometry.apply(fix)
    parcels = parcels[parcels.geometry.notna() & ~parcels.geometry.is_empty].copy()

    i2 = unary_union(zoning[zoning["ZONECLASS"] == "I-2"].geometry.values)
    cand_idx = sorted(parcels.sindex.query(i2, predicate="intersects"))
    cands = parcels.iloc[cand_idx].copy()
    cands["i2_geom"] = cands.geometry.intersection(i2).apply(fix)
    cands = cands[~cands["i2_geom"].is_empty].copy()
    print(f"  I-2 candidate parcels: {len(cands):,}")

    res_zones = zoning[zoning["ZONECLASS"].isin(RESIDENTIAL_DISTRICT_ZONECLASSES)]
    res_landuse = parcels[parcels["LANDUSE"].apply(landuse_res)]
    res_addr_mask = addresses["PLACETYPE"].astype(str).str.upper().isin(RESIDENTIAL_PLACETYPES)
    res_addrs = addresses[res_addr_mask]
    sj = gpd.sjoin(res_addrs[["geometry"]], parcels[["GISLINK", "geometry"]],
                   how="inner", predicate="within")
    res_addr_par = parcels[parcels["GISLINK"].isin(set(sj["GISLINK"].dropna().unique()))]
    res_union = unary_union(
        list(res_zones.geometry.values)
        + list(res_landuse.geometry.values)
        + list(res_addr_par.geometry.values)
    )
    print(f"  residential union built ({time.time()-t0:.1f}s)")

    # results[(setback_ft, building_sqft)] -> count
    results = {}
    strict_counts = {}
    for ft in SETBACKS_FT:
        t1 = time.time()
        excl = res_union.buffer(ft)
        sindex = gpd.GeoSeries([excl], crs=parcels.crs).sindex
        strict = 0
        siteable_per_parcel = []
        for row in cands.itertuples(index=False):
            parcel_g = row.geometry
            i2g = row.i2_geom
            hits = list(sindex.query(parcel_g, predicate="intersects"))
            if not hits:
                strict += 1
                eroded = fix(i2g.buffer(-DC_INTERNAL_SETBACK_FT))
            else:
                remnant = fix(i2g.difference(excl))
                if remnant.is_empty:
                    siteable_per_parcel.append(0.0)
                    continue
                eroded = fix(remnant.buffer(-DC_INTERNAL_SETBACK_FT))
            siteable_per_parcel.append(max_poly(eroded))
        strict_counts[ft] = strict
        for thr in BUILDING_THRESHOLDS:
            results[(ft, thr)] = sum(1 for s in siteable_per_parcel if s >= thr)
        print(f"  {ft}ft done strict={strict} ({time.time()-t1:.1f}s)")

    # Pretty matrix
    OUT.mkdir(exist_ok=True)
    out_path = OUT / "sweep_matrix.txt"
    with open(out_path, "w") as f:
        f.write("Johnson City Data Center site-eligibility sweep\n")
        f.write(f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"I-2 candidate parcels: {len(cands):,}\n")
        f.write(f"Setbacks: {SETBACKS_FT} (200 = ordinance, others exploratory)\n")
        f.write(f"DC internal yard setback: {DC_INTERNAL_SETBACK_FT}ft (side/rear)\n")
        f.write("\nStrict eligibility (entire parcel >= setback from residential):\n")
        for ft in SETBACKS_FT:
            f.write(f"  {ft:>4} ft: {strict_counts[ft]:>4}\n")
        f.write("\nSite-able count by min DC BUILDING footprint (sqft) x setback (ft):\n")
        header = "  building sqft  |" + "".join(f"  {ft:>5} ft" for ft in SETBACKS_FT) + "\n"
        f.write(header)
        f.write("  " + "-" * (len(header) - 3) + "\n")
        for thr in BUILDING_THRESHOLDS:
            row = f"  {thr:>13,} |"
            for ft in SETBACKS_FT:
                row += f"  {results[(ft, thr)]:>8}"
            f.write(row + "\n")
        f.write("\nNotes:\n")
        f.write("  - Site-able = largest contiguous compliant patch eroded by 50ft\n"
                "    (the DC's own side/rear yard per draft 6.20.3.2.I).\n"
                "  - Counts are NECESSARY-but-not-sufficient: every parcel still\n"
                "    needs BZA special-exception approval.\n")
    print(f"\nwrote {out_path}")
    print("\n--- sweep_matrix.txt ---")
    print(out_path.read_text())
    return 0


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