"""Exp 4 Phase 2 — training-crop exporter (DOGIS 2025 imagery + supervision rasters).

Samples residential Douglas County parcels in LiDAR-tile-aligned clusters and, for
each parcel, exports one training crop:

    images/{id}.jpg        L21 (~5.6 cm/px) DOGIS mosaic, CHIP_M meters square,
                           centered on the parcel centroid
    rasters/{id}_parcel.png   0=bg, 120=neighbor parcels, 255=the parcel
    rasters/{id}_road.png     street-centerline buffer (255 inside)
    rasters/{id}_lidar.png    RGB: R=building (class 6), G=ground (class 2),
                              B=height-above-ground (0.1 m units, clipped 25.5 m)
                              at RASTER_DOWN-pixel cells on the SAME crop grid
    meta/{id}.json         provenance + geotransform (Web-Mercator)
    manifest.csv           one row per crop (append-only; resume skips done ids)

All supervision is pre-rasterized into the crop's own pixel grid, so the SAM-3
labeler downstream needs NO geo dependencies (torch/transformers/PIL/numpy only).

Eval contamination guard: every address in EXCLUSION_FILES is geocoded (cached)
and any parcel within EXCLUDE_RADIUS_M of one is never sampled.

Runs on the Windows dev box or the Mac — needs the repo installed
(`pip install -e ".[api]"`) plus `data/ne_parcels/ne_parcels.gpkg`
(private HF dataset TempuraML/ne-parcels). Unattended, idempotent, resumable;
per-stage errors are recorded and the run always ends with a REPORT block.

    python scripts/exp4/export_training_crops.py --out data/exp4/crops \
        --n-crops 50 --per-cluster 17 --prefer-cached-tiles --seed 42
"""

from __future__ import annotations

import argparse
import csv
import io
import json
import math
import re
import time
import traceback
from pathlib import Path

import numpy as np
import requests
from PIL import Image, ImageDraw
from pyproj import Transformer
from scipy import ndimage
from shapely.geometry import box as shapely_box
from shapely.ops import transform as shapely_transform

from lawn_estimator.sources.lidar import DouglasS3LidarSource

REPO_ROOT = Path(__file__).resolve().parents[2]

SERVICE = "https://dcgis.org/server/rest/services/2025_Douglas_County_NE_Imagery/MapServer"
GEOCODER = "https://dcgis.org/server/rest/services/vector/Address_Points/FeatureServer/0/query"
STREETS = "https://dcgis.org/server/rest/services/vector/Street_Centerlines/FeatureServer/0/query"
PARCELS_GPKG = REPO_ROOT / "data" / "ne_parcels" / "ne_parcels.gpkg"

# DOGIS 2025 service fullExtent (EPSG:3857) — the hard coverage boundary.
IMAGERY_EXTENT = (-10740413.9352, 5040042.0973, -10671890.2482, 5071228.3064)

TILE_SIZE = 256
LEVEL = 21                  # ~5.6 cm/px ground at this latitude (probe-verified native detail)
CHIP_M = 96.0               # ground meters per crop side
RASTER_DOWN = 12            # supervision-raster cell = 12 crop px (~0.67 ground m)
ROAD_HALF_WIDTH_M = 5.5     # ground meters — generous centerline buffer prior
EXCLUDE_RADIUS_M = 120.0    # no training crop this close to an eval/QA address
MIN_SPACING_M = 60.0        # between sampled parcel centroids within a cluster
EARTH = 20037508.342787

EXCLUSION_FILES = [
    REPO_ROOT / "data" / "evals" / "addresses" / "exp5_green_reclaim.csv",
    REPO_ROOT / "data" / "lawn_care_schedule.csv",
    REPO_ROOT / "data" / "regression" / "douglas_qa.csv",
]

# BuildingYear strata (targets enforced across clusters): established canopy-heavy,
# mid-vintage, and new construction (dormant-turf-rich — the exp5 failure domain).
YEAR_BUCKETS = [(0, 1985), (1985, 2013), (2013, 3000)]

TO_3857 = Transformer.from_crs("EPSG:4326", "EPSG:3857", always_xy=True)
TO_4326 = Transformer.from_crs("EPSG:3857", "EPSG:4326", always_xy=True)


# ---------------------------------------------------------------- mercator math
def merc_to_pixel(mx: float, my: float, level: int) -> tuple[float, float]:
    world = TILE_SIZE * 2.0 ** level
    px = (mx + EARTH) / (2 * EARTH) * world
    py = (EARTH - my) / (2 * EARTH) * world
    return px, py


