# -*- coding: utf-8 -*-
"""
Created on Sat Sep 16 16:34:49 2017

@author: pavel
"""

import struct
import numpy as np
import scipy as sp
import networkx as nx
import matplotlib.pyplot as plt
import math
import time
import spharmonics

'''
    Definition of the Node class
    Duplicate of the node class in network
    Stores the physical position, outgoing edges list and incoming edges list.
'''
class Node:
    def __init__(self, point, outgoing, incoming):
        self.p = point
        self.o = outgoing
        self.i = incoming

#    def p():
#        return self.p

'''
    Definition of the Fiber class.
    Duplicate of the Node class in network
    Stores the starting vertex, the ending vertex, the points array and the radius array
'''
class Fiber:
    
       
    def __init__ (self, p1, p2, pois, rads):
        self.v0 = p1
        self.v1 = p2
        self.points = pois
        self.radii = rads
    
    '''
        return the length of the fiber.
    '''        
    def length(self):
        length = 0
        for i in range(len(self.points)-1):
            length = length + math.sqrt(pow(self.points[i][0]- self.points[i+1][0],2) + pow(self.points[i][1]- self.points[i+1][1],2) + pow(self.points[i][2]- self.points[i+1][2],2))

        return length
        
    '''
        returns the turtuosity of the fiber.
    '''    
    def turtuosity(self):
        turtuosity = 0
        distance = math.sqrt(math.pow(self.points[0][0]- self.points[len(self.points)-1][0],2) + math.pow(self.points[0][1]- self.points[len(self.points)-1][1],2) + math.pow(self.points[0][2]- self.points[len(self.points)-1][2],2))
        turtuosity = self.length()/distance
        #print(turtuosity)
        
        return turtuosity
        
    '''
        returns the volume of the fiber.
    '''
    def volume(self):
        volume = 0
        for i in range(len(self.points)-1):
            volume = volume + 1.0/3.0 * math.pi * (math.pow(self.radii[i],2) + math.pow(self.radii[i+1],2) + self.radii[i]*self.radii[i+1]) * math.sqrt(math.pow(self.points[i][0]- self.points[i+1][0],2) + math.pow(self.points[i][1]- self.points[i+1][1],2) + math.pow(self.points[i][2]- self.points[i+1][2],2))

        #print(volume)
        return volume
    
class NWT:   
    
    '''
        Writes the header given and open file descripion, number of verticies and number of edges.
    '''
    def writeHeader(open_file, numVerts, numEdges):
        txt = "nwtFileFormat fileid(14B), desc(58B), #vertices(4B), #edges(4B): bindata"
        b = bytearray()
        b.extend(txt.encode())
        open_file.write(b)
        open_file.write(struct.pack('i', numVerts))
        open_file.write(struct.pack('i', numEdges))
        
    
    '''
        Writes a single vertex to a file.
    '''
    def writeVertex(open_file, vertex):
        open_file.write(struct.pack('<f',vertex.p[0]))
        open_file.write(struct.pack('<f',vertex.p[1]))
        open_file.write(struct.pack('<f',vertex.p[2]))
        open_file.write(struct.pack('i', len(vertex.o)))
        open_file.write(struct.pack('i', len(vertex.i)))
        for j in range(len(vertex.o)):
            open_file.write(struct.pack('i',vertex.o[j]))
            
        for j in range(len(vertex.i)):
            open_file.write(struct.pack('i', vertex.i[j]))    
            
        return
    
    '''
        Writes a single fiber to a file.
    '''
    def writeFiber(open_file, edge):
        open_file.write(struct.pack('i',edge.v0))
        open_file.write(struct.pack('i',edge.v1))
        open_file.write(struct.pack('i',len(edge.points)))
        for j in range(len(edge.points)):
            open_file.write(struct.pack('<f', edge.points[j][0]))
            open_file.write(struct.pack('<f', edge.points[j][1]))
            open_file.write(struct.pack('<f', edge.points[j][2]))
            open_file.write(struct.pack('<f', edge.radii[j]))
            
        return
    
    '''
        Writes the entire network to a file in str given the vertices array and the edges array.
    '''
    def exportNWT(str, vertices, edges):
        with open(str, "wb") as file:
            NWT.writeHeader(file, len(vertices), len(edges))
            for i in range(len(vertices)):
                NWT.writeVertex(file, vertices[i])
                
            for i in range(len(edges)):
                NWT.writeFiber(file, edges[i])
                
        return
    
    
    '''
        Reads a single vertex from an open file and returns a node Object.
    '''
    def readVertex(open_file):
        points = np.tile(0., 3)
        bytes = open_file.read(4)
        points[0] = struct.unpack('f', bytes)[0]
        bytes = open_file.read(4)
        points[1] = struct.unpack('f', bytes)[0]
        bytes = open_file.read(4)
        points[2] = struct.unpack('f', bytes)[0]
        bytes = open_file.read(4)
        
        numO = int.from_bytes(bytes, byteorder='little')
        outgoing = np.tile(0, numO)
        bts = open_file.read(4)
        numI = int.from_bytes(bts, byteorder='little')
        incoming = np.tile(0, numI)
        for j in range(numO):
            bytes = open_file.read(4)
            outgoing[j] = int.from_bytes(bytes, byteorder='little')
            
        for j in range(numI):
            bytes = open_file.read(4)
            incoming[j] = int.from_bytes(bytes, byteorder='little')
            
        node = Node(points, outgoing, incoming)    
        return node
        
        
    '''
        Reads a single fiber from an open file and returns a Fiber object .   
    '''
    def readFiber(open_file):
        bytes = open_file.read(4)
        vtx0 = int.from_bytes(bytes, byteorder = 'little')
        bytes = open_file.read(4)
        vtx1 = int.from_bytes(bytes, byteorder = 'little')
        bytes = open_file.read(4)
        numVerts = int.from_bytes(bytes, byteorder = 'little')
        pts = []
        rads = []
        
        for j in range(numVerts):
            point = np.tile(0., 3)
            bytes = open_file.read(4)
            point[0] = struct.unpack('f', bytes)[0]
            bytes = open_file.read(4)
            point[1] = struct.unpack('f', bytes)[0]
            bytes = open_file.read(4)
            point[2] = struct.unpack('f', bytes)[0]
            bytes = open_file.read(4)
            radius = struct.unpack('f', bytes)[0]
            pts.append(point)
            rads.append(radius)
            
        F = Fiber(vtx0, vtx1, pts, rads)
            
        return F
    
    '''
        Imports a NWT file at location str.
        Returns a list of Nodes objects and a list of Fiber objects.
    '''

