# /// script
# requires-python = ">=3.10"
# dependencies = ["shapely>=2", "pyproj", "numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Filter and refine, and the bounding-box false positive.

100,000 shapes in a projected plane (metres): 50,000 compact circles and
50,000 long diagonal lines buffered by 20 m. For 1,000 random query
rectangles: matches by bounding box alone and after the exact intersects test,
by shape kind and fill ratio (area / bounding-box area). Time with an STRtree
(filter only, filter and refine) against brute-force exact tests. Then the
distance trap: two points 1 km apart measured in raw degrees at latitude 0
and 60, against the geodesic distance.
"""
import argparse, json, time
import numpy as np, shapely
from shapely import STRtree, box
from pyproj import Geod

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--shapes", type=int, default=100_000, help="Shapes (half circles, half lines).")
p.add_argument("--queries", type=int, default=1000, help="Query rectangles.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

rng = np.random.default_rng(args.seed); half = args.shapes // 2; W = 100_000
cx, cy = rng.uniform(0, W, half), rng.uniform(0, W, half)
circles = shapely.buffer(shapely.points(cx, cy), rng.uniform(50, 300, half))
x0, y0, L = rng.uniform(0, W, half), rng.uniform(0, W, half), rng.uniform(2000, 8000, half)
lines = shapely.buffer(shapely.linestrings(np.stack([np.stack([x0, y0], 1), np.stack([x0 + L / 1.414, y0 + L / 1.414], 1)], 1)), 20)
geoms = np.concatenate([circles, lines]); kind = np.array(["circle"] * half + ["line"] * half)
fill = shapely.area(geoms) / shapely.area(shapely.envelope(geoms))
qx, qy = rng.uniform(0, W, args.queries), rng.uniform(0, W, args.queries); qs = [box(a, b, a + 1000, b + 1000) for a, b in zip(qx, qy)]
tree = STRtree(geoms)
t = time.perf_counter(); cand = [tree.query(q) for q in qs]; t_filter = time.perf_counter() - t
t = time.perf_counter(); exact = [c[shapely.intersects(geoms[c], q)] for c, q in zip(cand, qs)]; t_refine = time.perf_counter() - t + t_filter
t = time.perf_counter(); brute = [np.flatnonzero(shapely.intersects(geoms, q)) for q in qs[:50]]; t_brute = (time.perf_counter() - t) * args.queries / 50
assert all(set(a) == set(b) for a, b in zip(exact[:50], brute))
rows = []
for k in ("circle", "line"):
    bb = sum(int((kind[c] == k).sum()) for c in cand); ex = sum(int((kind[e] == k).sum()) for e in exact)
    rows.append({"shapes": k, "fill_ratio": float(fill[kind == k].mean()), "bbox_matches": bb, "exact_matches": ex, "false_positive_pct": (1 - ex / bb) * 100 if bb else None})
timing = [{"method": "STRtree, bounding boxes only", "seconds": t_filter}, {"method": "STRtree, filter and refine", "seconds": t_refine},
          {"method": "brute-force exact test (extrapolated from 50 queries)", "seconds": t_brute}]
g = Geod(ellps="WGS84"); trap = []
for lat in (0.0, 60.0):
    lon2, lat2, _ = g.fwd(0.0, lat, 90, 1000)                                # a point 1 km due east
    raw_deg = float(np.hypot(lon2 - 0.0, lat2 - lat)); naive_km = raw_deg * 111.32
    _, _, true_m = g.inv(0.0, lat, lon2, lat2)
    trap.append({"latitude": lat, "raw_degrees": raw_deg, "naive_km": naive_km, "geodesic_km": true_m / 1000})
for r in rows: print(f"{r['shapes']:<7} fill ratio {r['fill_ratio']:.3f}: bounding-box matches {r['bbox_matches']:>7,}, exact {r['exact_matches']:>7,}, false positives {r['false_positive_pct']:5.1f}%")
for r in timing: print(f"{r['method']:<54} {r['seconds']:7.2f} s for {args.queries:,} queries")
for r in trap: print(f"latitude {r['latitude']:>4.0f}: 1 km east is {r['raw_degrees']:.5f} raw degrees; × 111.32 km = {r['naive_km']:.3f} km; geodesic {r['geodesic_km']:.3f} km")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"shapes": rows, "timing": timing, "distance": trap}}, f)