"""Compare 200ft vs 500ft (vs 1000ft) setbacks in terms that matter for
political argument: how many homes get protected, and which DC sites disappear.

Outputs:
  out/comparison_summary.txt        plain-text headline numbers
  out/homes_exposure.csv            per-dwelling distance to nearest DC patch
  out/marginal_parcels_500.geojson  parcels eligible at 200 but NOT at 500
  out/marginal_parcels_1000.geojson parcels eligible at 200 but NOT at 1000
  web/data/comparison.json          consumed by the web page for live display
"""
from __future__ import annotations
import json
import sys
import time
from pathlib import Path

import geopandas as gpd
import pandas as pd
from shapely import make_valid
from shapely.geometry import MultiPolygon, Polygon
from shapely.strtree import STRtree

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

RESIDENTIAL_LANDUSE_CODES = {"11", "12", "13", "14"}
RESIDENTIAL_PLACETYPES = {
    "SF RESIDENCE", "APARTMENT", "APARTMENTS", "MOBILE HOME",
    "CONDOMINIUM", "TOWNHOUSE", "DUPLEX", "PUBLIC HOUSING", "DORMITORY",
}
# Distance buckets reported in the summary table (ft). 5280 = 1 mile.
EXPOSURE_BUCKETS = [200, 500, 1000, 1500, 2000, 2640, 5280]


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


def stamp(msg, t0=None):
    now = time.time()
    if t0 is None:
        print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
    else:
        print(f"[{time.strftime('%H:%M:%S')}] {msg} ({now - t0:.1f}s)", flush=True)
    return now