class Network:
    
    def __init__(self, filename, clock=False):
        if clock:
            start_time = time.time()
            
        with open(filename, "rb") as file:
            header = file.read(72)
            bytes = file.read(4)
            numVertex = int.from_bytes(bytes, byteorder='little')
            bytes = file.read(4)
            numEdges = int.from_bytes(bytes, byteorder='little')
            
            self.N = []
            self.F = []
            for i in range(numVertex):
                node = NWT.readVertex(file)
                self.N.append(node)
    
            for i in range(numEdges):
                edge = NWT.readFiber(file)
                self.F.append(edge)
        if clock:
            print("Network initialization: "  + str(time.time() - start_time) + "s")
    
    '''
    Creates a graph from a list of nodes and a list of edges.
    Uses edge length as weight.
    Returns a NetworkX Object.
    '''
#    def createLengthGraph(self):
#        G = nx.Graph()
#        for i in range(len(self.nodeList)):
#            G.add_node(i, p=V[i].p)
#        for i in range(len(self.edgeList)):
#            G.add_edge(self.edgeList[i].v0, self.edgeList[i].v1, weight = E[i].length())
#            
#        return G
#    '''
#    Creates a graph from a list of nodes and a list of edges.
#    Uses edge turtuosity as weight.
#    Returns a NetworkX Object.
#    '''    
#    def createTortuosityGraph(nodeList, edgeList):
#        G = nx.Graph()
#        for i in range(len(nodeList)):
#            G.add_node(i, p=V[i].p)
#        for i in range(len(edgeList)):
#            G.add_edge(edgeList[i].v0, edgeList[i].v1, weight = E[i].turtuosity())
#            
#        return G
    
