# -*- coding: utf-8 -*-
"""
Created on Mon Mar 19 14:45:39 2018

@author: david
"""

import nwt
import matplotlib.pyplot as plt
import numpy


class fibernet:
    
    def __init__(self, filename):
        self.nwt(filename)
        self.buildgraph()
    
    #load a network from a NWT file
    def nwt(self, filename):
        net = nwt.NWT(filename)
        
        #first insert the points where fibers are connected
        self.P = []
        for vi in range(0, len(net.v)):                                         #store each vertex in the point list
            self.P.append( (net.v[vi].p[0], net.v[vi].p[1], net.v[vi].p[2]) )
        
        #now insert each fiber, connecting them appropriately using previously inserted points (vertices)
        self.F = []
        for ei in range(0, len(net.e)):
            f = [numpy.int(net.e[ei].v[0])]                                     #create a new fiber to store point indices, initialize with the first vertex ID
            for pi in range(1, net.e[ei].p.shape[0]):
                self.P.append( (net.e[ei].p[pi,0], net.e[ei].p[pi,1], net.e[ei].p[pi,2]))    #append the current point to the global point list
                f.append(len(self.P) - 1)                                       #append the current point ID
            f.append(numpy.int(net.e[ei].v[1]))                                 #append the last vertex ID
            self.F.append(f)
    
    #build a graph from the available geometry
    def buildgraph(self):
        p_to_v = numpy.ones((self.npoints()), numpy.int) * (-1)                            #create an array that maps a point ID to a vertex ID
        
        self.V = []                                                                  #create an empty vertex list
        self.E = []                                                                  #create an empty edge list
        for fi in range(0, self.nfibers()):                                     #for each fiber in the network
            if p_to_v[self.F[fi][0]] == -1:                                          #if the first point hasn't been encountered
                p_to_v[self.F[fi][0]] = len(self.V)                                       #add it to the mapping
                self.V.append( (self.F[fi][0], []) )                                              #add a new vertex representing this point
            if p_to_v[self.F[fi][-1]] == -1:
                p_to_v[self.F[fi][-1]] = len(self.V)
                self.V.append( (self.F[fi][-1], []) )
                
            v0 = p_to_v[self.F[fi][0]]                                          #get the two vertex indices corresponding to this edge
            v1 = p_to_v[self.F[fi][-1]]
            
            self.E.append( (v0, v1) )                                                #create an edge representing this fiber and add it to the edge list
            self.V[v0][1].append(fi)                                                 #add this edge ID to the corresponding vertices to update the adjacency list
            self.V[v1][1].append(fi)
            
    #return the length of a fiber specified by fiber ID fid
    def length(self, fid):
        p = self.points(fid)                                                 #get the points corresponding to the fiber
        l = 0                                                                   #initialize the length to zero
        for pi in range(1, len(p)):                                             #for each line segment
            v = p[pi] - p[pi-1]                                                 #calculate the vector representing the line segment
            l = l + numpy.linalg.norm(numpy.array(v))                                  #add the length of the line segment
        return l                                                                #return the total length   

    #return a list of vertices adjacent to the vertex specified by vid
    def adjacent(self, vid):
        adj = []                                                                #initialize an empty adjacency list
        for e in self.V[vid][1]:                                                #for each edge connected to vid
            if self.E[e][0] != vid:
                adj.append(self.E[e][0])
            else:
                adj.append(self.E[e][1])                
        return adj
    
    #return the fiber network graph as an adjacency list
    def adjacencylist(self):
        A = []
        for vi in range(0, len(self.V)):                                                        #for each vertex in the graph
            A.append( self.adjacent(vi) )                                       #add the adjacency list for the vertex into the main adjacency list
        return A

    #return the fiber network graph as an adjacency matrix
    # possible values for the matrix include:
    #           binary - 1 when there is a connection, 0 otherwise
    #           edgeID - stores the edge ID corresponding to a connection, and -1 otherwise
    #           length - stores the length of the connecting fiber, and -1 if there is no connection
    def adjacencymatrix(self, value="binary"):
        if value == "binary":
            M = numpy.ones( (len(self.V), len(self.V)), numpy.bool )
        elif value == "edgeID":
            M = numpy.ones( (len(self.V), len(self.V)), numpy.int ) * -1
        elif value == "length":
            M = numpy.ones( (len(self.V), len(self.V)) ) * -1
        else:
            return []
            
        for vi in range(0, len(self.V)):                                        #for each vertex in the graph
            A = self.adjacent(vi)                                               #get the adjacency matrix
            for a in A:
                if value == "binary":
                    M[vi,a] = M[a,vi] = 1
                if value == "edgeID":
                    M[vi,a] = M[a,vi] = self.findedge(a, vi)
                if value == "length":
                    M[vi,a] = M[a,vi] = self.length(self.findedge(a,vi))
        return M
    
    #find the edge ID corresponding to two vertices
    def findedge(self, v0, v1):
        
        edges = self.V[v0][1]                                                   #get the list of candidate edges connected to v0
        for e in edges:                                                         #for each edge
            if self.E[e][0] == v1 or self.E[e][1] == v1:
                return e
        
        return -1                                                               #return -1 if the edge doesn't exist
    
    #get the points associated with a specified fiber ID fid    
    def points(self, fid):                                                   #get the points corresponding to a specified fiber index
        f = self.F[fid]                                                         #get the list of point indices
        return numpy.array(self.P)[f]
    
    #get the total number of points in the network
    def npoints(self):
        return len(self.P)
    
    #get the total number of fibers in the network
    def nfibers(self):
        return len(self.F)
    
    #plot the network as a 3D line plot
    def plot(self):
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        
        for fi in range(0, self.nfibers()):
            test = fibers.points(fi)
            ax.plot(test[:, 0], test[:, 1], test[:, 2])
        return fig