def main() -> int:
    t0 = stamp("loading layers")
    addresses = gpd.read_parquet(CACHE / "addresses.parquet")
    parcels = gpd.read_parquet(CACHE / "parcels.parquet")
    crs = parcels.crs

    # Eligible (buildable + frontage-validated) GeoJSONs from analyze.py.
    # These are in EPSG:4326; reproject back to EPSG:2274 for distance math.
    elig = {}
    for ft in (200, 500, 1000):
        p = OUT / f"eligible_{ft}ft_buildable.geojson"
        if p.exists() and p.stat().st_size > 100:
            g = gpd.read_file(p).to_crs(crs)
            elig[ft] = g
        else:
            elig[ft] = gpd.GeoDataFrame({"geometry": []}, geometry="geometry", crs=crs)
    stamp("loaded", t0)
    for ft, g in elig.items():
        print(f"  buildable@{ft}ft: {len(g)} parcels")

    # ----- Residential dwellings = address points with residential PLACETYPE
    t1 = stamp("identifying dwellings")
    res_mask = addresses["PLACETYPE"].astype(str).str.upper().isin(RESIDENTIAL_PLACETYPES)
    homes = addresses[res_mask].copy()
    print(f"  dwelling address points: {len(homes):,}")
    stamp("done", t1)

    # ----- For each home, distance (ft) to nearest buildable patch per setback
    summary = {}
    home_dist_cols = {}
    for ft in (200, 500, 1000):
        t1 = stamp(f"computing per-home distance to nearest {ft}ft buildable patch")
        if len(elig[ft]) == 0:
            home_dist_cols[ft] = [float("inf")] * len(homes)
            stamp("  (no eligible patches)", t1)
            continue
        tree = STRtree(elig[ft].geometry.values)
        polys = elig[ft].geometry.values
        dists = []
        for pt in homes.geometry.values:
            # nearest polygon, then exact distance from point to that polygon.
            i = tree.nearest(pt)
            d = polys[i].distance(pt)
            dists.append(d)
        home_dist_cols[ft] = dists
        stamp("done", t1)

    homes["dist_200_ft"] = home_dist_cols[200]
    homes["dist_500_ft"] = home_dist_cols[500]
    homes["dist_1000_ft"] = home_dist_cols[1000]

    # ----- Exposure buckets per setback
    for ft in (200, 500, 1000):
        d = homes[f"dist_{ft}_ft"]
        summary[f"setback_{ft}"] = {
            "buildable_patches": int(len(elig[ft])),
            "homes_within_ft": {
                str(b): int((d <= b).sum()) for b in EXPOSURE_BUCKETS
            },
        }

    # ----- Homes "protected" by tightening from 200 -> 500 (or -> 1000)
    # Definition: home is within X ft of any patch at 200ft setback, but is
    # >= X ft from any patch at 500ft setback. I.e., the patches that pushed
    # the home into the exposure ring are exactly the ones that go away at
    # the tighter setback.
    protection = {}
    for tighter in (500, 1000):
        rows = {}
        for b in EXPOSURE_BUCKETS:
            exposed_at_200 = homes["dist_200_ft"] <= b
            still_exposed = homes[f"dist_{tighter}_ft"] <= b
            rows[str(b)] = int((exposed_at_200 & ~still_exposed).sum())
        protection[f"to_{tighter}"] = {
            "newly_protected_within_ft": rows,
        }
    summary["protection"] = protection

    # ----- Homes newly protected: within 1000ft of a possible DC at 200ft
    # setback, but >= 1000ft at the tighter setback. The 1000ft threshold is
    # the headline "close-range exposure" distance.
    PROTECT_BUCKET_FT = 1000
    for tighter in (500, 1000):
        exposed_200 = homes["dist_200_ft"] <= PROTECT_BUCKET_FT
        still_exposed = homes[f"dist_{tighter}_ft"] <= PROTECT_BUCKET_FT
        newly = homes[exposed_200 & ~still_exposed].copy()
        # Reproject to EPSG:4326 for the web
        newly_web = newly.to_crs("EPSG:4326")
        keep_cols = ["FULLADDR", "PLACETYPE", "dist_200_ft", f"dist_{tighter}_ft"]
        newly_web = newly_web[[c for c in keep_cols if c in newly_web.columns] + ["geometry"]]
        out_p = OUT / f"homes_protected_{tighter}.geojson"
        if out_p.exists():
            out_p.unlink()
        if len(newly_web):
            newly_web.to_file(out_p, driver="GeoJSON")
        print(f"  homes newly protected -> {tighter}ft (within {PROTECT_BUCKET_FT}ft): {len(newly_web)}")

    # ----- Marginal parcels: eligible at 200 but lost at 500 / 1000
    e200 = elig[200]
    for tighter in (500, 1000):
        et = elig[tighter]
        tighter_ids = set(et["GISLINK"].dropna().astype(str).tolist()) if "GISLINK" in et.columns else set()
        if "GISLINK" not in e200.columns:
            marginal = gpd.GeoDataFrame(columns=e200.columns)
        else:
            marginal = e200[~e200["GISLINK"].astype(str).isin(tighter_ids)].copy()
        out_p = OUT / f"marginal_parcels_{tighter}.geojson"
        if out_p.exists():
            out_p.unlink()
        if len(marginal):
            marginal.to_file(out_p, driver="GeoJSON")
        print(f"  marginal at -> {tighter}ft: {len(marginal)} parcels lost")
        summary[f"marginal_to_{tighter}"] = int(len(marginal))

    # ----- Write outputs
    OUT.mkdir(exist_ok=True)
    homes_out = homes[["FULLADDR", "PLACETYPE", "dist_200_ft", "dist_500_ft", "dist_1000_ft"]].copy()
    for c in ("dist_200_ft", "dist_500_ft", "dist_1000_ft"):
        homes_out[c] = homes_out[c].round(1)
    homes_out.to_csv(OUT / "homes_exposure.csv", index=False)
    print(f"  wrote out/homes_exposure.csv ({len(homes_out):,} rows)")

    # Plain-text headline
    with open(OUT / "comparison_summary.txt", "w") as f:
        f.write("Johnson City data-center setback comparison\n")
        f.write(f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Dwelling address points considered: {len(homes):,}\n\n")
        f.write("Buildable-subdividable patches:\n")
        for ft in (200, 500, 1000):
            f.write(f"  {ft:>4} ft setback: {summary[f'setback_{ft}']['buildable_patches']:>4} patches\n")
        f.write("\nHomes within distance D of any potential data center:\n")
        f.write(f"  {'D (ft)':>8} | {'200ft':>8} | {'500ft':>8} | {'1000ft':>8}\n")
        f.write("  " + "-" * 44 + "\n")
        for b in EXPOSURE_BUCKETS:
            r = [summary[f"setback_{ft}"]["homes_within_ft"][str(b)] for ft in (200, 500, 1000)]
            f.write(f"  {b:>8,} | {r[0]:>8,} | {r[1]:>8,} | {r[2]:>8,}\n")
        f.write("\nHomes NEWLY PROTECTED by tightening setback from 200ft:\n")
        f.write(f"  {'D (ft)':>8} | {'200->500':>10} | {'200->1000':>10}\n")
        f.write("  " + "-" * 36 + "\n")
        for b in EXPOSURE_BUCKETS:
            a = protection["to_500"]["newly_protected_within_ft"][str(b)]
            c = protection["to_1000"]["newly_protected_within_ft"][str(b)]
            f.write(f"  {b:>8,} | {a:>10,} | {c:>10,}\n")
        f.write(f"\nMarginal DC sites lost by tightening 200 -> 500:  {summary['marginal_to_500']}\n")
        f.write(f"Marginal DC sites lost by tightening 200 -> 1000: {summary['marginal_to_1000']}\n")
    print(f"  wrote out/comparison_summary.txt")
    print("\n--- comparison_summary.txt ---")
    print((OUT / "comparison_summary.txt").read_text())

    # JSON for the web page
    WEB.mkdir(parents=True, exist_ok=True)
    with open(WEB / "comparison.json", "w") as f:
        json.dump({
            "generated_at": time.strftime("%Y-%m-%d %H:%M:%S"),
            "total_homes": int(len(homes)),
            "buckets_ft": EXPOSURE_BUCKETS,
            **summary,
        }, f, indent=2)
    print(f"  wrote web/data/comparison.json")

    # Copy comparison artifacts into web/
    for stem in ("marginal_parcels_500", "marginal_parcels_1000",
                 "homes_protected_500", "homes_protected_1000"):
        src = OUT / f"{stem}.geojson"
        if src.exists():
            dst = WEB / src.name
            dst.write_bytes(src.read_bytes())
            print(f"  copied -> web/data/{src.name}")

    return 0


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