def ground_mpp(lat: float, level: int) -> float:
    return (2 * EARTH / (TILE_SIZE * 2.0 ** level)) * math.cos(math.radians(lat))


def fetch_tile(session: requests.Session, z: int, row: int, col: int, retries: int = 3) -> Image.Image:
    url = f"{SERVICE}/tile/{z}/{row}/{col}"
    last = None
    for attempt in range(retries):
        try:
            r = session.get(url, timeout=30)
            if r.status_code == 200 and r.headers.get("Content-Type", "").startswith("image"):
                return Image.open(io.BytesIO(r.content)).convert("RGB")
            last = f"HTTP {r.status_code}"
        except requests.RequestException as e:
            last = repr(e)
        time.sleep(1.5 * (attempt + 1))
    raise RuntimeError(f"tile {z}/{row}/{col}: {last}")


def fetch_chip_3857(session: requests.Session, bbox: tuple[float, float, float, float]) -> Image.Image:
    """Mosaic covering an EPSG:3857 bbox at LEVEL, cropped exactly to it."""
    x0, y0 = merc_to_pixel(bbox[0], bbox[3], LEVEL)   # top-left
    x1, y1 = merc_to_pixel(bbox[2], bbox[1], LEVEL)   # bottom-right
    x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1)
    c0, c1, r0, r1 = x0 // TILE_SIZE, x1 // TILE_SIZE, y0 // TILE_SIZE, y1 // TILE_SIZE
    mosaic = Image.new("RGB", ((c1 - c0 + 1) * TILE_SIZE, (r1 - r0 + 1) * TILE_SIZE))
    for row in range(r0, r1 + 1):
        for col in range(c0, c1 + 1):
            mosaic.paste(fetch_tile(session, LEVEL, row, col),
                         ((col - c0) * TILE_SIZE, (row - r0) * TILE_SIZE))
    return mosaic.crop((x0 - c0 * TILE_SIZE, y0 - r0 * TILE_SIZE,
                        x1 - c0 * TILE_SIZE, y1 - r0 * TILE_SIZE))


# ------------------------------------------------------------------ exclusions
def load_exclusions(session: requests.Session, cache_path: Path) -> list[tuple[float, float]]:
    """Geocode every eval/QA address once (cached). Returns EPSG:3857 points."""
    addresses: set[str] = set()
    for f in EXCLUSION_FILES:
        text = f.read_text(encoding="utf-8", errors="replace")
        for m in re.finditer(r'"?([0-9]+ [^,"]+)(?:,| in )\s*(?:Omaha|Elkhorn|Bennington|Valley|Waterloo|Bellevue|Papillion|Gretna|La Vista|Springfield)', text):
            addresses.add(m.group(1).strip().upper())
    cache: dict = json.loads(cache_path.read_text()) if cache_path.exists() else {}
    pts = []
    for addr in sorted(addresses):
        if addr not in cache:
            try:
                r = session.get(GEOCODER, params={
                    "where": f"FULLADDR LIKE '{addr}%'", "outFields": "FULLADDR",
                    "returnGeometry": "true", "outSR": "3857", "f": "json"}, timeout=30)
                feats = r.json().get("features") or []
                cache[addr] = [feats[0]["geometry"]["x"], feats[0]["geometry"]["y"]] if feats else None
            except Exception:
                cache[addr] = None
            cache_path.write_text(json.dumps(cache, indent=1))
        if cache[addr]:
            pts.append(tuple(cache[addr]))
    misses = [a for a in sorted(addresses) if not cache.get(a)]
    print(f"exclusions: {len(pts)} geocoded, {len(misses)} not in Douglas Address_Points "
          f"(Sarpy/bad addresses — outside the sampling universe anyway): {misses}")
    return pts


# ------------------------------------------------------------------- sampling
def residential_parcels(bbox_3857: tuple):
    import geopandas as gpd
    g = gpd.read_file(PARCELS_GPKG, bbox=bbox_3857)
    if g.empty:
        return g
    try:
        g["year"] = g["BuildingYear"].replace("", np.nan).astype(float)
    except Exception:
        g["year"] = np.nan
    has_building = (g["ImpSF"].fillna(0) > 400) | (g["Improvements_Value"].fillna(0) > 20000)
    size_ok = g["GIS_Acres"].fillna(0).between(0.07, 0.7)
    # Residential only — calibrated on pilot parcels (Douglas): single-family homes
    # carry assessor Classification_Code prefix "0101"; the commercial lot that
    # leaked through the building+size proxy was "0103…".
    residential = g["Classification_Code"].fillna("").str.startswith("0101")
    return g[has_building & size_ok & residential & g.geometry.notna()].copy()


