
import os

os.chdir('C:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing')
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
import db_connect
from shapely import wkb
import sqlalchemy as sal

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

# load complete TIME_EDGES files obtained from VIASAT analysis
os.chdir('C:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing')


conn_HAIG = db_connect.connect_HAIG_Viasat_CT()
cur_HAIG = conn_HAIG.cursor()

# Function to generate WKB hex
def wkb_hexer(line):
    return line.wkb_hex

# Create an SQL connection engine to the output DB
engine = sal.create_engine('postgresql://postgres:vaxcrio1@localhost:5432/HAIG_Viasat_CT')

os.chdir('C:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing')
os.getcwd()


# load file containing times and speed, RECORDS and FREQUENCIES from the Database (from the map-matching)
gdf_all_EDGES = pd.read_sql_query(
    ''' SELECT *
        FROM public.paths_postprocess_temp''', conn_HAIG)

## transform Geometry from text to LINESTRING
def wkb_tranformation(line):
   return wkb.loads(line.geom, hex=True)

gdf_all_EDGES['geometry'] = gdf_all_EDGES.apply(wkb_tranformation, axis=1)
gdf_all_EDGES.drop(['geom'], axis=1, inplace= True)

## drop 'level_0' column
gdf_all_EDGES.drop(['level_0'], axis=1, inplace = True)

## add "ORIGIN" and "DESTINATION"
## for each idtrajectory....get first (ORIGIN) and last element (DESTINATION) from min(sequence) and max(sequence)
# df_origin = gdf_all_EDGES.groupby('idtrajectory').agg({'sequenza': min}).reset_index()
# df_origin = pd.merge(df_origin, gdf_all_EDGES, on = ['idtrajectory', 'sequenza'], how = 'left')
# df_origin = df_origin[['idtrajectory', 'sequenza', 'u']]
# df_destination = gdf_all_EDGES.groupby('idtrajectory').agg({'sequenza': max}).reset_index()
# df_destination = pd.merge(df_destination, gdf_all_EDGES, on = ['idtrajectory', 'sequenza'], how = 'left')
# df_destination = df_destination[['idtrajectory', 'sequenza', 'v']]
# df_origin=df_origin.rename(columns = {'u':'ORIGIN'})
# df_destination=df_destination.rename(columns = {'v':'DESTINATION'})
# df_destination.drop(['sequenza'], axis=1, inplace= True)
# df_origin.drop(['sequenza'], axis=1, inplace= True)
#
# OD = pd.merge(df_origin, df_destination, on=['idtrajectory'], how='left')
# gdf_all_EDGES = pd.merge(gdf_all_EDGES,OD, on = ['idtrajectory'], how = 'left')
gdf_all_EDGES = gpd.GeoDataFrame(gdf_all_EDGES)
gdf_all_EDGES.plot()


os.chdir('C:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing\\vulnerability')
# load all hexagonal cells (hex_grid)
hex_grid = gpd.read_file("hex_grid_DB.geojson")

### load grafo with cost obtained from VIASAT data (map-matching)
os.chdir('C:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing')
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('C:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing\\vulnerability')

####################################################################################
####################################################################################
# create basemap
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_DB.npy',allow_pickle='TRUE').item()
# print(read_dictionary['hello']) # displays "world"

###################################################################################
## 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))
gdf_all_EDGES['importance'] = 0
hex_grid['importance'] = 0

## chose one ORIGIN and DESTINATION
# i = 5049834302
# j = 318758878

## choose one element for the map "hex_grid_DB.html"
# element = elements[439]
# index_elements = [439]


# find all ORIGINS and DESTINATIONS
# OD = TIME_EDGES.drop_duplicates(['ORIGIN', 'DESTINATION'])
# O = list(OD.ORIGIN)
# D = list(OD.DESTINATION)


for index_elements, element in list(elements.items()):
    DELAY = 0
    DELAY_CELL = []
    print("CELL:", index_elements)
    gdf_elements = element
    gdf_elements = gdf_elements.drop_duplicates(['ORIGIN', 'DESTINATION'])
    gdf_elements.reset_index(level=0, inplace=True)
    O = list(gdf_elements.ORIGIN)
    D = list(gdf_elements.DESTINATION)
    zipped_OD = zip(O, D)
    # loop ever each ORIGIN --> DESTINATION pair
    for (i, j) in zipped_OD:
        print(i, j)
        try:
            ### chose one ORIGIN i for now....and its DESTINATION j (it is a pair)
            # i = O[0]
            # j = D[0]

            ## 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)
            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')
            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']
                # 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_cost = gdf_edges[index_gdf_edges.isin(index_df_nodes)]

                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)

                # keys = list(df_nodes.columns.values)
                # index_TIME_EDGES = gdf_all_EDGES.set_index(keys).index
                # index_df_nodes = df_nodes.set_index(keys).index
                # edges_matched_route_OD = gdf_all_EDGES[index_TIME_EDGES.isin(index_df_nodes)]

                edges_matched_route_OD = pd.merge(df_nodes, gdf_all_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")

                # calculate travel demand (flux) (NULL SCENARIO, no penalty)
                if len(edges_matched_route_OD) > 0:
                    df_flux = edges_matched_route_OD
                    df_flux['timedate'] = pd.to_datetime(df_flux['timedate'])
                    ## use 15 minutes average
                    df_flux = df_flux.resample(rule='15Min', on='timedate').mean()
                    df_flux['records'] = df_flux['records'].ffill()  # 15 minutes flux
                    df_flux['records'] = df_flux['records'] / 0.25
                    BBB = pd.DataFrame(df_flux['records'])
                    TRAVEL_DEMAND_OD = BBB.records.mean()  # veichles/hour



            ## 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)]

                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)
            else:
                DELAY = DELAY + TRAVEL_DEMAND_OD * (((penalty / 3600) ** 2) / 2)  # vehicles*hours
                DELAY_CELL.append(DELAY)
                print("TOTAL DELAY:", DELAY)

            # 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 = ((gdf_all_EDGES['u'] == U) & (gdf_all_EDGES['v'] == V))
                gdf_all_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))
            # save TIME_EDGES files containing the total delay time in each element (links in each hexagonal cells)
            gdf_all_EDGES['importance'] = gdf_all_EDGES['importance'].replace(np.nan, 0)
            gdf_all_EDGES.to_file(filename='TIME_EDGES_importance_DB.geojson', driver='GeoJSON')
            hex_grid.to_file(filename='hex_grid_importance_DB.geojson', driver='GeoJSON')

        except (nx.NodeNotFound, nx.exception.NetworkXNoPath):
            print('O-->D NodeNotFound', 'i:', i, 'j:', j)






