"""Fast vector geoprocessing with GeoPandas and Shapely 2.

Benchmarks point-in-polygon joins (naive loop vs STRtree vs sjoin), buffers,
overlays and dissolves, then compares Shapefile, GeoPackage, GeoParquet and
FlatGeobuf for writing, reading and bounding-box queries.
Run with:  python vector_performance.py
"""
import time
import urllib.request
from contextlib import contextmanager
from pathlib import Path

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import shapely
from shapely.geometry import box

HERE = Path(__file__).parent
DATA, OUT = HERE / "data", HERE / "output"
DATA.mkdir(exist_ok=True)
OUT.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 120)
rng = np.random.default_rng(42)
timings = {}


def section(title: str) -> None:
    print(f"\n=== {title} ===")


@contextmanager
def timer(label: str):
    t0 = time.perf_counter()
    yield
    timings[label] = time.perf_counter() - t0
    print(f"{label:<42s} {timings[label]:8.3f} s")


def fetch(name: str) -> Path:
    """Download a Natural Earth zip (scale taken from the file name) once into ./data."""
    dest = DATA / name
    if not dest.exists():
        scale = name.split("_")[1]  # e.g. "10m"
        urllib.request.urlretrieve(f"https://naciscdn.org/naturalearth/{scale}/cultural/{name}", dest)
    return dest


section("versions")
print("geopandas", gpd.__version__, "| shapely", shapely.__version__, "| GEOS", shapely.geos_version_string)

# --- Step 1: build a realistic test dataset -------------------------------------
# 250,000 "incident" points clustered around UK towns (weighted by population),
# plus a 5 km analysis grid over Great Britain, all in British National Grid.
countries = gpd.read_file(fetch("ne_50m_admin_0_countries.zip"), engine="pyogrio")
uk = countries.loc[countries.ADM0_A3 == "GBR", ["geometry"]].to_crs(27700)
gb = uk.explode(index_parts=False).reset_index(drop=True)  # parts share one label until reset
gb = gb.loc[[gb.area.idxmax()]]  # the island of Great Britain only
places = gpd.read_file(fetch("ne_10m_populated_places.zip"), engine="pyogrio")
towns = places.loc[places.ADM0_A3 == "GBR"].to_crs(27700).clip(gb)

N = 250_000
w = towns.POP_MAX.to_numpy(float)
centre = rng.choice(len(towns), size=N, p=w / w.sum())
xy = np.column_stack([towns.geometry.x.to_numpy()[centre], towns.geometry.y.to_numpy()[centre]])
xy += rng.normal(scale=8_000, size=(N, 2))
noise = rng.random(N) < 0.25  # a quarter of points spread uniformly
minx, miny, maxx, maxy = gb.total_bounds
xy[noise] = np.column_stack([rng.uniform(minx, maxx, noise.sum()), rng.uniform(miny, maxy, noise.sum())])
points = gpd.GeoDataFrame(
    {"id": np.arange(N), "category": rng.choice(["burglary", "vehicle", "criminal damage", "theft"], N),
     "value": rng.gamma(2.0, 50.0, N).round(2)},
    geometry=gpd.points_from_xy(xy[:, 0], xy[:, 1]), crs=27700)
points = points[points.intersects(gb.geometry.iloc[0])].reset_index(drop=True)

