"""
06 — Temporal queries: filter(time_range=...), edge_traffic_at(time).

This example creates walk data with timestamp (and optionally dwell_time) in
superset schema, writes to Parquet, then demonstrates filter(time_range=(start, end))
with datetime or time objects, and edge_traffic_at(time) to see which edges were
active at a given time (e.g. traffic_count per edge).
"""

import shutil
import sys
from datetime import datetime, time
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_walks_with_timestamps,
    write_to_temp_parquet,
)


def main():
    print("=" * 60)
    print("06 — Temporal queries")
    print("=" * 60)

    nodes_gdf, edges_gdf = make_simple_network_gdf()
    walks_df = make_walks_with_timestamps()
    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 time range (datetime)
        start = datetime(2024, 1, 1, 8, 0, 0)
        end = datetime(2024, 1, 1, 8, 12, 0)
        in_range = walks.filter(time_range=(start, end)).compute()
        print("\n1. filter(time_range=(start, end)) with datetime")
        print(in_range[["walk_id", "sequence_idx", "node_id", "timestamp"]].to_string(index=False))

        # filter by time range (time objects — use current date in backend)
        in_range_time = walks.filter(time_range=(time(8, 0), time(8, 15))).compute()
        print("\n2. filter(time_range=(time(8,0), time(8,15)))")
        print(in_range_time[["walk_id", "sequence_idx", "timestamp"]].to_string(index=False))

        # edge_traffic_at(time): which edges were active at time T
        t = datetime(2024, 1, 1, 8, 7, 0)
        traffic = walks.edge_traffic_at(t).compute()
        print("\n3. edge_traffic_at(datetime(2024,1,1,8,7,0)) — traffic per edge at that time")
        print(traffic.to_string(index=False))

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


if __name__ == "__main__":
    main()
