# -*- coding: utf-8 -*-
"""
This script takes the graph data created by Script 1 and uses it to create a distance matrix for every CBS user scenario. 
This distance matrix is the key input for the Google-OR Tools Route solver (see description for form of matrix)

It is written to be run cell-by-cell.

@author: Charlie Ferguson
"""


#%% Import packages

from __future__ import print_function
import pandas as pd
import numpy as np
import networkx as nx
from multiprocessing import Pool
from pathlib import Path

#%% Define Type of Solver

Solver="local" #or "server"

if Solver=="local":
    FilePath="{Path to local server}"
elif Solver=="server":
    FilePath="{Path to remote server}"
del Solver


#%% Define Variables

Location="Lima" #Location identifier
SpaCRS="WGS 1984 UTM Zone 37S" #WGS1984 UTM Co-ordinate Reference System

NumSetsArray=np.array([[10, 20, 30,50,70,100,250,500,1000,3000,5000],  #This defines the number of CBS users and the number of iterations for each
                       [50, 50, 50,50,50,50 ,25, 25, 10,  10,10]])

CBSUsersFileSets=["CBSUsers_Num_10","CBSUsers_Num_20","CBSUsers_Num_30","CBSUsers_Num_50","CBSUsers_Num_70", #This defines the names of the folders (following from previous Python Script 1)
                  "CBSUsers_Num_100","CBSUsers_Num_250","CBSUsers_Num_500","CBSUsers_Num_1000",
                  "CBSUsers_Num_3000","CBSUsers_Num_5000"]


#%% Create Distance Matrix (this is written to run in a HPC, can be converted in for loop for local machine use)

def PrimaryFunction(i):
    NumSets=NumSetsArray[1,i]
    for j in range(0,NumSets):
        my_file = Path(FilePath+"/InputFiles/"+Location+"_Inputs/"+CBSUsersFileSets[i]+"/"+CBSUsersFileSets[i]+"_Set_"+str(j)+ "_DistanceMatrix.csv")
        if my_file.is_file():
               # The file already exists
               pass
        else:
            #Identify which edges are connected to which nodes
            df = pd.read_csv (FilePath+"/InputFiles/"+Location+"_Inputs/"+CBSUsersFileSets[i]+"/"+CBSUsersFileSets[i]+"_Set_"+str(j)+ "_ProximityTableCSV.csv")
            df=df[["IN_FID","NEAR_FID"]]
            a = np.array([1,2])
            df['node'] = np.resize(a, df.shape[0])
            dfB=df.pivot(index="IN_FID",columns="node",values="NEAR_FID")
            dfB.columns = ['node1', 'node2']
            ###
            #Get length of each of the edges
            df2 = pd.read_csv (FilePath+"/InputFiles/"+Location+"_Inputs/"+CBSUsersFileSets[i]+"/"+CBSUsersFileSets[i]+"_Set_"+str(j)+ "_LinesCSV.csv")
            df2=df2[["FID","Length"]]
            dfB[['distance']]=df2[["Length"]].to_numpy()
            ###
            #Bring in the nodes
            dfpoints = pd.read_csv (FilePath+"/InputFiles/"+Location+"_Inputs/"+CBSUsersFileSets[i]+"/"+CBSUsersFileSets[i]+"_Set_"+str(j)+ "_PointsCSV.csv")
            dfpoints=dfpoints[["FID","POINT_X","POINT_Y"]].to_numpy()
            dfpoints=pd.DataFrame(dfpoints,columns=["id","X","Y"])
            ###
            edgelist=dfB
            nodelist=dfpoints
            ##
            #Create empty graph in NetworkX
            g=nx.Graph()
            #Populate graph with pre-processed data
            for k, elrow in edgelist.iterrows():
                g.add_edge(elrow[0], elrow[1], **elrow[2:].to_dict())
                ##
            for k, nlrow in nodelist.iterrows():
                nx.set_node_attributes(g, {nlrow['id']:  nlrow[1:].to_dict()})
                ##
            #Create distance matrix from Dijkstra algorthm
            B=dict(nx.all_pairs_dijkstra_path_length(g,weight="distance"))
            DistanceMatrix= pd.DataFrame(index=range(len(dfpoints)),columns=range(len(dfpoints)))
            for node1 in B:
                for node2 in B[node1]:
                    DistanceMatrix.iloc[int(node1),int(node2)]=B[node1][node2]
                    ##
                ##
            #Depending on whether the resulting matrix contains any NaNs, this is saved
            if DistanceMatrix.isnull().sum().sum()==0:
                DistanceMatrix2=  DistanceMatrix.round(0).astype(int)
                np.savetxt(FilePath+"/InputFiles/"+Location+"_Inputs/"+CBSUsersFileSets[i]+"/"+CBSUsersFileSets[i]+"_Set_"+str(j)+ "_DistanceMatrix.csv",DistanceMatrix2, delimiter=",")
            else:
                np.savetxt(FilePath+"/InputFiles/"+Location+"_Inputs/"+CBSUsersFileSets[i]+"/"+CBSUsersFileSets[i]+"_Set_"+str(j)+ "_DistanceMatrixFAIL.csv",DistanceMatrix2, delimiter=",")
        ##
        #####
#The following 'if' statement allows parallelisation of computation in a HPC
if __name__ == '__main__':
     pool = Pool()  #os.cpu_count() - 1
     pool.map(PrimaryFunction, range(0,len(CBSUsersFileSets)))