import os
import networkx as nx
import numpy as np
import pandas as pd
import geopandas as gpd
import time
import math
from shapely import wkt
from shapely.geometry import Point, Polygon
import multiprocessing as mp

from helpers import from_p_file
from ref_systems import osm_crs, projected_crs


categories =['healthcare', 'services', 'transport', 'outdoor', 'supplies', 'restaurant', 'culture', 'education', 'physical']


# function to create a hexagonal grid
def create_hexagonal_grid(polygon, side_length):
    """
    Create a hexagonal grid inside a given polygon with specified side length.

    Parameters:
        polygon (shapely.geometry.Polygon): input polygon.
        side_length (float): length of each hexagon side in meters.

    Returns:
        GeoDataFrame: hexagonal grid clipped to the polygon.
    """
    # spacing between hexagon centers
    dx = (math.sqrt(3) * side_length)  # horizontal distance between centers
    dy = 1.5 * side_length  # vertical distance between centers

    # bounding box of the polygon
    minx, miny, maxx, maxy = polygon.bounds

    # generate hexagon centers
    hex_centers = []
    y = miny
    row = 0
    while y <= maxy + dy:
        x = minx
        if row % 2 == 1:  # offset odd rows
            x += (math.sqrt(3) * side_length) / 2

        while x <= maxx + dx:
            hex_centers.append(Point(x, y))
            x += dx
        y += dy
        row += 1

    # create hexagons around each center
    hexagons = []
    for center in hex_centers:
        hexagon = Polygon([
            (center.x + side_length * math.cos(angle), center.y + side_length * math.sin(angle))
            for angle in [math.pi / 6 + i * math.pi / 3 for i in range(6)]
        ])
        hexagons.append(hexagon)

    # convert to GeoDataFrame
    grid = gpd.GeoDataFrame(geometry=hexagons)
    grid['centroids'] = hex_centers

    # clip the grid to the input polygon
    clipped_grid = grid[grid.intersects(polygon)].copy()
    clipped_grid.reset_index(drop=True, inplace=True)

    return clipped_grid


def hexagon_path_computation(hex_grid, G, pois):
    hex_grid_copy = hex_grid.copy()
    hex_grid_copy['time'] = float('NaN')
    print(f"Index {hex_grid.index[0]} to index {hex_grid.index[-1]} is being processed by PID {os.getpid()}")
    
    # compute the fastest path to nearest X POIs
    start = time.time()
    duration, indices = [], []
    for ix, n in hex_grid.iterrows():
        trip_dur = []
        current_pois = pois.copy()

        while len(trip_dur) < 40:
            sindex = current_pois.sindex
            _, neighbors_ix = sindex.nearest(n.centroids, return_all=False)
            dest = current_pois.iloc[neighbors_ix]

            current_pois.drop(dest.index, axis=0, inplace=True)
            if len(dest) == 0:
                break

            try:
                dur = nx.shortest_path_length(G, n.origin, dest.nearest.iloc[0], 'time')
            except nx.exception.NetworkXNoPath:
                continue

            trip_dur.append(dur)

        # average the 20 nearest durations
        if trip_dur:
            if all(x == float('inf') for x in trip_dur):
                avg_dur = float('NaN')
            else:
                trip_dur_20 = sorted(trip_dur)[:20]
                avg_dur = np.nanmean(trip_dur_20)
        else:
            avg_dur = float('NaN')

        duration.append(avg_dur)
        indices.append(ix)
    hex_grid_copy.loc[indices, 'time'] = duration
    print(time.time() - start)
    return hex_grid_copy


