# -*- 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

'''
    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):
        self.v0 = 0
        self.v1 = 0
        self.points = []
        self.radii = []
        
    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
'''
    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:
        writeHeader(file, len(vertices), len(edges))
        for i in range(len(vertices)):
            writeVertex(file, vertices[i])
            
        for i in range(len(edges)):
            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.
'''
def importNWT(str):
    with open(str, "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')
        
        nodeList = []
        fiberList = []
        for i in range(numVertex):
            node = readVertex(file)
            nodeList.append(node)

        for i in range(numEdges):
            edge = readFiber(file)
            fiberList.append(edge)
        
    #exportNWT("/home/pavel/Documents/Python/NetLayout/from_python_seg.nwt", nodeList, fiberList)        
    print(str)
    return nodeList, fiberList;
    
'''
Creates a graph from a list of nodes and a list of edges.
Uses edge length as weight.
Returns a NetworkX Object.
'''
 
def createLengthGraph(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].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(nodeList, edgeList):

    lower = nodeList[0].p.copy()
    upper = lower.copy()
    for i in nodeList:
        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(nodeList, edgeList, R=(100, 100, 100)):
    
    #get a list of all node positions in the network
    P = []
    for e in edgeList:
        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 = aabb(nodeList, edgeList)           #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