"""
Shared data helpers for tutorial examples.

Provides small, reproducible networks and walk DataFrames so each
example can stay self-contained while avoiding copy-paste.
"""

import tempfile
from pathlib import Path

import pandas as pd
import geopandas as gpd
from shapely.geometry import Point, LineString


def make_simple_network_gdf():
    """
    Create a small network as (nodes_gdf, edges_gdf) in the unit square.
    Nodes A(0,0), B(1,0), C(1,1), D(0,1); edges AB, BC, CD, DA, BE (E at 2,0).
    """
    nodes_data = [
        {"node_id": 0, "x": 0, "y": 0},
        {"node_id": 1, "x": 1, "y": 0},
        {"node_id": 2, "x": 1, "y": 1},
        {"node_id": 3, "x": 0, "y": 1},
        {"node_id": 4, "x": 2, "y": 0},
    ]
    nodes_df = pd.DataFrame(nodes_data)
    nodes_gdf = gpd.GeoDataFrame(
        nodes_df,
        geometry=[Point(row["x"], row["y"]) for _, row in nodes_df.iterrows()],
    )

    edges_data = [
        {"edge_id": 0, "source_node": 0, "target_node": 1},
        {"edge_id": 1, "source_node": 1, "target_node": 2},
        {"edge_id": 2, "source_node": 2, "target_node": 3},
        {"edge_id": 3, "source_node": 3, "target_node": 0},
        {"edge_id": 4, "source_node": 1, "target_node": 4},
    ]
    edges_df = pd.DataFrame(edges_data)

    def line(row):
        s = nodes_gdf[nodes_gdf["node_id"] == row["source_node"]].iloc[0]
        t = nodes_gdf[nodes_gdf["node_id"] == row["target_node"]].iloc[0]
        return LineString([(s.geometry.x, s.geometry.y), (t.geometry.x, t.geometry.y)])

    edges_gdf = gpd.GeoDataFrame(
        edges_df,
        geometry=[line(row) for _, row in edges_df.iterrows()],
    )
    return nodes_gdf, edges_gdf


def make_simple_walks_df():
    """
    Create a small walks DataFrame in superset schema.
    walk1: A -> B -> C; walk2: B -> E; walk3 (cycle): A -> B -> C -> D -> A.
    """
    return pd.DataFrame([
        {"walk_id": "walk1", "sequence_idx": 0, "node_id": 0, "exit_edge": 0},
        {"walk_id": "walk1", "sequence_idx": 1, "node_id": 1, "exit_edge": 1},
        {"walk_id": "walk1", "sequence_idx": 2, "node_id": 2, "exit_edge": None},
        {"walk_id": "walk2", "sequence_idx": 0, "node_id": 1, "exit_edge": 4},
        {"walk_id": "walk2", "sequence_idx": 1, "node_id": 4, "exit_edge": None},
        {"walk_id": "walk3", "sequence_idx": 0, "node_id": 0, "exit_edge": 0},
        {"walk_id": "walk3", "sequence_idx": 1, "node_id": 1, "exit_edge": 1},
        {"walk_id": "walk3", "sequence_idx": 2, "node_id": 2, "exit_edge": 2},
        {"walk_id": "walk3", "sequence_idx": 3, "node_id": 3, "exit_edge": 3},
        {"walk_id": "walk3", "sequence_idx": 4, "node_id": 0, "exit_edge": None},
    ])


def make_walks_with_timestamps():
    """Walks in superset schema with timestamp (and optional dwell_time)."""
    from datetime import datetime, timedelta
    base = datetime(2024, 1, 1, 8, 0, 0)
    return pd.DataFrame([
        {"walk_id": "w1", "sequence_idx": 0, "node_id": 0, "exit_edge": 0, "timestamp": base},
        {"walk_id": "w1", "sequence_idx": 1, "node_id": 1, "exit_edge": 1, "timestamp": base + timedelta(minutes=5)},
        {"walk_id": "w1", "sequence_idx": 2, "node_id": 2, "exit_edge": None, "timestamp": base + timedelta(minutes=10)},
        {"walk_id": "w2", "sequence_idx": 0, "node_id": 1, "exit_edge": 4, "timestamp": base + timedelta(minutes=15)},
        {"walk_id": "w2", "sequence_idx": 1, "node_id": 4, "exit_edge": None, "timestamp": base + timedelta(minutes=20)},
    ])


