
import struct
import numpy as np
import os
import math
import scipy as sp
import scipy.spatial
import scipy.signal
import matplotlib.pyplot as plt

class vertex:
    def __init__(self, x, y, z, e_out, e_in):
        self.p = np.array([x, y, z])
        self.Eout = e_out
        self.Ein = e_in

class edge:
    def __init__(self, v0, v1, p):
        self.v = (v0, v1)
        self.p = p

class linesegment:
    def __init__(self, p0, p1):
        self.p = (p0, p1)

    #return a set of points sampling the line segment at the specified spacing
    def pointcloud(self, spacing):
        v = self.p[1] - self.p[0]
        l = math.sqrt(v[0] ** 2 + v[1] ** 2 + v[2] ** 2)
        n = int(np.ceil(l / spacing))
        if n == 0:
            return [self.p[0], self.p[1]]
        
        pc = []
        for i in range(n+1):
            pc.append(self.p[0] + v * (i/n))
        return pc

class NWT:
    def __init__(self, filename):        
        [_, fext] = os.path.splitext(filename)                          #get the file extension so that we know the file type
        if fext == ".nwt":                                              #if the file extension is NWT
            self.load_nwt(filename)                                     #load a NWT file
        elif fext == ".obj":                                            #if the file extension is OBJ
            self.load_obj(filename)                                     #load an OBJ file
        else:                                                           #otherwise raise an exception
            raise ValueError("file type is unsupported as a network")

    def load_nwt(self, filename):        
        fid = open(filename, "rb")                                      #open a binary file for reading
        self.header = fid.read(14).decode("utf-8")                      #load the header
        self.desc = fid.read(58).decode("utf-8")                        #load the description
        nv = struct.unpack("I", fid.read(4))[0]                         #load the number of vertices and edges
        ne = struct.unpack("I", fid.read(4))[0]

        self.v = []                                                     #create an empty list to store the vertices        
        for _ in range(nv):                                             #iterate across all vertices
            p = np.fromfile(fid, np.float32, 3)                         #read the vertex position
            E = np.fromfile(fid, np.uint32, 2)                          #read the number of edges
            e_out = np.fromfile(fid, np.uint32, E[0])                   #read the indices of the outgoing edges
            e_in = np.fromfile(fid, np.uint32, E[1])                    #read the indices of the incoming edges
            v = vertex(p[0], p[1], p[2], e_out, e_in)                   #create a vertex
            self.v.append(v)

        self.e = []                                                     #create an empty array to store the edges        
        for _ in range(ne):                                             #iterate over all edges            
            v = np.fromfile(fid, np.uint32, 2)                          #load the vertex indices that this edge connects            
            npts = struct.unpack("I", fid.read(4))[0]                   #read the number of points defining this edge            
            pv = np.fromfile(fid, np.float32, 4*npts)                   #read the array of points            
            p = [(pv[i],pv[i+1]) for i in range(0,npts,2)]              #conver the point values to an array of 4-element tuples            
            self.e.append(edge(v[0], v[1], p))                          #create and append the edge to the edge list

    def load_obj(self, filename):        
        fid = open(filename, "r")                                       #open the file for reading        
        vertices = []                                                   #create an array of vertices
        lines = []                                                      #create an array of lines
        for line in fid:                                                #for each line in the file
            elements = line.split(" ")                                  #split it into token elements            
            if elements[0] == "v":                                      #if the element is a vertex               
                c = [float(i) for i in elements[1:]]                    #get the point coordinates                
                vertices.append(c)                                      #add the coordinates to the vertex list            
            if elements[0] == "l":                                      #if the element is a line                
                idx = [int(i) for i in elements[1:]]                    #get the indices for the points that make up the line                
                lines.append(idx)                                       #add this line to the line list

        self.header = "nwtfileformat "                                  #assign a header and description
        self.desc = "File generated from OBJ"
                                                                        #insert the first and last vertex ID for each line into a set
        vertex_set = set()                                              #create an empty set
        for line in lines:                                              #for each line in the list of lines
            vertex_set.add(line[0])                                     #add the first and last vertex to the vertex set (this will remove redundancies)
            vertex_set.add(line[-1])
        
        obj2nwt = dict()                                                #create a new dictionary - will be used to map vertex IDs in the OBJ to IDs in the NWT object

        #create a mapping between OBJ vertex indices and NWT vertex indices
        vi = 0                                                          #initialize a vertex counter to zero
        for si in vertex_set:                                           #for each vertex in the set of vertices
            obj2nwt[si] = vi                                            #assign the mapping
            vi = vi + 1                                                 #increment the vertex counter

        #iterate through each line (edge), assigning them to their starting and ending vertices
        v_out = [list() for _ in range(len(vertex_set))]                #create an array of empty lists storing the inlet and outlet edges for each vertex
        v_in = [list() for _ in range(len(vertex_set))]

        self.e = []                                                     #create an empty list storing the NWT vertex IDs for each edge (inlet and outlet)
        for li in range(len(lines)):                                    #for each line
            v0 = obj2nwt[lines[li][0]]                                  #get the NWT index for the starting and ending points (vertices)
            v1 = obj2nwt[lines[li][-1]]

            v_out[v0].append(li)                                        #add the line index to a list of inlet edges
            v_in[v1].append(li)                                         #add the line index to a list of outlet edges

            p = []                                                      #create an emptu array of points used to store point positions in the NWT graph
            for pi in range(1, len(lines[li]) - 1):                     #for each point in the line that is not an end point (vertex)
                p.append(np.array(vertices[lines[li][pi]-1]))              #add the coordinates for that point as a tuple into the point list
            self.e.append(edge(v0, v1, p))                              #create an edge, specifying the inlet and outlet vertices and all defining points

        #for each vertex in the set, create a NWT vertex containing all of the necessary edge information
        self.v = []                                                     #create an empty list to store the vertices
        for s in vertex_set:                                            #for each OBJ vertex in the vertex set
            vi = obj2nwt[s]                                             #calculate the corresponding NWT index
            self.v.append(vertex(vertices[s-1][0], vertices[s-1][1], vertices[s-1][2], v_out[vi], v_in[vi]))    #create a vertex object, consisting of a position and attached edges

    #return a set of line segments connecting all points in the network
    def linesegments(self):

        s = []                                                          #create an empty list of line segments
        for e in self.e:                                                #for each edge in the graph
            p0 = self.v[e.v[0]].p                                       #load the first point (from the starting vertex)

            for p in e.p:                                               #for each point in the edge
                p1 = np.array([p[0], p[1], p[2]])                       #get the second point for the line segment
                s.append(linesegment(p0, p1))                           #append the line segment to the list of line segments
                p0 = p1                                                 #update the start point for the next segment to the end point of this one
            
            p1 = self.v[e.v[1]].p                                       #load the last point (from the ending vertex)
            s.append(linesegment(p0, p1))                               #append the last line segment for this edge to the list
        return s

    #return a point cloud sampling the centerline of the network at the given spacing
    def pointcloud(self, spacing):
        ls = self.linesegments()
        pc = []
        for l in ls:
            pc = pc + l.pointcloud(spacing)
        return pc

