"""
03 — WalkFrame basics: filter(), limit(), to_pandas(), to_geopandas().

This example builds a WalkFrame from temp Parquet and demonstrates filtering
by walk_id or node_id, limiting rows, and exporting with .to_pandas() and
.to_geopandas(). You can also build a WalkFrame with nodes= and edges=
directly; use .network to access the underlying Network.
"""

import shutil
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import walkingpandas as wp

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("03 — WalkFrame basics")
    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)

        # filter by walk_id
        w1 = walks.filter(walk_id="walk1").compute()
        print("\n1. .filter(walk_id='walk1').compute()")
        print(w1[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # filter by node_id (walks that pass through that node)
        through_b = walks.filter(node_id=1).compute()
        print("\n2. .filter(node_id='B').compute() — walks that visit B")
        print(through_b[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # limit
        limited = walks.limit(3).compute()
        print("\n3. .limit(3).compute()")
        print(limited[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # to_pandas() is an alias for compute()
        df = walks.limit(2).to_pandas()
        print(f"\n4. .to_pandas() → same as .compute(), got {len(df)} rows")

        # to_geopandas() runs the query and returns a GeoDataFrame (geometry enrichment can be added later)
        gdf = walks.limit(2).to_geopandas()
        print(f"5. .to_geopandas() → GeoDataFrame with {len(gdf)} rows")

        # Chaining: filter + limit + compute
        chained = walks.filter(walk_id="walk3").limit(2).compute()
        print("\n6. Chained: .filter(walk_id='walk3').limit(2).compute()")
        print(chained[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # --- String IDs: feed strings in, get strings back ---
        print("\n--- String-ID workflow ---")
        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)

        # filter by string node_id
        through_b = str_wf.filter(node_id="Station_B").compute()
        print("7. filter(node_id='Station_B') with string IDs")
        print(through_b[["walk_id", "sequence_idx", "node_id"]].to_string(index=False))

        # compute(decode=False) reveals the raw integers
        raw = str_wf.filter(walk_id="walk1").compute(decode=False)
        print("\n8. compute(decode=False) — raw integers")
        print(raw[["walk_id", "node_id", "exit_edge"]].to_string(index=False))

        # compute(decode=True) (default) returns original strings
        decoded = str_wf.filter(walk_id="walk1").compute(decode=True)
        print("\n9. compute(decode=True) — original string labels")
        print(decoded[["walk_id", "node_id", "exit_edge"]].to_string(index=False))

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


if __name__ == "__main__":
    main()
