#!/usr/bin/env python3
"""
Map-Matching Sandbox

This script tests the new walkingpandas map-matching architecture step-by-step.
It loads a raw GPS file, downloads the exact OSM topology, runs the routing engine
(Phase 1: Trace), snaps to the graph (Phase 2: Snap), enriches with attributes
(Phase 3: Enrich), and exports geometries to GeoPackage files for QGIS inspection.
"""

import pandas as pd
import geopandas as gpd
from shapely.geometry import LineString
from pathlib import Path

import walkingpandas as wp
from walkingpandas.io.match import ValhallaMatcher, download_region, decode_valhalla


def prepare_raw_data(csv_path: str) -> pd.DataFrame:
    """Loads the raw GPS CSV and adds required placeholder columns if missing."""
    print(f"--- 1. Loading Raw Data: {csv_path} ---")
    df = pd.read_csv(csv_path)
    
    # Keep only the rows with valid coordinates
    df = df.dropna(subset=["latitude", "longitude"]).copy()
    
    # The new pipeline expects standard column names: trip_id, timestamp, lon, lat
    df = df.rename(columns={"longitude": "lon", "latitude": "lat"})
    
    # Add dummy trip and time if they don't exist in this raw file
    if "trip_id" not in df.columns:
        df["trip_id"] = "trip_1"
    if "timestamp" not in df.columns:
        df["timestamp"] = pd.date_range("2018-11-05 08:00:00", periods=len(df), freq="s")
        
    print(f"Loaded {len(df)} GPS points.")
    return df


def get_padded_bbox(df: pd.DataFrame, buffer_deg: float = 0.005) -> tuple:
    """Calculates a bounding box from the dataframe and adds a small buffer."""
    north = df["lat"].max() + buffer_deg
    south = df["lat"].min() - buffer_deg
    east = df["lon"].max() + buffer_deg
    west = df["lon"].min() - buffer_deg
    return (north, south, east, west)


def main():
    # Setup paths
    print("--- Start Testing ---")
    # gps_file = 'gps/example2/2018-11-05_sgv_23_peu_1_processed.csv'
    gps_file = "gps/example2/2018-11-05_sgv_0101_PEU_processed.csv"
    if not Path(gps_file).exists():
        print(f"File not found: {gps_file}")
        print("Please run this script from the root of the walkingpandas repository.")
        return

    # 1. Prepare Data & Region
    gps_df = prepare_raw_data(gps_file)
    bbox = get_padded_bbox(gps_df)
    print(f"Region Bounding Box (N, S, E, W): {bbox}")

    # 2. Download the Single Source of Truth
    print("\n--- 2. Downloading OSM Data ---")
    osm_xml_path = "sandbox_map.osm"
    osm_pbf_path = "sandbox_map.osm.pbf"
    
    if Path(osm_xml_path).exists() and Path(osm_pbf_path).exists():
        print(f"OSM data already exists locally ({osm_xml_path}, {osm_pbf_path}). Skipping download.")
        osm_xml, osm_pbf = osm_xml_path, osm_pbf_path
    else:
        osm_xml, osm_pbf = download_region(
            bbox=bbox, 
            output_xml=osm_xml_path, 
            output_pbf=osm_pbf_path
        )

    # 3. Initialize the Matcher & Network
    print("\n--- 3. Initializing ValhallaMatcher & wp.Network ---")
    # This automatically builds tiles from the .pbf and exposes the C++ actor
    matcher = ValhallaMatcher(osm_pbf=osm_pbf, osm_xml=osm_xml, costing="pedestrian")
    
    # Build the strict walkingpandas topological graph from the .osm XML
    network = matcher.get_network()
    nodes_count = len(network.get_nodes(as_geopandas=False))
    edges_count = len(network.get_edges(as_geopandas=False))
    print(f"Network synchronized: {nodes_count} nodes, {edges_count} edges.")

    # 4. Phase 1: TRACE
    print("\n--- 4. Phase 1: Trace (Valhalla Engine) ---")
    raw_traces = matcher.trace(gps_df)
    
    if raw_traces.empty:
        print("Valhalla failed to route this trajectory.")
        return
        
    trace = raw_traces.iloc[0]
    print(f"Valhalla returned {len(trace['valhalla_edges'])} candidate edges.")

    # Save the Raw Valhalla Shape for QGIS debugging
    decoded_coords = decode_valhalla(trace["raw_shape"])
    path_geom = LineString(decoded_coords)
    raw_gdf = gpd.GeoDataFrame([{"name": "Valhalla Raw Shape"}], geometry=[path_geom], crs="EPSG:4326")
    raw_gdf.to_file("sandbox_raw_shape.gpkg", layer="raw_shape", driver="GPKG")
    print("-> Saved raw engine geometry to 'sandbox_raw_shape.gpkg'")

    # 5. Phase 2: SNAP (Topological Alignment)
    print("\n--- 5. Phase 2: Snap (Topological Alignment) ---")
    snapped = matcher.snap(raw_traces, network)

    print(f"Topological snapping successful: {len(snapped)} walk rows.")

    # 6. Phase 3: ENRICH (Attribute Projection)
    print("\n--- 6. Phase 3: Enrich (Attribute Projection) ---")
    final_walk = matcher.enrich(
        snapped_df=snapped,
        raw_traces=raw_traces,
        node_attributes=["timestamp"],
    )
    # Export topology to GPKG for visual verification of edge_ids
    edges_gdf = network.get_edges(as_geopandas=True)
    walked_edge_ids = final_walk["exit_edge"].dropna().unique()
    matched_edges_gdf = edges_gdf[edges_gdf["edge_id"].isin(walked_edge_ids)]
    print(f"Extracted {len(matched_edges_gdf)} physical edge geometries from the network to save.")
    matched_edges_gdf.to_file("sandbox_topology_check.gpkg", layer="matched_edges", driver="GPKG")
    print("Final enriched walk frame:")
    print(final_walk)
    
    print("\nDone! Verify edges in 'sandbox_matched_topology.gpkg' then check 'final_walk' dataframe.")


if __name__ == "__main__":
    main()