if __name__ == "__main__":
    # ask the user to select the speed model, yes: slope-dependent walking and biking speed is used, no: length / avg_speed (4.8km/h walk or 21.6 km/h bike)
    terrain = input("Include slope-dependent walking and biking speed? Please enter either 'yes' or 'no': ")

    # number of multiprocesses in parallel
    num_processes = 4
    
    # full path to AGILE_2025_data_code folder
    fullpath = os.path.abspath(os.path.join(os.getcwd(), '..'))
    
    # input paths
    path_roi = os.path.join(fullpath, 'data', 'BEZIRKSGRENZEOGD', 'BEZIRKSGRENZEOGDPolygon.shp')
    path_pois = os.path.join(fullpath, 'results', 'pois')
    path_g_walk = os.path.join(fullpath, 'data', 'graphs_small', 'walk')

    for mode in ['walk', 'bike', 'walk_PT']:
        print(mode)
        
        # results for 'walking plus public transport' AND no terrain are not needed --> skip
        if mode == 'walk_PT' and terrain == 'no':
            continue
        
        # input path (mode dependent)
        path_g = os.path.join(fullpath, 'data', 'graphs_small', mode)
        
        # output path (mode dependent)
        path_results = os.path.join(fullpath, 'results', 'accessibility_small', mode, 'terrain_' + terrain)
        
        # create the output directory if it doesn't exist
        if not os.path.exists(path_results):
            os.makedirs(path_results)

        # read ROI shapefile
        borders = gpd.read_file(path_roi)
        borders.to_crs(projected_crs, inplace=True)
        polygon = borders.unary_union

        # create hexagonal grid
        side_length = 200  # side length of each hexagon in meters
        hex_grid = create_hexagonal_grid(polygon, side_length)
        hex_grid.set_crs(projected_crs, inplace=True)

        # read graph
        G = from_p_file(path_g)
        
        # convert to average speeds if needed
        if terrain == 'no':
            if mode == 'walk':
                for u, v, i in G.edges:
                    temp = G[u][v][i].copy()
                    avg_time = temp['length'] / (4.8 *(1000/60))  # 4.8km/h to m/min
                    G[u][v][i]['time'] = avg_time
            if mode == 'bike':
                for u, v, i in G.edges:
                    temp = G[u][v][i].copy()
                    avg_time = temp['length'] / (21.6 *(1000/60))  # 21.6km/h to m/min
                    G[u][v][i]['time'] = avg_time

        # check largest connected component to find reasonable origins
        if mode == 'walk_PT':
            G_walk = from_p_file(path_g_walk)  # additional walking graph needed to find origins
            G_undir = nx.MultiGraph(G_walk)
            temp_G_nodes = max(nx.connected_components(G_undir), key=len)
            temp_G_graph = G_walk.subgraph(temp_G_nodes)
        else:
            G_undir = nx.MultiGraph(G)
            temp_G_nodes = max(nx.connected_components(G_undir), key=len)
            temp_G_graph = G.subgraph(temp_G_nodes)

        # graph nodes to dataframe
        temp_G = pd.DataFrame.from_dict(dict(temp_G_graph.nodes(data=True)), orient='index')
        temp_G = gpd.GeoDataFrame(temp_G, geometry='geom', crs=osm_crs)
        temp_G.to_crs(projected_crs, inplace=True)
        temp_G.reset_index(inplace=True)
        temp_G.rename(columns={'index': 'node_id'}, inplace=True)

        # assign each centroid the nearest node ON the graph
        sindex = temp_G.sindex
        _, neighbors_ix = sindex.nearest(hex_grid.centroids, return_all=False)
        hex_grid['origin'] = temp_G.iloc[neighbors_ix].node_id.values
        hex_grid = hex_grid[hex_grid.within(borders[borders.BEZNR == 1].unary_union)]  # small extent

        # compute accessibility per category
        for category in categories:
            print(category)
            start_category = time.time()
            
            # input path (category dependent)
            path_poi = os.path.join(path_pois, category + '.csv')
            
            # read POI data
            pois = pd.read_csv(path_poi)
            pois['geom'] = pois['geom'].apply(wkt.loads)
            pois = gpd.GeoDataFrame(pois, geometry='geom', crs=osm_crs)
            pois.to_crs(projected_crs, inplace=True)

            # assign each POI the nearest node ON the graph
            sindex = temp_G.sindex
            _, neighbors_ix = sindex.nearest(pois.geom, return_all=False)
            pois['nearest'] = temp_G.iloc[neighbors_ix].node_id.values

            # split hexgrid in n parts for multiprocessing
            split_size = int(np.ceil(len(hex_grid) / num_processes))
            hex_grid_split = [hex_grid.iloc[i:i + split_size] for i in range(0, len(hex_grid), split_size)]

            # start multiprocessing
            pool = mp.Pool(processes=num_processes)
            p1 = pool.starmap_async(hexagon_path_computation, [(hex_part, G, pois) for hex_part in hex_grid_split])
            result = p1.get()
            final_res = pd.concat(result)

            # store results
            final_res.to_csv(os.path.join(path_results, category + '.csv'))
            
            # close the pool
            pool.close()
            pool.join()
            print(f'category: {category}, ', time.time() - start_category)
        print('done with mode ', mode)
    print('done')