def merc_m(ground_m: float, lat: float) -> float:
    return ground_m / math.cos(math.radians(lat))


# -------------------------------------------------------------------- rasters
def rasterize_polys(polys_px: list[list[tuple[float, float]]], values: list[int],
                    size: int) -> Image.Image:
    img = Image.new("L", (size, size), 0)
    d = ImageDraw.Draw(img)
    for pts, v in zip(polys_px, values, strict=True):
        if len(pts) >= 3:
            d.polygon(pts, fill=v)
    return img


def lidar_raster(points: np.ndarray, bbox: tuple, cells: int, cell_merc: float) -> Image.Image:
    """RGB supervision raster on the crop grid: R=building, G=ground, B=HAG*10."""
    x0, y0, x1, y1 = bbox  # 3857; y grows north, raster row 0 = north edge
    sel = (points["X"] >= x0) & (points["X"] < x1) & (points["Y"] > y0) & (points["Y"] <= y1)
    p = points[sel]
    out = np.zeros((cells, cells, 3), dtype=np.uint8)
    if len(p) == 0:
        return Image.fromarray(out)
    ci = np.clip(((p["X"] - x0) / cell_merc).astype(int), 0, cells - 1)
    ri = np.clip(((y1 - p["Y"]) / cell_merc).astype(int), 0, cells - 1)

    building = np.zeros((cells, cells), dtype=bool)
    ground = np.zeros((cells, cells), dtype=bool)
    np.logical_or.at(building, (ri[p["Classification"] == 6], ci[p["Classification"] == 6]), True)
    np.logical_or.at(ground, (ri[p["Classification"] == 2], ci[p["Classification"] == 2]), True)
    building = ndimage.binary_closing(building, iterations=1)
    ground = ndimage.binary_closing(ground, iterations=1)

    zmax = np.full((cells, cells), -np.inf)
    np.maximum.at(zmax, (ri, ci), p["Z"])
    gmin = np.full((cells, cells), np.inf)
    g = p[p["Classification"] == 2]
    np.minimum.at(gmin, (np.clip(((y1 - g["Y"]) / cell_merc).astype(int), 0, cells - 1),
                         np.clip(((g["X"] - x0) / cell_merc).astype(int), 0, cells - 1)), g["Z"])
    have_g = np.isfinite(gmin)
    if have_g.any():
        _, idx = ndimage.distance_transform_edt(~have_g, return_indices=True)
        gfill = gmin[idx[0], idx[1]]
        hag = np.where(np.isfinite(zmax), np.clip(zmax - gfill, 0, 25.5), 0)
    else:
        hag = np.zeros((cells, cells))
    out[..., 0] = building * 255
    out[..., 1] = ground * 255
    out[..., 2] = (hag * 10).astype(np.uint8)
    return Image.fromarray(out)


def fetch_roads(session: requests.Session, bbox: tuple) -> list[list[tuple[float, float]]]:
    r = session.get(STREETS, params={
        "geometry": json.dumps({"xmin": bbox[0], "ymin": bbox[1], "xmax": bbox[2],
                                "ymax": bbox[3], "spatialReference": {"wkid": 3857}}),
        "geometryType": "esriGeometryEnvelope", "spatialRel": "esriSpatialRelIntersects",
        "outFields": "OBJECTID", "returnGeometry": "true", "outSR": "3857", "f": "json"},
        timeout=60)
    r.raise_for_status()
    paths = []
    for f in r.json().get("features", []):
        for path in (f.get("geometry") or {}).get("paths", []):
            paths.append([(x, y) for x, y in path])
    return paths


def road_raster(paths: list, bbox: tuple, size_px: int, mpp_merc: float, lat: float) -> Image.Image:
    img = Image.new("L", (size_px, size_px), 0)
    d = ImageDraw.Draw(img)
    width_px = max(2, int(2 * merc_m(ROAD_HALF_WIDTH_M, lat) / mpp_merc))
    x0, _, _, y1 = bbox
    for path in paths:
        pts = [((x - x0) / mpp_merc, (y1 - y) / mpp_merc) for x, y in path]
        if len(pts) >= 2:
            d.line(pts, fill=255, width=width_px)
    return img


