#!/usr/bin/env python3
"""
Build Road Network Graph for Wuhan Respiratory Disease Risk Prediction Platform
Extracts Wuhan OSM road network and builds graph structure
"""

import json
import os
import numpy as np
import pandas as pd
import geopandas as gpd
from shapely.geometry import shape, MultiPolygon, Polygon
from scipy.sparse import csr_matrix, lil_matrix
import networkx as nx
import pyrosm
import warnings
warnings.filterwarnings('ignore')

# Paths
PBF_PATH = '/home/akiba/CA/Datas/地图/hubei-260129.osm.pbf'
WUHAN_GEOJSON = '/home/akiba/CA/Datas/武汉市.geojson'
OUTPUT_DIR = '/home/akiba/CA/processed/graph'

def load_wuhan_boundary():
    """Load Wuhan boundary from geojson"""
    with open(WUHAN_GEOJSON, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # Combine all district polygons into one
    geometries = []
    for feat in data['features']:
        geom = shape(feat['geometry'])
        geometries.append(geom)

    # Create union of all geometries
    boundary = geometries[0]
    for g in geometries[1:]:
        boundary = boundary.union(g)

    return boundary, data['features']

def get_district_for_point(point, features):
    """Find which district a point belongs to"""
    for feat in features:
        geom = shape(feat['geometry'])
        if geom.contains(point):
            return feat['properties']['name']
    return 'unknown'

def build_road_graph():
    """Build road network graph from OSM data"""
    print("Loading Wuhan boundary...")
    boundary, district_features = load_wuhan_boundary()
    print(f"  Boundary type: {boundary.geom_type}")

    print("Reading OSM data...")
    # Initialize OSM reader with Wuhan boundary
    print("  Initializing OSM reader...")
    osm = pyrosm.OSM(PBF_PATH, bounding_box=boundary)

    # Get all drivable roads (more comprehensive than just primary/secondary)
    print("Extracting roads within Wuhan boundary...")
    # Filter to Wuhan boundary using bounding box first (faster)
    bounds = boundary.bounds
    print(f"  Bounding box: {bounds}")

    # Read roads using pyrosm with custom filter
    # Get all highways first, then filter to boundary
    print("  Reading highways...")
    highways = osm.get_data_by_custom_criteria({
        'highway': ['motorway', 'trunk', 'primary', 'secondary', 'tertiary',
                    'unclassified', 'residential', 'living_street', 'pedestrian',
                    'track', 'service', 'road']
    })
    print(f"  Total highway elements: {len(highways)}")

    if len(highways) == 0:
        print("ERROR: No highways found. Trying alternative approach...")
        return None

    # Convert to GeoDataFrame
    gdf = gpd.GeoDataFrame(highways, geometry='geometry', crs='EPSG:4326')
    print(f"  GeoDataFrame size: {len(gdf)}")

    # Filter to Wuhan boundary
    print("  Clipping to Wuhan boundary...")
    gdf_clipped = gdf[gdf.geometry.is_valid].copy()
    gdf_clipped = gdf_clipped[gdf_clipped.intersects(boundary)]
    gdf_clipped = gdf_clipped.geometry.apply(lambda g: g.intersection(boundary) if g.is_valid else None)
    gdf_clipped = gdf_clipped.dropna()

    # Explode MultiLineStrings to LineStrings
    def explode_geom(g):
        if g.geom_type == 'MultiLineString':
            return list(g.geoms)
        elif g.geom_type == 'LineString':
            return [g]
        elif g.geom_type == 'MultiPolygon':
            # Get all polygon exteriors as LineStrings
            result = []
            for poly in g.geoms:
                result.append(poly.exterior)
            return result
        elif g.geom_type == 'Polygon':
            # Intersection of a LineString with boundary can return Polygon
            return [g.exterior]
        elif g.geom_type == 'GeometryCollection':
            result = []
            for geom in g.geoms:
                result.extend(explode_geom(geom))
            return result
        return []

    all_geoms = []
    for g in gdf_clipped.geometry:
        all_geoms.extend(explode_geom(g))

    print(f"  Total line segments after clipping: {len(all_geoms)}")

    if len(all_geoms) == 0:
        print("ERROR: No geometries after clipping")
        return None

    # Build graph
    print("Building graph structure...")
    G = nx.MultiDiGraph()

    node_id_counter = 0
    node_info = {}  # osmid -> (lat, lon, district, road_type)

    # First pass: collect all unique points
    all_points = set()
    point_to_node = {}

    for i, geom in enumerate(all_geoms):
        coords = list(geom.coords)
        for coord in coords:
            all_points.add(coord)

    print(f"  Total unique points: {len(all_points)}")

    # Map points to node IDs
    for pt in all_points:
        point_to_node[pt] = node_id_counter
        node_id_counter += 1

    # Add nodes to graph
    for pt, nid in point_to_node.items():
        G.add_node(nid, osmid=nid, x=pt[0], y=pt[1])

    # Second pass: create edges from line segments
    edge_count = 0
    edges_data = []

    for geom in all_geoms:
        coords = list(geom.coords)
        for i in range(len(coords) - 1):
            u = point_to_node[coords[i]]
            v = point_to_node[coords[i+1]]

            # Calculate edge weight (1/length_km)
            dx = coords[i+1][0] - coords[i][0]
            dy = coords[i+1][1] - coords[i][1]
            length_deg = np.sqrt(dx**2 + dy**2)
            # Approximate conversion at Wuhan latitude (30N)
            length_km = length_deg * 111.32 * np.cos(np.radians(30))
            length_km = max(length_km, 0.0001)  # avoid division by zero

            weight = 1.0 / length_km

            G.add_edge(u, v, weight=weight, length=length_km)
            edges_data.append((u, v, length_km, weight))
            edge_count += 1

    print(f"  Graph nodes: {G.number_of_nodes()}")
    print(f"  Graph edges: {G.number_of_edges()}")

    # Check node count and apply fallback if needed
    if G.number_of_nodes() > 70000:
        print("\nNode count exceeds 70k, applying highway filter...")
        # Filter to major roads only
        major_roads = osm.get_data_by_custom_criteria({
            'highway': ['motorway', 'trunk', 'primary', 'secondary', 'tertiary']
        })
        gdf_major = gpd.GeoDataFrame(major_roads, geometry='geometry', crs='EPSG:4326')
        gdf_major = gdf_major[gdf_major.geometry.is_valid].copy()
        gdf_major = gdf_major[gdf_major.intersects(boundary)]

        # Rebuild graph
        G = nx.MultiDiGraph()
        node_id_counter = 0
        point_to_node = {}

        all_geoms = []
        for g in gdf_major.geometry:
            all_geoms.extend(explode_geom(g))

        all_points = set()
        for geom in all_geoms:
            coords = list(geom.coords)
            for coord in coords:
                all_points.add(coord)

        for pt in all_points:
            point_to_node[pt] = node_id_counter
            node_id_counter += 1

        for pt, nid in point_to_node.items():
            G.add_node(nid, osmid=nid, x=pt[0], y=pt[1])

        for geom in all_geoms:
            coords = list(geom.coords)
            for i in range(len(coords) - 1):
                u = point_to_node[coords[i]]
                v = point_to_node[coords[i+1]]
                dx = coords[i+1][0] - coords[i][0]
                dy = coords[i+1][1] - coords[i][1]
                length_deg = np.sqrt(dx**2 + dy**2)
                length_km = length_deg * 111.32 * np.cos(np.radians(30))
                length_km = max(length_km, 0.0001)
                weight = 1.0 / length_km
                G.add_edge(u, v, weight=weight, length=length_km)

        print(f"  Filtered graph nodes: {G.number_of_nodes()}")
        print(f"  Filtered graph edges: {G.number_of_edges()}")

    node_count = G.number_of_nodes()
    if node_count < 15000 or node_count > 70000:
        print(f"WARNING: Node count {node_count} outside target range 15k-70k")

    # Check connectivity
    print("\nChecking graph connectivity...")
    if G.number_of_nodes() > 0:
        # Get largest weakly connected component
        if G.is_directed():
            connected = list(nx.weakly_connected_components(G))
        else:
            connected = list(nx.connected_components(G))
        largest_cc = max(connected, key=len)
        print(f"  Total components: {len(connected)}")
        print(f"  Largest component size: {len(largest_cc)}")
        print(f"  Largest component ratio: {len(largest_cc)/G.number_of_nodes():.2%}")

        # Keep only largest component
        nodes_to_remove = set(G.nodes()) - set(largest_cc)
        G.remove_nodes_from(nodes_to_remove)
        print(f"  After pruning to largest CC: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")

        # Relabel nodes to consecutive integers 0..n-1 for adjacency matrix
        old_nodes = list(G.nodes())
        new_nodes = range(len(old_nodes))
        mapping = dict(zip(old_nodes, new_nodes))
        G = nx.relabel_nodes(G, mapping, copy=False)
        print(f"  Relabeled nodes to consecutive IDs 0..{G.number_of_nodes()-1}")

    # Build output files
    print("\nGenerating output files...")

    # 1. Node metadata
    node_data = []
    for nid in G.nodes():
        props = G.nodes[nid]
        # Approximate lat/lon
        lat = props.get('y', 0)
        lon = props.get('x', 0)
        node_data.append({
            'osmid': nid,
            'lat': lat,
            'lon': lon,
            'district': 'unknown',  # Would need reverse geocoding
            'road_type': 'unknown'
        })

    node_df = pd.DataFrame(node_data)
    node_df.to_parquet(f'{OUTPUT_DIR}/node_metadata.parquet', index=False)
    print(f"  Saved node_metadata.parquet: {len(node_df)} nodes")

    # 2. Edge list
    edge_data = []
    for u, v, data in G.edges(data=True):
        edge_data.append({
            'source': u,
            'target': v,
            'weight': data.get('weight', 1.0),
            'length_km': data.get('length', 0)
        })

    edge_df = pd.DataFrame(edge_data)
    edge_df.to_csv(f'{OUTPUT_DIR}/edge_list.csv', index=False)
    print(f"  Saved edge_list.csv: {len(edge_df)} edges")

    # 3. Adjacency matrix (sparse CSR)
    print("  Building adjacency matrix...")
    n = G.number_of_nodes()
    adj = lil_matrix((n, n), dtype=np.float32)

    for u, v, data in G.edges(data=True):
        adj[u, v] = data.get('weight', 1.0)
        # Make it symmetric for undirected use
        adj[v, u] = data.get('weight', 1.0)

    adj_csr = adj.tocsr()
    np.savez(f'{OUTPUT_DIR}/adjacency_matrix.npz', data=adj_csr.data, indices=adj_csr.indices, indptr=adj_csr.indptr, shape=adj_csr.shape)
    print(f"  Saved adjacency_matrix.npz: {adj_csr.shape}")

    # Verify outputs
    print("\n=== VERIFICATION ===")
    print(f"Node count: {G.number_of_nodes()}")
    print(f"Edge count: {G.number_of_edges()}")
    print(f"Target range: 15,000 - 70,000")

    # Check components
    if G.number_of_nodes() > 0:
        if G.is_directed():
            components = list(nx.weakly_connected_components(G))
        else:
            components = list(nx.connected_components(G))
        print(f"Connected components: {len(components)}")

    # Verify files exist
    for fname in ['adjacency_matrix.npz', 'edge_list.csv', 'node_metadata.parquet']:
        fpath = f'{OUTPUT_DIR}/{fname}'
        if os.path.exists(fpath):
            size = os.path.getsize(fpath)
            print(f"  {fname}: {size/1024:.1f} KB")
        else:
            print(f"  {fname}: MISSING")

    print("\nDone!")
    return G

if __name__ == '__main__':
    G = build_road_graph()