
import os

os.chdir('D:\\ENEA_CAS_WORK\\Catania_RAFAEL\\viasat_data')
os.getcwd()

import numpy as np
import pandas as pd
import geopandas as gpd
from geopandas import GeoDataFrame
from shapely.geometry import Point
import folium
import osmnx as ox
import networkx as nx
import math
import momepy
# from funcs_network_FK import roads_type_folium
from shapely import geometry
from shapely.geometry import Point, Polygon
import psycopg2
# import db_connect
import datetime
import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from folium_stuff_FK_map_matching import plot_graph_folium_FK
import glob
from funcs_network_FK import cost_assignment
import statistics


#####
##########
###################
##########################
###############################
#####################################
#############################################
#####################################################
#########################################################
###############################################################
###############################################################
### VULNERABILITY ALGORITHM ###################################
###############################################################
###############################################################

# load complete TIME_EDGES files obtained from VIASAT analysis
os.chdir('D:\\ENEA_CAS_WORK\\Catania_RAFAEL\\viasat_data')
# load all EDGES obtained from the map-matching algorithm

# TIME_EDGES = gpd.read_file("sottorete_speed.geojson")   ### all vehicles
TIME_EDGES = gpd.read_file("sottorete_speed_counts_veicoli_pesanti.geojson")   ### only veicoli pesanti


## load a grafo for a sottorete
# grafo_sottorete = ox.load_graphml("Catania__Italy_60km_sottorete.graphml")
# grafo_sottorete_nodes, grafo_sottorete_edges = ox.graph_to_gdfs(grafo_sottorete)
# grafo_sottorete_edges.plot()
# TIME_EDGES = pd.merge(grafo_sottorete_edges[['u', 'v']],TIME_EDGES, on=['u', 'v'], how='left')
# ## drop all rows with NA values
# TIME_EDGES = TIME_EDGES.dropna()
# TIME_EDGES.reset_index(level=0, inplace=True)
# TIME_EDGES = gpd.GeoDataFrame(TIME_EDGES)
# TIME_EDGES.plot()

# load all hexagonal cells (hex_grid)
hex_grid = gpd.read_file("hex_grid.geojson")

### load grafo with cost obtained from VIASAT data (map-matching)
file_graphml = 'CATANIA_VIASAT_cost.graphml'
grafo = ox.load_graphml(file_graphml)
# make attr "cost" as float!!! (units of "cost" are SECONDS!)
for u, v, key, attr in grafo.edges(keys=True, data=True):
    if len(attr['VIASAT_cost']) > 0:
        attr['VIASAT_cost'] = float(attr.get("VIASAT_cost"))
    if len(attr['cost']) >0:
        attr['cost'] = float(attr.get("cost"))

# get geodatframes for nodes and edges
gdf_nodes, gdf_edges = ox.graph_to_gdfs(grafo)

# change directory
# os.chdir('D:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing\\vulnerability')

####################################################################################
####################################################################################
# create basemap around Catania
ave_LAT = 37.53988692816245
ave_LON = 15.044971594798902
my_map = folium.Map([ave_LAT, ave_LON], zoom_start=8.5, tiles='cartodbpositron')
####################################################################################
####################################################################################

# # load all 'elements' (these are the indices of the hex cells associated to all edges crossing each of them)
elements = np.load('elements.npy',allow_pickle='TRUE').item()

###################################################################################
## calculate the initial travel timeS and shortest paths between the OD pairs #####
###################################################################################

penalty = 18000  # PENALTY time or DISRUPTION time (seconds of closure of the link (u,v))
## set 'importance' of an edges to 0 at the NULL scenario (no disruption)
TIME_EDGES['importance'] = 0
hex_grid['importance'] = 0


## load al ORIGINS and DESTINATIONS of the "sottorete"
all_catania_OD = pd.read_csv("all_catania_OD.csv", delimiter=',')

## chose one ORIGIN and DESTINATION

# element = elements[25]
# index_elements = [25]