###########################################################
##### MAP with different weights on the links #############
###########################################################


os.chdir('C:\\ENEA_CAS_WORK\\Catania_RAFAEL\\postprocessing\\vulnerability')
# load TIME_EDGES files and all hexagonal cells (hex_grid)
TIME_EDGES = gpd.read_file("TIME_EDGES_importance.geojson")
hex_grid = gpd.read_file("hex_grid_importance.geojson")

# TIME_EDGES = TIME_EDGES[( TIME_EDGES.importance < 5000) & (TIME_EDGES.importance > 1)]

########################
### 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)'})
## save GeoJson file
# TIME_EDGES = gpd.read_file("TIME_EDGES_importance.geojson")

# AAA = pd.DataFrame( TIME_EDGES[TIME_EDGES['stroke-width'] > 0])

folium.GeoJson(
TIME_EDGES[['u','v', 'importance (vei*hour)', 'color', '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")



# ## make a color map for the link importance (vehicles*hours) of each element in each hexagonal cell
# # add colors based on 'importance' (vehicles*hours)
# vmin = min(TIME_EDGES.importance)
# vmax = max(TIME_EDGES.importance)
# # Try to map values to colors in hex
# norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
# mapper = plt.cm.ScalarMappable(norm=norm, cmap=plt.cm.YlOrRd)  # scales of Reds (or "coolwarm" , "bwr", °cool°)
# TIME_EDGES['color'] = TIME_EDGES['importance'].apply(lambda x: mcolors.to_hex(mapper.to_rgba(x)))
# TIME_EDGES['importance'] = round(TIME_EDGES['importance'], 0)
#
# ### add colors to map
# my_map = plot_graph_folium_FK(TIME_EDGES, graph_map=None, popup_attribute=None,
#                               zoom=1, fit_bounds=True, edge_width=2, edge_opacity=1)
# style = {'fillColor': '#00000000', 'color': '#00000000'}
# # add 'u' and 'v' as highligths for each edge (in blue)
# folium.GeoJson(
#     # data to plot
#     TIME_EDGES[['u', 'v', 'importance', 'geometry']].to_json(),
#     show=True,
#     style_function=lambda x:style,
#     highlight_function=lambda x: {'weight':3,
#         'color':'blue',
#         'fillOpacity':1
#     },
#     # fields to show
#     tooltip=folium.features.GeoJsonTooltip(
#         fields=['u', 'v', 'importance']
#     ),
# ).add_to(my_map)
# folium.TileLayer('cartodbdark_matter').add_to(my_map)
# folium.LayerControl().add_to(my_map)
#
# my_map.save("VULNERABILITY_map.html")


########################
### hex_grid file #####
########################

## make a color map for the link importance (vehicles*hours) of each element in each hexagonal cell
# add colors based on 'importance' (vehicles*hours)
vmin = min(hex_grid.importance)
vmax = max(hex_grid.importance)
# Try to map values to colors in hex
norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
mapper = plt.cm.ScalarMappable(norm=norm, cmap=plt.cm.Reds)  # scales of Reds (or "coolwarm" , "bwr", °cool°)
hex_grid['colors'] = hex_grid['importance'].apply(lambda x: mcolors.to_hex(mapper.to_rgba(x)))
hex_grid['importance'] = round(hex_grid['importance'], 0)


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

hex_grid=hex_grid.rename(columns = {'importance':'importance (vei*hour)'})

folium.GeoJson(
hex_grid[['importance (vei*hour)', 'colors', 'geometry']].to_json(),
    style_function=lambda x: {
        'fillColor': x['properties']['colors'],
        'color': x['properties']['colors'],
        'weight': 1,
        'fillOpacity': 0.6,
        },
highlight_function=lambda x: {'weight':3,
        'color':'blue',
        'fillOpacity':0.5
    },
    # fields to show
    tooltip=folium.features.GeoJsonTooltip(
        fields=['importance (vei*hour)']),
    ).add_to(my_map)


folium.TileLayer('cartodbdark_matter').add_to(my_map)
folium.LayerControl().add_to(my_map)

my_map.save("cells_VULNERABILITY_map.html")