cell = 5_000
gx, gy = np.meshgrid(np.arange(minx // cell * cell, maxx, cell), np.arange(miny // cell * cell, maxy, cell))
grid = gpd.GeoDataFrame(geometry=shapely.box(gx.ravel(), gy.ravel(), gx.ravel() + cell, gy.ravel() + cell),
                        crs=27700)
grid = grid[grid.intersects(gb.geometry.iloc[0])].reset_index(drop=True)
grid["cell_id"] = np.arange(len(grid))
# 100 km "region" code, handy for dissolving later
grid["region"] = ((grid.geometry.bounds.minx // 100_000).astype(int).astype(str) + "-"
                  + (grid.geometry.bounds.miny // 100_000).astype(int).astype(str))

section("test data")
print(f"towns used as cluster centres: {len(towns)}")
print(f"points: {len(points):,}  grid cells: {len(grid):,}  regions: {grid.region.nunique()}")
print(points.head().to_string())

# --- Step 2: point-in-polygon, three ways -----------------------------------------
section("point-in-polygon")
sample = points.iloc[:2_000]
with timer("naive loop (2,000 points)"):
    naive = []
    for pt in sample.geometry:
        hit = grid.index[grid.contains(pt)]
        naive.append(hit[0] if len(hit) else -1)
per_point = timings["naive loop (2,000 points)"] / len(sample)
print(f"{'  -> extrapolated to all points':<42s} {per_point * len(points):8.1f} s")
timings["naive loop (extrapolated)"] = per_point * len(points)

with timer("STRtree.query, all points"):
    tree = shapely.STRtree(grid.geometry.values.data)
    pt_idx, poly_idx = tree.query(points.geometry.values.data, predicate="within")

with timer("gpd.sjoin, all points"):
    joined = gpd.sjoin(points, grid[["cell_id", "geometry"]], how="inner", predicate="within")

assert np.array_equal(np.sort(pt_idx), np.sort(joined.index.to_numpy()))
check = pd.Series(poly_idx, index=pt_idx).sort_index()
assert (check.iloc[:2_000].to_numpy() == np.array(naive)[check.index[check.index < 2_000]]).all()
print(f"matched {len(joined):,} of {len(points):,} points; all three methods agree")
speedup = timings["naive loop (extrapolated)"] / timings["gpd.sjoin, all points"]
print(f"sjoin speed-up over naive loop: ~{speedup:,.0f}x")

counts = joined.groupby("cell_id").size().rename("n_points")
grid = grid.join(counts, on="cell_id").fillna({"n_points": 0}).astype({"n_points": int})
print(grid.sort_values("n_points", ascending=False).head(5)[["cell_id", "region", "n_points"]].to_string())

# --- Step 3: buffers ----------------------------------------------------------------
section("buffers")
subset = points.iloc[:50_000]
with timer("buffer 500 m, quad_segs=16 (default)"):
    b16 = subset.buffer(500)
with timer("buffer 500 m, quad_segs=4"):
    b4 = subset.buffer(500, resolution=4)
print(f"vertices per buffer: {shapely.get_num_coordinates(b16.values.data[0])} vs "
      f"{shapely.get_num_coordinates(b4.values.data[0])}; area error of quad_segs=4: "
      f"{(1 - b4.area.mean() / (np.pi * 500 ** 2)) * 100:.2f}%")

# --- Step 4: overlay & clip -----------------------------------------------------------
section("overlay")
burglary = points[points.category == "burglary"].iloc[:20_000]
with timer("buffer + unary_union (20,000 x 1 km)"):
    zones = gpd.GeoDataFrame(geometry=[burglary.buffer(1_000, resolution=8).unary_union], crs=27700)
    zones = zones.explode(index_parts=False).reset_index(drop=True)
print(f"merged into {len(zones):,} hotspot polygons")
with timer("overlay(grid, zones, 'intersection')"):
    inter = gpd.overlay(grid[["cell_id", "region", "geometry"]], zones, how="intersection")
inter["area_km2"] = inter.area / 1e6
cover = inter.groupby("cell_id").area_km2.sum()
print(f"{len(inter):,} intersection pieces; {(cover > 12.5).sum():,} cells are more than half covered")

# --- Step 5: dissolve ---------------------------------------------------------------
section("dissolve")
with timer("dissolve(by='region')  [unary_union]"):
    regions = grid.dissolve(by="region", aggfunc={"n_points": "sum"})
with timer("coverage_union_all per region"):
    cov = {r: shapely.coverage_union_all(g.geometry.values.data) for r, g in grid.groupby("region")}
diff = max(abs(regions.loc[r].geometry.area - g.area) for r, g in cov.items())
print(f"{len(regions)} regions; max area difference between methods: {diff:.6f} m²")

# --- Step 6: file formats -------------------------------------------------------------
section("file formats")
formats = {
    "Shapefile": ("points.shp", "ESRI Shapefile"),
    "GeoPackage": ("points.gpkg", "GPKG"),
    "FlatGeobuf": ("points.fgb", "FlatGeobuf"),
    "GeoParquet": ("points.parquet", None),
}
query_box = box(420_000, 420_000, 450_000, 450_000)  # 30 x 30 km around Sheffield/Manchester
rows = []
for name, (fname, driver) in formats.items():
    path = OUT / fname
    for p in OUT.glob(path.stem + ".*"):
        if p.suffix in {".shp", ".shx", ".dbf", ".prj", ".cpg", ".gpkg", ".fgb", ".parquet"}:
            p.unlink()
    t0 = time.perf_counter()
    if driver:
        points.to_file(path, driver=driver, engine="pyogrio")
    else:
        points.to_parquet(path)
    t_write = time.perf_counter() - t0

    t0 = time.perf_counter()
    back = gpd.read_file(path, engine="pyogrio") if driver else gpd.read_parquet(path)
    t_read = time.perf_counter() - t0
    assert len(back) == len(points)

    t0 = time.perf_counter()
    if driver:
        sub = gpd.read_file(path, engine="pyogrio", bbox=query_box.bounds)
    else:  # geopandas 0.13 has no bbox filter for Parquet: read, then filter
        sub = gpd.read_parquet(path)
        sub = sub[sub.intersects(query_box)]
    t_bbox = time.perf_counter() - t0

    size = sum(p.stat().st_size for p in OUT.glob(path.stem + ".*")
               if p.suffix in {".shp", ".shx", ".dbf", ".prj", ".cpg", ".gpkg", ".fgb", ".parquet"})
    rows.append({"format": name, "size_MB": size / 1e6, "write_s": t_write, "read_s": t_read,
                 "bbox_read_s": t_bbox, "bbox_rows": len(sub)})

fmt = pd.DataFrame(rows).set_index("format").round({"size_MB": 1, "write_s": 2, "read_s": 2, "bbox_read_s": 3})
print(fmt.to_string())
fmt.to_html(OUT / "formats_table.html", border=0)

# --- Step 7: figures --------------------------------------------------------------------
fig, ax = plt.subplots(figsize=(7.5, 2.8))
labels = ["naive loop (extrapolated)", "STRtree.query, all points", "gpd.sjoin, all points"]
vals = [timings[k] for k in labels]
bars = ax.barh(["Naive loop\n(extrapolated)", "shapely.STRtree\n.query()", "geopandas\n.sjoin()"], vals,
               color=["#adb5bd", "#2f9e44", "#1d6fa3"])
ax.set_xscale("log")
ax.invert_yaxis()
for b, v in zip(bars, vals):
    ax.text(v * 1.15, b.get_y() + b.get_height() / 2, f"{v:,.2f} s", va="center")
ax.set_xlim(min(vals) / 3, max(vals) * 8)
ax.set_xlabel("Seconds (log scale)")
ax.set_title(f"Point-in-polygon: {len(points):,} points x {len(grid):,} grid cells")
fig.savefig(OUT / "01-pip-timings.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(6, 7.5))
grid.plot(column="n_points", ax=ax, cmap="magma_r", scheme="NaturalBreaks", k=7, linewidth=0,
          legend=True, legend_kwds={"title": "Points per 5 km cell", "loc": "upper left", "fontsize": 7, "fmt": "{:.0f}"})
gb.boundary.plot(ax=ax, color="#495057", linewidth=0.4)
ax.set_axis_off()
ax.set_title(f"{len(joined):,} points aggregated to 5 km cells with gpd.sjoin")
fig.savefig(OUT / "02-grid-counts.png")
plt.close(fig)

fig, ax = plt.subplots(figsize=(7, 6))
view = box(380_000, 360_000, 480_000, 460_000)
g_view = grid[grid.intersects(view)]
g_view.boundary.plot(ax=ax, color="#ced4da", linewidth=0.5)
inter[inter.intersects(view)].plot(ax=ax, column="area_km2", cmap="viridis", linewidth=0.2, edgecolor="white",
                                   legend=True, legend_kwds={"label": "Intersected area per cell (km²)",
                                                             "shrink": 0.7})
burglary[burglary.intersects(view)].plot(ax=ax, color="#e8590c", markersize=1)
ax.set_xlim(view.bounds[0], view.bounds[2])
ax.set_ylim(view.bounds[1], view.bounds[3])
ax.set_axis_off()
ax.set_title("overlay(): 1 km hotspot zones cut by the 5 km grid\n(Manchester, Leeds/Bradford and Sheffield)")
fig.savefig(OUT / "03-overlay.png")
plt.close(fig)

fig, axes = plt.subplots(1, 3, figsize=(11, 3))
for ax, col, title in zip(axes, ["size_MB", "write_s", "read_s"], ["File size (MB)", "Write time (s)",
                                                                   "Full read time (s)"]):
    ax.bar(fmt.index, fmt[col], color=["#adb5bd", "#1d6fa3", "#2f9e44", "#e8590c"])
    ax.set_title(title)
    ax.tick_params(axis="x", rotation=30)
    ax.grid(axis="y", alpha=0.3)
fig.suptitle(f"{len(points):,} points with 3 attribute columns", y=1.03)
fig.savefig(OUT / "04-formats.png")
plt.close(fig)
print("\nfigures written to", OUT)