#    '''
#    Creates a graph from a list of nodes and a list of edges.
#    Uses edge volume as weight.
#    Returns a NetworkX Object.
#    '''    
#    def createVolumeGraph(nodeList, edgeList):
#        G = nx.Graph()
#        for i in range(len(nodeList)):
#            G.add_node(i, p=V[i].p)
#        for i in range(len(edgeList)):
#            G.add_edge(edgeList[i].v0, edgeList[i].v1, weight = E[i].volume())
#            
#        return G
#'''
#Returns the positions dictionary for the Circular layout.
#'''    
#def getCircularLayout(graph, dim, radius):
#    return nx.circular_layout(graph, dim, radius)
#
#'''
#Return the positions dictionary for the Spring layout.
#'''    
#def getSpringLayout(graph, pos, iterations, scale):
#    return nx.spring_layout(graph, 2, None, pos, iterations, 'weight', scale, None)
#        
#'''
#Draws the graph.
#'''        
#def drawGraph(graph, pos):
#    nx.draw(graph, pos)
#    return

    def aabb(self):
    
        lower = self.N[0].p.copy()
        upper = lower.copy()
        for i in self.N:
            for c in range(len(lower)):
                if lower[c] > i.p[c]:
                    lower[c] = i.p[c]
                if upper[c] < i.p[c]:
                    upper[c] = i.p[c]
        return lower, upper
    
    #calculate the distance field at a given resolution
    #   R (triple) resolution along each dimension
    def distancefield(self, R=(100, 100, 100)):
        
        #get a list of all node positions in the network
        P = []
        for e in self.F:
            for p in e.points:
                P.append(p)
                
        #turn that list into a Numpy array so that we can create a KD tree
        P = np.array(P)
        
        #generate a KD-Tree out of the network point array
        tree = sp.spatial.cKDTree(P)
        
        plt.scatter(P[:, 0], P[:, 1])
        
        #specify the resolution of the ouput grid
        R = (200, 200, 200)
        
        #generate a meshgrid of the appropriate size and resolution to surround the network
        lower, upper = self.aabb(self.N, self.F)    #get the space occupied by the network
        x = np.linspace(lower[0], upper[0], R[0])   #get the grid points for uniform sampling of this space
        y = np.linspace(lower[1], upper[1], R[1])
        z = np.linspace(lower[2], upper[2], R[2])
        X, Y, Z = np.meshgrid(x, y, z)
        #Z = 150 * numpy.ones(X.shape)
        
        
        Q = np.stack((X, Y, Z), 3)
        
        
        D, I = tree.query(Q)
        
        return D
    
    #returns the number of points in the network
    def npoints(self):                              
        n = 0                                       #initialize the counter to zero
        for f in self.F:                            #for each fiber
            n = n + len(f.points) - 2               #count the number of points in the fiber - ignoring the end points
        n = n + len(self.N)                         #add the number of nodes (shared points) to the node count
        return n                                    #return the number of nodes
    
    #returns all of the points in the network
    def points(self):
        k = self.npoints()
        P = np.zeros((3, k))                        #allocate space for the point list
        
        idx = 0
        for f in self.F:                            #for each fiber in the network
            for ip in range(1, len(f.points)-1):    #for each point in the network
                P[:, idx] = f.points[ip]            #store the point in the raw point list
                idx = idx + 1
        return P                                    #return the point array        
    
    #returns the number of linear segments in the network
    def nsegments(self):
        n = 0                                       #initialize the segment counter to 0
        for f in self.F:                            #for each fiber
            n = n + len(f.points) - 1               #calculate the number of line segments in the fiber (points - 1)
        return n                                    #return the number of line segments
    
    #return a list of line segments representing the network
    def segments(self, dtype=np.float32):
        k = self.nsegments()                        #get the number of line segments
        start = np.zeros((k, 3),dtype=dtype)                    #start points for the line segments
        end = np.zeros((k, 3), dtype=dtype)                      #end points for the line segments
        
        idx = 0                                     #initialize the index counter to zero
        for f in self.F:                            #for each fiber in the network
            for ip in range(0, len(f.points)-1):    #for each point in the network
                start[idx, :] = f.points[ip]            #store the point in the raw point list
                idx = idx + 1
        
        idx = 0
        for f in self.F:                            #for each fiber in the network
            for ip in range(1, len(f.points)):      #for each point in the network
                end[idx, :] = f.points[ip]            #store the point in the raw point list
                idx = idx + 1
                
        return start, end
    
    #function returns the fiber associated with a given 1D line segment index
    def segment2fiber(self, idx):        
        i = 0
        for f in range(len(self.F)):                #for each fiber in the network
            i = i + len(self.F[f].points)-1         #add the number of points in the fiber to i
            if i > idx:                             #if we encounter idx in this fiber
                return self.F[f].points, f          #return the fiber associated with idx and the index into the fiber array
        
    def vectors(self, clock=False, dtype=np.float32):
        if clock:
            start_time = time.time()
        start, end = self.segments(dtype)                #retrieve all of the line segments
        v = end - start                             #calculate the resulting vectors
        l = np.sqrt(v[:, 0]**2 + v[:,1]**2 + v[:,2]**2) #calculate the fiber lengths
        z = l==0                                    #look for any zero values
        nz = z.sum()
        if nz > 0:
            print("WARNING: " + str(nz) + " line segment(s) of length zero were found in the network and will be removed" )
            
        if clock:
            print("Network::vectors: " + str(time.time() - start_time) + "s")
            
        return np.delete(v, np.where(z), 0)
    
    #scale all values in the network by tuple S = (sx, sy, sz)
    def scale(self, S):
        for f in self.F:
            for p in f.points:
                p[0] = p[0] * S[0]
                p[1] = p[1] * S[1]
                p[2] = p[2] * S[2]
                
        for n in self.N:
            n.p[0] = n.p[0] * S[0]
            n.p[1] = n.p[1] * S[1]
            n.p[2] = n.p[2] * S[2]
        
    
    #calculate the adjacency weighting function for the network given a set of vectors X = (x, y, z) and weight exponent k
    def adjacencyweight(self, P, k=200, length_threshold = 25, dtype=np.float32):
        V = self.vectors(dtype)                                                 #get the vectors representing each segment
        #V = V[0:n_vectors, :]
        L = np.expand_dims(np.sqrt((V**2).sum(1)), 1)                           #calculate the length of each vector
        
        outliers = L > length_threshold                                         #remove outliers based on the length_threshold
        V = np.delete(V, np.where(outliers), 0)
        L = np.delete(L, np.where(outliers))
        V = V/L[:,None]                                                         #normalize the vectors
        
        P = np.stack(spharmonics.sph2cart(1, P[0], P[1]), P[0].ndim)        
        PV = P[...,None,:] * V
        cos_alpha = PV.sum(PV.ndim-1)
        W = np.abs(cos_alpha) ** k

        return W, L