# i = 418506072
# j = 3914961551




# elements.keys()
element = elements[25]
index_elements = [25]
# element.plot()
## merge element with ORIGINS and DESTINATIONS
# element = pd.merge(element, all_catania_OD[['u','v','ORIGIN', 'DESTINATION']], on=['u', 'v'], how='left')
## drop all rows with NA values
# element = element.dropna()



for index_elements, element in list(elements.items()):
    DELAY = 0
    DELAY_CELL = []
    print("CELL:", index_elements)
    gdf_elements = element
    ## merge element with ORIGINS and DESTINATIONS
    gdf_elements = pd.merge(gdf_elements, all_catania_OD[['u', 'v', 'ORIGIN', 'DESTINATION']], on=['u', 'v'], how='left')
    ## drop all rows with NA values
    gdf_elements = gdf_elements.dropna()
    # gdf_elements = gdf_elements.drop_duplicates(['ORIGIN', 'DESTINATION'])
    gdf_elements.reset_index(level=0, inplace=True)
    # O = list(gdf_elements.ORIGIN)
    O = list(gdf_elements.ORIGIN.unique())
    # D = list(gdf_elements.DESTINATION)
    D = list(gdf_elements.DESTINATION.unique())
    zipped_OD = zip(O, D)
    # loop ever each ORIGIN --> DESTINATION pair
    for (i, j) in zipped_OD:
        print(i,j)
        try:
            ## find shortest path based on the "cost" (time)
            ## NULL scenario
            init_shortest_OD_path_cost = nx.shortest_path(grafo, i, j, weight='VIASAT_cost')  # using cost (time)
            print("DO IT++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++")
            path_edges = list(zip(init_shortest_OD_path_cost,init_shortest_OD_path_cost[1:]))
            lr = nx.shortest_path_length(grafo, i,j, weight='VIASAT_cost') ## this is a time
            lunghezza=[]
            if lr != 0:
                for l in path_edges:
                  lunghezza.append(grafo [l[0]] [l[1]] [0]['length'])  # get only the length for each arch between 2 path edges, [0] it the key = 0
                print("km:{0:.3f} h:{1:.3f} vm:{2:.0f}".format(sum(lunghezza)/1000, lr/3600, sum(lunghezza)/1000/lr*3600))  # units == km
                time_OD_NULL = lr/3600  # hours
                length_OD_NULL = sum(lunghezza)/1000  #km

                ## update MAP with paths
                df_nodes = pd.DataFrame(path_edges)
                df_nodes.columns = ['u', 'v']
                ## merge 'df_nodes' with 'gdf_edges'
                edges_shortest_route_VIASAT_cost = pd.merge(df_nodes, gdf_edges, on=['u', 'v'], how='left')
                edges_shortest_route_VIASAT_cost = gpd.GeoDataFrame(edges_shortest_route_VIASAT_cost)
                edges_shortest_route_VIASAT_cost.drop_duplicates(['u', 'v'], inplace=True)


                ## merge 'df_nodes' with 'TIME_EDGES' (sottorete)
                edges_matched_route_OD = pd.merge(df_nodes, TIME_EDGES, on=['u', 'v'], how='left')
                edges_matched_route_OD = gpd.GeoDataFrame(edges_matched_route_OD)
                edges_matched_route_OD.drop_duplicates(['u', 'v'], inplace=True)

                 ## add plot in Folium map
                 ## save first as geojson file
                 # style_VIASAT = {'fillColor': '#00000000', 'color': '#ff0000'}
                 # edges_shortest_route_VIASAT_cost.geometry.to_file(filename='shortest_route_VIASAT_NULL.geojson',
                 #                                                  driver='GeoJSON')
                 # folium.GeoJson('shortest_route_VIASAT_NULL.geojson', style_function=lambda x: style_VIASAT).add_to((my_map))
                 # my_map.save("VULNERABILITY_map.html")

                ## get the OD travel demand (counts/travelled time) (NULL SCENARIO, no penalty)
                ## https://aetransport.org/public/downloads/QIoR6/498-514ec4e5be18d.pdf
                if len(edges_matched_route_OD) > 0:
                    ## only use vehicle number (counts, MEAN from each EDGE)  and then divide by the time_OD_NULL....this is the TRAVEL_DEMAND_OD
                    counts = (pd.DataFrame(edges_matched_route_OD))['counts']
                    counts = counts.dropna()
                    ## Travel demand: divide by 4 (months), then by 30 (days), then by 16 (hours of higher traffic flow) (hours)
                    TRAVEL_DEMAND_OD = ((((counts.sum())/4)/30)/16)  # veichles/hour all over the path (trajectory)



            ## for each edge (LINK) in the "element" assign a penalty time (disruption)
            # close a LINK (u,v pair) by adding sufficiently large penalty M (time in seconds ~ 5 hours = 18000 secs)
            for u, v, key, attr in grafo.edges(keys=True, data=True):
                attr['VIASAT_cost_penalty'] = attr.get("VIASAT_cost")
                zipped = zip(list(gdf_elements.u), list(gdf_elements.v))
                if (u, v) in zipped:
                    print(u,v)
                    print("gotta!=============================================================================================")
                    # break
                    attr['VIASAT_cost_penalty'] = float(attr['VIASAT_cost']) + penalty
                    print(attr['VIASAT_cost_penalty'])
                    grafo.add_edge(u, v, key, attr_dict=attr)
            # get shortest path again...but now with the PENALTY
            shortest_OD_path_VIASAT_penalty = nx.shortest_path(grafo, i, j,
                                                                    weight='VIASAT_cost_penalty')
            path_edges = list(zip(shortest_OD_path_VIASAT_penalty, shortest_OD_path_VIASAT_penalty[1:]))
            lr = nx.shortest_path_length(grafo, i, j, weight='VIASAT_cost_penalty')
            if lr !=0:
                lunghezza = []
                for l in path_edges:
                    lunghezza.append(grafo[l[0]][l[1]][0][
                                         'length'])  # get only the length for each arch between 2 path edges, [0] it the key = 0
                print("km:{0:.3f} h:{1:.3f} vm:{2:.0f}".format(sum(lunghezza) / 1000, lr / 3600,
                                                               sum(lunghezza) / 1000 / lr * 3600))  # units == km
                # total time from O--D with penalty
                time_OD_penalty = lr/3600
                length_OD_penalty = sum(lunghezza) / 1000  # km

                # update MAP with paths
                # add shortest path (by "VIASAT cost" to the my_map
                df_nodes = pd.DataFrame(path_edges)
                df_nodes.columns = ['u', 'v']
                # filter 'grafo' with 'df_nodes'
                # keys = list(df_nodes.columns.values)
                # index_gdf_edges = gdf_edges.set_index(keys).index
                # index_df_nodes = df_nodes.set_index(keys).index
                # edges_shortest_route_VIASAT_penalty = gdf_edges[index_gdf_edges.isin(index_df_nodes)]

                ## merge 'df_nodes' with 'gdf_edges'
                edges_shortest_route_VIASAT_penalty = pd.merge(df_nodes, gdf_edges, on=['u', 'v'], how='left')
                edges_shortest_route_VIASAT_penalty = gpd.GeoDataFrame(edges_shortest_route_VIASAT_penalty)
                edges_shortest_route_VIASAT_penalty.drop_duplicates(['u', 'v'], inplace=True)

                # add plot in Folium map
                ## save first as geojson file
                # style_VIASAT_penalty = {'fillColor': '#00000000', 'color': '#0000ff'}
                # edges_shortest_route_VIASAT_penalty.geometry.to_file(
                #     filename='shortest_route_VIASAT_penalty.geojson',
                #     driver='GeoJSON')
                # folium.GeoJson('shortest_route_VIASAT_penalty.geojson',
                #                style_function=lambda x: style_VIASAT_penalty).add_to((my_map))
                # my_map.save("VULNERABILITY_map.html")

            ## calculate "closure impact" for each DESTINATION
            ## difference time between closure and normal conditions
            DT = time_OD_penalty - time_OD_NULL  # (hours)
            ## closure impact
            if DT < penalty / 3600:
                DELAY = DELAY + TRAVEL_DEMAND_OD * DT * ((penalty / 3600) - DT / 2)  # vehicles*hours
                DELAY_CELL.append(DELAY)
                print("TOTAL DELAY:", DELAY)
                print("max(DELAY_CELL):", max(DELAY_CELL), "CELL:", index_elements)
            else:
                DELAY = DELAY + TRAVEL_DEMAND_OD * (((penalty / 3600) ** 2) / 2)  # vehicles*hours
                DELAY_CELL.append(DELAY)
                print("TOTAL DELAY:", DELAY)
                print("max(DELAY_CELL):", max(DELAY_CELL), "CELL:", index_elements)

            # create a new field "importance" in the TIME_EDGES geodataframe
            U = list(element.u)
            V = list(element.v)
            zipped_UV = zip(U, V)
            for (U, V) in zipped_UV:
                # print(U, V)
                mask = ((TIME_EDGES['u'] == U) & (TIME_EDGES['v'] == V))
                TIME_EDGES.loc[mask, 'importance'] = DELAY
                mask_cell = ((hex_grid.index == index_elements))
                hex_grid.loc[mask_cell, 'importance'] = max(DELAY_CELL)

            # restore initial travel time for each edge (LINK)....basically we need to remove the penalty time
            for u, v, key, attr in grafo.edges(keys=True, data=True):
                zipped = zip(list(gdf_elements.u), list(gdf_elements.v))
                if (u, v) in zipped:
                    # print(u,v)
                    # print("gotta!=============================================================================================")
                    attr['VIASAT_cost_penalty'] = float(attr['VIASAT_cost'])
                    grafo.add_edge(u, v, key, attr_dict=attr)

            with open("last_CELL_ID.txt", "w") as text_file:
                    text_file.write("last CELL ID: %s" % (index_elements))
            TIME_EDGES['importance'] = TIME_EDGES['importance'].replace(np.nan, 0)
        except (nx.NodeNotFound, nx.exception.NetworkXNoPath):
            print('O-->D NodeNotFound', 'i:', i, 'j:', j)

