Using STR-packed R-trees for fast trajectory lookups

import numpy as np
import geopandas as gpd
from shapely import STRtree, box
from shapely.geometry.base import BaseGeometry


class TripIndex:
    """
    A packed R-tree over trip envelopes, with the exact-predicate step.

    The tree answers "which trips MIGHT intersect this?" in microseconds.
    The exact test then runs on tens of candidates rather than millions of
    rows — that two-stage shape is the entire point.
    """

    def __init__(self, trips: gpd.GeoDataFrame, geom_col: str = "geometry"):
        if trips.empty:
            raise ValueError("Cannot index an empty trip set.")
        if trips.crs is None or trips.crs.is_geographic:
            raise ValueError(
                f"trips must be in a PROJECTED metric CRS; got {trips.crs}. "
                "A tree built in degrees returns candidates whose distances "
                "are latitude-dependent and cannot be compared to metres."
            )
        if trips[geom_col].isna().any():
            raise ValueError("Null geometries cannot be indexed; drop them first.")

        self.trips = trips.reset_index(drop=False)   # keep the original ids
        self.crs = trips.crs
        # STRtree packs at construction — there is no incremental insert.
        self.tree = STRtree(self.trips[geom_col].values)

    def candidates(self, query: BaseGeometry, buffer_m: float = 0.0) -> np.ndarray:
        """Positional indices whose ENVELOPE intersects the query envelope."""
        g = query.buffer(buffer_m) if buffer_m > 0 else query
        return self.tree.query(g)

    def intersects(self, query: BaseGeometry, buffer_m: float = 0.0) -> gpd.GeoDataFrame:
        """Rows that actually intersect — coarse filter, then exact predicate."""
        idx = self.candidates(query, buffer_m)
        if idx.size == 0:
            return self.trips.iloc[[]]
        cand = self.trips.iloc[idx]
        g = query.buffer(buffer_m) if buffer_m > 0 else query
        # The exact test is the expensive one; it now runs on len(idx) rows.
        return cand[cand.geometry.intersects(g)]

    def in_window(self, bbox: tuple, t0, t1,
                  start_col: str = "t_start", end_col: str = "t_end") -> gpd.GeoDataFrame:
        """
        Space AND time. The tree only knows geometry, so the temporal
        predicate is applied to the candidate set — which is exactly the
        right order, because the spatial filter is the selective one here.
        """
        for c in (start_col, end_col):
            if c not in self.trips.columns:
                raise ValueError(f"Missing temporal column: {c}")
        hits = self.intersects(box(*bbox))
        overlaps = (hits[start_col] < t1) & (hits[end_col] > t0)
        return hits[overlaps]


def build_and_time(trips: gpd.GeoDataFrame) -> dict:
    """Report build cost and tree quality so a rebuild schedule can be set."""
    import time
    t0 = time.perf_counter()
    idx = TripIndex(trips)
    build_s = time.perf_counter() - t0

    # Mean envelope area relative to the total extent is a proxy for how
    # selective the tree can be: long diagonal trips have huge envelopes and
    # a tree over them prunes poorly no matter how well it is packed.
    env = trips.geometry.envelope.area
    total = trips.total_bounds
    extent = (total[2] - total[0]) * (total[3] - total[1])
    return {
        "n": len(trips),
        "build_s": round(build_s, 3),
        "mean_envelope_frac": float((env / max(extent, 1e-9)).mean()),
        "index": idx,
    }

def validate_tree(stats: dict, sample_query, expected_max_candidates: int = 500) -> None:
    """Three checks that separate a useful tree from an expensive one."""
    idx = stats["index"]

    # 1. Envelope bloat. If the average trip envelope covers a large share of
    #    the extent, every query hits most of the tree and the index cannot
    #    help — split long trips into sub-trips before indexing.
    assert stats["mean_envelope_frac"] < 0.05, (
        f"mean envelope is {stats['mean_envelope_frac']:.1%} of the extent — "
        "long diagonal trips defeat any R-tree; index sub-trips instead"
    )

    # 2. Candidate count. A coarse filter that returns thousands has not
    #    filtered; the exact predicate then dominates the runtime.
    n_cand = len(idx.candidates(sample_query))
    assert n_cand <= expected_max_candidates, (
        f"{n_cand} candidates for one query — check the envelope bloat above"
    )

    # 3. Precision of the coarse filter: how many candidates survive the
    #    exact test. Below ~10% means the envelopes are a poor proxy.
    n_exact = len(idx.intersects(sample_query))
    precision = n_exact / max(n_cand, 1)
    print(f"OK — build {stats['build_s']}s, {n_cand} candidates, "
          f"{precision:.0%} survive the exact test")