def gaussian(X, sigma):
    return np.exp(-0.5 * (X ** 2 / sigma ** 2))

#set the input parameters
sigma = 10
GTfile = "04_GT.obj"
Tfile = "04_Ta.obj"

#tunable constants
subdiv = 4                                                            #fraction of sigma used to sample each network
shadow = 10                                                             #thickness of the shadow network when displaying geometric results

#load the ground truth and test case networks
GT = NWT(GTfile)
T = NWT(Tfile)

#generate point clouds representing both networks
P_T = np.array(T.pointcloud(sigma/subdiv))
P_GT = np.array(GT.pointcloud(sigma/subdiv))

#generate KD trees for each network
GT_tree = sp.spatial.cKDTree(P_GT)
T_tree = sp.spatial.cKDTree(P_T)

#query each KD tree to get the corresponding geometric distances
[T_dist, _] = GT_tree.query(P_T)
[GT_dist, _] = T_tree.query(P_GT)

#convert distances to Gaussian metrics
T_metric = gaussian(T_dist, sigma)
GT_metric = gaussian(GT_dist, sigma)

#calculate the TPR and FPR
print("FNR = " + str(1 - np.mean(GT_metric)))
print("FPR = " + str(1 - np.mean(T_metric)))

plt.subplot(1, 2, 1)
plt.scatter(P_GT[:, 0], P_GT[:, 1], s=sigma*shadow, c="grey")
plt.scatter(P_T[:, 0], P_T[:, 1], s=sigma, c=T_metric, cmap = "plasma")

plt.subplot(1, 2, 2)
plt.scatter(P_T[:, 0], P_T[:, 1], s=sigma*shadow, c="grey")
plt.scatter(P_GT[:, 0], P_GT[:, 1], s=sigma, c=GT_metric, cmap = "plasma")