## save geojson files
TIME_EDGES.to_file(filename='TIME_EDGES_importance.geojson', driver='GeoJSON')


###################################
##### MAP #########################
###################################

# load TIME_EDGES files
TIME_EDGES = gpd.read_file("TIME_EDGES_importance.geojson")

########################
### TIME_EDGES file ####
########################

#############################################################################################
# create basemap
ave_LAT = 37.53988692816245
ave_LON = 15.044971594798902
my_map = folium.Map([ave_LAT, ave_LON], zoom_start=11, tiles='cartodbpositron')
#############################################################################################

## add width or weight to each LINESTRING based on the TIME_EDGES on the 'importance' field
TIME_EDGES['stroke-width'] = (round(TIME_EDGES['importance']/max(TIME_EDGES['importance']), 2) *3) +0.1
TIME_EDGES['importance'] = round(TIME_EDGES['importance'], 0)
TIME_EDGES=TIME_EDGES.rename(columns = {'importance':'importance (vei*hour)'})

folium.GeoJson(
TIME_EDGES[['u','v', 'importance (vei*hour)', 'stroke-width', 'geometry']].to_json(),
    style_function=lambda x: {
        'fillColor': 'black',
        'color': 'black',
        'weight':  x['properties']['stroke-width'],
        'fillOpacity': 1,
        },
highlight_function=lambda x: {'weight':3,
        'color':'blue',
        'fillOpacity':1
    },
    # fields to show
    tooltip=folium.features.GeoJsonTooltip(
        fields=['u', 'v', 'importance (vei*hour)']),
    ).add_to(my_map)

## save map
my_map.save("VULNERABILITY_map.html")