# ------------------------------------------------------------------------ main
def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default=str(REPO_ROOT / "data" / "exp4" / "crops"))
    ap.add_argument("--n-crops", type=int, default=50)
    ap.add_argument("--per-cluster", type=int, default=17)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--prefer-cached-tiles", action="store_true",
                    help="sample only LiDAR tiles already in data/lidar (fast pilot)")
    ap.add_argument("--purge-lidar", action="store_true",
                    help="delete LiDAR tiles downloaded per cluster after use "
                         "(pre-existing cached tiles are kept; caps disk at ~1 cluster)")
    args = ap.parse_args()

    out = Path(args.out)
    for sub in ("images", "rasters", "meta"):
        (out / sub).mkdir(parents=True, exist_ok=True)
    manifest_path = out / "manifest.csv"
    done_ids = set()
    if manifest_path.exists():
        with open(manifest_path, newline="", encoding="utf-8") as f:
            done_ids = {row["id"] for row in csv.DictReader(f)}

    print("=" * 66)
    print("CROP EXPORT REPORT BEGIN — copy everything down to REPORT END")
    print("=" * 66)
    print(f"out: {out} | target: {args.n_crops} crops | per-cluster: {args.per_cluster} "
          f"| seed: {args.seed} | already done: {len(done_ids)}")

    rng = np.random.default_rng(args.seed)
    session = requests.Session()
    session.headers["User-Agent"] = "lawn-estimator-exp4-export/1.0"
    report = {"clusters": 0, "crops": 0, "skipped_done": 0, "failures": 0}

    try:
        exclusions = load_exclusions(session, out / "exclusions_geocoded.json")
    except Exception:
        print("EXCLUSION LOAD FAILED — refusing to sample without the contamination guard")
        print(traceback.format_exc())
        print("REPORT END")
        return
    excl = np.array(exclusions) if exclusions else np.zeros((0, 2))

    lidar_src = DouglasS3LidarSource(session=session)
    tiles = lidar_src._load_index().to_crs("EPSG:3857")
    ext = shapely_box(*IMAGERY_EXTENT).buffer(-500)  # stay inside imagery coverage
    tiles = tiles[tiles.geometry.within(ext)]
    if args.prefer_cached_tiles:
        cached = {p.stem for p in (REPO_ROOT / "data" / "lidar").glob("*.las")}
        tiles = tiles[tiles["NAME"].isin(cached)]
        print(f"cached-tile mode: {len(tiles)} candidate tiles")
    order = rng.permutation(len(tiles))
    bucket_counts = {i: 0 for i in range(len(YEAR_BUCKETS))}
    per_bucket_target = math.ceil(args.n_crops / args.per_cluster / len(YEAR_BUCKETS))

    for tix in order:
        if report["crops"] + report["skipped_done"] >= args.n_crops:
            break
        tile = tiles.iloc[tix]
        try:
            inner = tile.geometry.buffer(-merc_m(CHIP_M / 2 + 10, 41.26))
            if inner.is_empty:
                continue
            parcels = residential_parcels(tile.geometry.bounds)
            if len(parcels) < args.per_cluster:
                continue
            parcels = parcels[parcels.geometry.centroid.within(inner)]
            if len(parcels) < args.per_cluster:
                continue
            med = np.nanmedian(parcels["year"]) if np.isfinite(parcels["year"]).any() else 1990
            bucket = next(i for i, (a, b) in enumerate(YEAR_BUCKETS) if a <= med < b)
            if bucket_counts[bucket] >= per_bucket_target and not args.prefer_cached_tiles:
                continue

            # LiDAR once per cluster (pass the shrunk tile so neighbors don't download)
            las_before = {p.name for p in (REPO_ROOT / "data" / "lidar").glob("*.las")}
            geom_wgs84 = shapely_transform(TO_4326.transform, tile.geometry.buffer(-1))
            points = lidar_src.points_for_geometry(geom_wgs84, "EPSG:3857")
            if len(points) == 0:
                print(f"cluster {tile['NAME']}: no LiDAR points — skipped")
                continue
            roads_all = fetch_roads(session, tile.geometry.bounds)

            # spaced random sample of parcels, exclusion-guarded
            chosen = []
            cents = parcels.geometry.centroid
            for pix in rng.permutation(len(parcels)):
                c = cents.iloc[pix]
                if excl.shape[0] and np.min(np.hypot(excl[:, 0] - c.x, excl[:, 1] - c.y)) \
                        < merc_m(EXCLUDE_RADIUS_M, 41.26):
                    continue
                if any(c.distance(cents.iloc[j]) < merc_m(MIN_SPACING_M, 41.26) for j in chosen):
                    continue
                chosen.append(pix)
                if len(chosen) >= args.per_cluster:
                    break

            made = 0
            for pix in chosen:
                row = parcels.iloc[pix]
                pid = str(row.get("Parcel_ID") or row.get("State_PID") or f"{tile['NAME']}_{pix}")
                crop_id = f"{tile['NAME']}_{pid}".replace("/", "_").replace(" ", "")
                if crop_id in done_ids:
                    report["skipped_done"] += 1
                    continue
                try:
                    c = cents.iloc[pix]
                    lat = TO_4326.transform(c.x, c.y)[1]
                    half = merc_m(CHIP_M / 2, lat)
                    bbox = (c.x - half, c.y - half, c.x + half, c.y + half)
                    chip = fetch_chip_3857(session, bbox)
                    size = chip.size[0]
                    chip = chip.crop((0, 0, size, size))  # enforce square
                    mpp_merc = 2 * half / size

                    def to_px(geom, bbox=bbox, mpp_merc=mpp_merc):
                        # statewide gpkg is MultiPolygon Z — coords are (x, y, z)
                        ext_pts = list(geom.exterior.coords) if geom.geom_type == "Polygon" \
                            else [pt for part in geom.geoms for pt in part.exterior.coords]
                        return [((pt[0] - bbox[0]) / mpp_merc, (bbox[3] - pt[1]) / mpp_merc)
                                for pt in ext_pts]

                    neighbors = parcels[parcels.geometry.intersects(shapely_box(*bbox))]
                    polys, vals = [], []
                    for _, nb in neighbors.iterrows():
                        if str(nb.get("Parcel_ID")) != str(row.get("Parcel_ID")):
                            polys.append(to_px(nb.geometry))
                            vals.append(120)
                    polys.append(to_px(row.geometry))
                    vals.append(255)
                    rasterize_polys(polys, vals, size).save(out / "rasters" / f"{crop_id}_parcel.png")

                    road_raster(roads_all, bbox, size, mpp_merc, lat).save(
                        out / "rasters" / f"{crop_id}_road.png")

                    cells = max(1, size // RASTER_DOWN)
                    lidar_raster(points, bbox, cells, 2 * half / cells).save(
                        out / "rasters" / f"{crop_id}_lidar.png")

                    chip.save(out / "images" / f"{crop_id}.jpg", quality=92)
                    meta = {
                        "id": crop_id, "parcel_id": pid,
                        "situs": str(row.get("Situs_Address", "")),
                        "year_built": None if not np.isfinite(row["year"]) else int(row["year"]),
                        "acres": float(row.get("GIS_Acres") or 0),
                        "bbox_3857": list(bbox), "level": LEVEL, "size_px": size,
                        "mpp_merc": mpp_merc, "ground_mpp": ground_mpp(lat, LEVEL),
                        "chip_m": CHIP_M, "raster_down": RASTER_DOWN,
                        "lidar_cells": cells, "cluster": str(tile["NAME"]),
                        "year_bucket": bucket,
                    }
                    (out / "meta" / f"{crop_id}.json").write_text(json.dumps(meta, indent=1))
                    new_file = not manifest_path.exists()
                    with open(manifest_path, "a", newline="", encoding="utf-8") as f:
                        w = csv.writer(f)
                        if new_file:
                            w.writerow(["id", "cluster", "situs", "year_built", "acres", "bucket"])
                        w.writerow([crop_id, tile["NAME"], meta["situs"],
                                    meta["year_built"], round(meta["acres"], 3), bucket])
                    done_ids.add(crop_id)
                    made += 1
                    report["crops"] += 1
                    if report["crops"] + report["skipped_done"] >= args.n_crops:
                        break
                except Exception:
                    report["failures"] += 1
                    print(f"crop {crop_id}: FAILED")
                    print(traceback.format_exc())
            if args.purge_lidar:
                del points  # release the mmap/array before unlinking on Windows
                for p in (REPO_ROOT / "data" / "lidar").glob("*.las"):
                    if p.name not in las_before:
                        try:
                            p.unlink()
                        except OSError as e:
                            print(f"  purge failed for {p.name}: {e}")
            if made:
                bucket_counts[bucket] += 1
                report["clusters"] += 1
                print(f"cluster {tile['NAME']}: +{made} crops (median year {med:.0f}, bucket {bucket})")
        except Exception:
            report["failures"] += 1
            print(f"cluster {tile.get('NAME', '?')}: FAILED")
            print(traceback.format_exc())

    print(f"\nsummary: {json.dumps(report)} | bucket_clusters: {bucket_counts}")
    print(f"manifest: {manifest_path} ({len(done_ids)} total crops)")
    print("=" * 66)
    print("REPORT END")
    print("=" * 66)


if __name__ == "__main__":
    main()