def make_string_id_network_gdf():
    """
    Same topology as make_simple_network_gdf() but with human-readable
    string IDs: nodes Station_A … Station_E, edges link_AB … link_BE.
    """
    nodes_data = [
        {"node_id": "Station_A", "x": 0, "y": 0},
        {"node_id": "Station_B", "x": 1, "y": 0},
        {"node_id": "Station_C", "x": 1, "y": 1},
        {"node_id": "Station_D", "x": 0, "y": 1},
        {"node_id": "Station_E", "x": 2, "y": 0},
    ]
    nodes_df = pd.DataFrame(nodes_data)
    nodes_gdf = gpd.GeoDataFrame(
        nodes_df,
        geometry=[Point(row["x"], row["y"]) for _, row in nodes_df.iterrows()],
    )

    edges_data = [
        {"edge_id": "link_AB", "source_node": "Station_A", "target_node": "Station_B"},
        {"edge_id": "link_BC", "source_node": "Station_B", "target_node": "Station_C"},
        {"edge_id": "link_CD", "source_node": "Station_C", "target_node": "Station_D"},
        {"edge_id": "link_DA", "source_node": "Station_D", "target_node": "Station_A"},
        {"edge_id": "link_BE", "source_node": "Station_B", "target_node": "Station_E"},
    ]
    edges_df = pd.DataFrame(edges_data)

    def line(row):
        s = nodes_gdf[nodes_gdf["node_id"] == row["source_node"]].iloc[0]
        t = nodes_gdf[nodes_gdf["node_id"] == row["target_node"]].iloc[0]
        return LineString([(s.geometry.x, s.geometry.y), (t.geometry.x, t.geometry.y)])

    edges_gdf = gpd.GeoDataFrame(
        edges_df,
        geometry=[line(row) for _, row in edges_df.iterrows()],
    )
    return nodes_gdf, edges_gdf


def make_string_id_walks_df():
    """
    Walks using the string-ID network: Station_A, link_AB, etc.
    walk1: A→B→C; walk2: B→E; walk3 (cycle): A→B→C→D→A.
    """
    return pd.DataFrame([
        {"walk_id": "walk1", "sequence_idx": 0, "node_id": "Station_A", "exit_edge": "link_AB"},
        {"walk_id": "walk1", "sequence_idx": 1, "node_id": "Station_B", "exit_edge": "link_BC"},
        {"walk_id": "walk1", "sequence_idx": 2, "node_id": "Station_C", "exit_edge": None},
        {"walk_id": "walk2", "sequence_idx": 0, "node_id": "Station_B", "exit_edge": "link_BE"},
        {"walk_id": "walk2", "sequence_idx": 1, "node_id": "Station_E", "exit_edge": None},
        {"walk_id": "walk3", "sequence_idx": 0, "node_id": "Station_A", "exit_edge": "link_AB"},
        {"walk_id": "walk3", "sequence_idx": 1, "node_id": "Station_B", "exit_edge": "link_BC"},
        {"walk_id": "walk3", "sequence_idx": 2, "node_id": "Station_C", "exit_edge": "link_CD"},
        {"walk_id": "walk3", "sequence_idx": 3, "node_id": "Station_D", "exit_edge": "link_DA"},
        {"walk_id": "walk3", "sequence_idx": 4, "node_id": "Station_A", "exit_edge": None},
    ])


def write_to_temp_parquet(nodes_gdf, edges_gdf, walks_df):
    """
    Write nodes, edges, and walks to a temporary directory as Parquet.
    Returns the Path to the temp dir (caller should use as context manager or delete).
    """
    temp_path = Path(tempfile.mkdtemp())
    nodes_gdf.to_parquet(temp_path / "nodes.parquet", index=False)
    edges_gdf.to_parquet(temp_path / "edges.parquet", index=False)
    walks_df.to_parquet(temp_path / "walks.parquet", index=False)
    return temp_path
