"""
05 — Spatial filtering: passing_through, within_bounds, intersecting, near_point.

This example uses a network with known coordinates (unit square). It demonstrates
passing_through (node IDs), within_bounds (bbox), intersecting (Shapely polygon),
and near_point (Point + distance). Reverse spatial optimization: query network
first, then IN on Parquet.
"""

import shutil
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import walkingpandas as wp
from shapely.geometry import box, Point

from _data_helpers import (
    make_simple_network_gdf,
    make_simple_walks_df,
    make_string_id_network_gdf,
    make_string_id_walks_df,
    write_to_temp_parquet,
)


def main():
    print("=" * 60)
    print("05 — Spatial filtering")
    print("=" * 60)

    nodes_gdf, edges_gdf = make_simple_network_gdf()
    walks_df = make_simple_walks_df()
    temp_dir = write_to_temp_parquet(nodes_gdf, edges_gdf, walks_df)

    try:
        network = wp.Network(nodes=temp_dir / "nodes.parquet", edges=temp_dir / "edges.parquet")
        walks = wp.WalkFrame(temp_dir / "walks.parquet", network=network)

        # passing_through: filter by one or several node IDs
        through_c = walks.passing_through(2).compute()
        print("\n1. passing_through('C') — walks that visit node C")
        print(through_c[["walk_id", "node_id"]].drop_duplicates().to_string(index=False))

        through_b_e = walks.passing_through([1, 4]).compute()
        print("\n2. passing_through(['B', 'E']) — walks visiting B or E")
        print(through_b_e[["walk_id", "node_id"]].drop_duplicates().to_string(index=False))

        # within_bounds: bbox (minx, miny, maxx, maxy)
        bbox = (0.5, -0.1, 1.5, 0.5)
        in_bbox = walks.within_bounds(bbox).compute()
        print(f"\n3. within_bounds({bbox}) — steps inside bbox")
        print(in_bbox[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # intersecting: Shapely polygon (e.g. box)
        poly = box(0.5, 0.5, 2, 1.5)
        inter = walks.intersecting(poly).compute()
        print("\n4. intersecting(box(0.5, 0.5, 2, 1.5)) — steps in polygon")
        print(inter[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # near_point: Point + distance (buffer)
        pt = Point(1, 0)
        near = walks.near_point(pt, distance=0.5).compute()
        print("\n5. near_point(Point(1,0), distance=0.5) — steps near (1,0)")
        print(near[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # --- String IDs: passing_through with station names ---
        print("\n--- String-ID spatial filtering ---")
        str_nodes, str_edges = make_string_id_network_gdf()
        str_walks = make_string_id_walks_df()

        str_net = wp.Network()
        str_net.add_nodes(str_nodes)
        str_net.add_edges(str_edges)
        str_wf = wp.WalkFrame.from_dataframe(str_walks, network=str_net)

        result = str_wf.passing_through("Station_C").compute()
        print("6. passing_through('Station_C') with string IDs")
        print(result[["walk_id", "node_id"]].drop_duplicates().to_string(index=False))

        result2 = str_wf.passing_through(["Station_B", "Station_E"]).compute()
        print("\n7. passing_through(['Station_B', 'Station_E'])")
        print(result2[["walk_id", "node_id"]].drop_duplicates().to_string(index=False))

        print("\nDone.")
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


if __name__ == "__main__":
    main()
