"""

Authors: Kristian Cho and Gillian Minnehan
Date: 11/28/20

"""
import numpy as np
import csv
from pca import pca
import utils
from sklearn import neighbors
import time
import matplotlib.pyplot as plt
from itertools import combinations
from hw4_icp import icp, icp2
from sklearn.preprocessing import normalize

def step(si, fi):
    """

    Returns either 0 or 1 depending on si and fi values

    si - center of the definition interval of fi
        (0 for alpha, phi, theta and 1 for d)
    fi - feature

    """
    if fi < si:
        return 0
    return 1

def calculateBin(alpha, phi, theta, d, r):
    """

    Returns the bin index for the pair of points

    alpha - np.dot(v, nt)
    phi - np.dot(u, (pt - ps) / d)
    theta - np.cross(u, (pt - ps) / d)
    d - np.linalg.norm(pt - ps)
    r - max distance in k neighbors

    """
    idx = 0
    idx += step(0, alpha)
    idx += step(r, d) * 2 #TODO: step paramater should be r, not 0
    idx += step(0, phi) * 4
    idx += step(0, theta) * 8
    return idx

def fcalculateBin(alpha, phi, theta):
    """

    Returns the bin index for the pair of points

    alpha - np.dot(v, nt)
    phi - np.dot(u, (pt - ps) / d)
    theta - np.cross(u, (pt - ps) / d)
    d - np.linalg.norm(pt - ps)
    r - max distance in k neighbors

    """
    idx = 0
    idx += step(0, alpha)
    idx += step(0, phi) * 2
    idx += step(0, theta) * 4
    return idx

def getAngle(v1, v2):
    """

    Returns the angle between two vectors in radians

    v1 - 3 x 1 numpy matrix
    v2 - 3 x 1 numpy matrix

    """
    uv1 = v1 / np.linalg.norm(v1)
    uv2 = v2 / np.linalg.norm(v2)
    dot_product = np.dot(uv2, uv1)
    angle = np.arccos(dot_product).item()
    return angle

def assignRoles(normA, normB, a, b):
    """

    Returns the assignments for the source and target points/normals
    
    normA - normal associated with point A
    normB - normal associated with point B
    a - point A
    b - point B

    """
    line = np.asmatrix(np.expand_dims(a - b, axis=0)) # line between two points
    aAlph = getAngle(normA, line)
    bAlph = getAngle(normB, line)
    if(aAlph < bAlph):
        return a, b, normA, normB
    else:
        return b, a, normB, normA

def getNeighborPoints(P, ind):
    """

    Returns a matrix of all the neighbor points

    P - a list of numpy 3 x 1 matrices that represent the points
    ind - a list of integers representing the indicies of the 
        neightbor points

    """
    pts = []
    for i in ind:
        pts.append(P[i])
    return pts

def load_pc(filename):
    """
    Load a csv PC.

    Loads a point cloud from a csv file and returns 
    a list of 3 by 1 numpy arrays that represent the points.

    filename - a string containing the files name.

    """
    pc = []
    with open(filename, 'r') as file:
        reader = csv.reader(file)
        for row in reader:
            point = np.array([float(x) for x in row])
            if len(point) == 0:
                continue
            pc.append(point)
    return pc

def load_sydney_obj_pc(filename):
    """
    Load a csv PC from the sydney urban objects dataset.

    Loads a point cloud from a csv file and returns 
    a list of 3 by 1 numpy arrays that represent the points.

    filename - a string containing the files name.

    """
    pc = []
    with open(filename, 'r') as file:
        reader = csv.reader(file)
        for row in reader:
            point = np.array([float(x) for x in row[3:6]])
            if len(point) == 0:
                continue
            pc.append(point)

    return pc

def pfh(P, k):
    """

    Returns the signatures for point cloud P using the point feature
    histogram method

    P - 1 x 3 list of points

    """
    #start = time.clock()
    # preprocessing
    tree = neighbors.KDTree(np.array(P)) # TODO: specify leaf_size?
    KNeighbors = []
    norms = []
    # experimenting with using r
    radii = []
    r = 0.2
    numCount = 0
    for i in range(len(P)):
        # get k nearest neighbors
        dist, ind = tree.query([P[i]], k=k+1) # TODO: change to radius of neighbors
        ind = np.delete(ind, 0, 1)
        ind = ind.reshape(ind.shape[1])
        radii.append(dist[0][-1])
        # experimenting begins here
        # ind, dist = tree.query_radius([P[i]], r=r, return_distance=True, sort_results=True)
        # numCount += len(list(combinations(ind[0], 2)))
        # ind = ind[0][1:]
        # experimenting ends here

        KNeighbors.append(ind)
        neighborPts = getNeighborPoints(P, ind)

        # find normal of neighbors
        newPc, vs, meanVec = pca(neighborPts, False)
        norm = np.asmatrix(np.expand_dims(np.cross(vs[0, :], vs[1, :]), axis=1)) # TODO: test with plotting
        norms.append(norm)
        neighborPtsMat = []
        # Uncomment to plot neighbors and fitted plane
        # if(i == 5): # modify to plot plane for a specific point
        #     fig = utils.view_pc_two([neighborPts], color='r')
        #     utils.draw_plane(fig, norm, meanVec, color=(0, 1, 0, 0.3), length=[-100,100], width=[-100,100])
        #     return

    # computing signatures
    currentIndex = 0
    count = 0
    allSignatures = []
    numCombinations = len(list(combinations(KNeighbors[0], 2)))
    for h in KNeighbors: # number of points
        listSig = [0] * 16
        for j in range(k):  # number of neighbors
            for i in range(j + 1, k):
                a = P[h[j]]
                b = P[h[i]]
                ps, pt, ns, nt = assignRoles(norms[h[j]], norms[h[i]], a, b)
                d = np.linalg.norm(pt - ps)
                u = ns.reshape(3) # norm
                v = np.cross(pt - ps, u)
                if (np.dot(v - ps, ns) / np.linalg.norm(v - a)) < 0:
                    u *= -1
                    v = np.cross(pt - ps, u)
                w = np.cross(u,v)
                alpha = np.dot(v,nt)
                phi = np.dot(u, (pt - ps) / d)
                theta = np.arctan2([np.dot(w, nt)], [np.dot(u, nt)])
                # bin calculation
                bin = calculateBin(alpha, phi, theta, d, radii[currentIndex])
                #print("point ", h[j], " and ", h[i], ": bin#",bin)
                listSig[bin] += 1
        currentIndex += 1
        listSig = [[x / numCombinations] for x in listSig]
        oneSignature = np.matrix(listSig)
        allSignatures.append(oneSignature)
    #end = time.clock()
    #print("Time: ", end - start)
    #print("Signatures: ", allSignatures)
    # Uncomment to plot signatures in bar graph
    # fig, ax = plt.subplots()
    # axes = plt.gca()
    # axes.set_ylim([0, 1])
    # xaxis = np.arange(16)
    # ax.bar(xaxis, signatures, 0.35)
    # ax.set_xticks(xaxis)
    # plt.show()

    return allSignatures

def fpfh(P, k):
    """

    Returns the signatures for point cloud P using the fast point feature
    histogram method

    P - 1 x 3 list of points

    """
    # preprocessing
    tree = neighbors.KDTree(np.array(P)) # TODO: specify leaf_size?
    KNeighbors = []
    norms = []
    allSPFH = []
    allDist = []
    # calculating Simplified Point Feature Histograms
    for i in range(len(P)):
        # get k nearest neighbors
        dist, ind = tree.query([P[i]], k=k+1) # TODO: change to radius of neighbors
        dist = np.delete(dist, 0, 1)
        dist = dist.reshape(dist.shape[1])
        ind = np.delete(ind, 0, 1)
        ind = ind.reshape(ind.shape[1])
        KNeighbors.append(ind)
        allDist.append(dist)
        neighborPts = getNeighborPoints(P, ind)

        # find normal of neighbors
        newPc, vs, meanVec = pca(neighborPts, False)
        norm = np.asmatrix(np.expand_dims(np.cross(vs[0, :], vs[1, :]), axis=1)) # TODO: test with plotting
        norms.append(norm)
    
    currentIndex = 0
    numCombinations = len(list(combinations(KNeighbors[0], 2)))
    for h in KNeighbors:
        listSig = [0] * 8
        ps = P[currentIndex]
        ns = norms[currentIndex]
        for j in range(k):
            pt = P[h[j]]
            nt = norms[h[j]]
            d = np.linalg.norm(pt - ps)
            u = ns.reshape(3) # norm
            v = np.cross(pt - ps, u)
            if (np.dot(v - ps, ns) / np.linalg.norm(v - ps)) < 0:
                u *= -1
                v = np.cross(pt - ps, u)
            w = np.cross(u,v)
            alpha = np.dot(v,nt)
            phi = np.dot(u, (pt - ps) / d)
            theta = np.arctan2([np.dot(w, nt)], [np.dot(u, nt)])
            bin = fcalculateBin(alpha, phi, theta)
            listSig[bin] += 1
        currentIndex += 1
        listSig = [x / k for x in listSig]
        # oneSignature = np.matrix(listSig)
        allSPFH.append(listSig)
    
    currentIndex = 0
    allFPFH = []
    for i in range(len(P)):
        currentSig = [0] * 8
        for j in range(k):
            currentSum = [x * (1 / allDist[i][j]) for x in allSPFH[KNeighbors[i][j]]]
            currentSig = [x + y for x,y in zip(currentSig, currentSum)]
        currentSig = [(1 / k) * x for x in currentSig]
        currentSig = [x + y for x,y in zip(currentSig, allSPFH[i])]
        currentSig = [[x] for x in currentSig]
        currentSig = np.matrix(currentSig)
        currentSig = normalize(currentSig, axis=0, norm='l1')
        allFPFH.append(np.matrix(currentSig))
    
    return allFPFH

def main():
    # Import the cloud
    #pc_source = load_sydney_obj_pc('/home/parallels/Documents/eecs498/final_project/data/sydney-urban-objects-dataset/objects/trash.1.2738.csv')
    pc_source = load_pc('/home/kristian/Desktop/EECS 498/EECS498Final/data/hw4/cloud_icp_source.csv')
    #utils.view_pc_two([pc_source]) # original
    pc_target = load_pc('/home/kristian/Desktop/EECS 498/EECS498Final/data/hw4/cloud_icp_target0.csv')
    print("Number of points in source: ", len(pc_source))
    pc_source1 = utils.load_pc('/home/kristian/Desktop/EECS 498/EECS498Final/data/hw4/cloud_icp_source.csv')
    pc_target1 = utils.load_pc('/home/kristian/Desktop/EECS 498/EECS498Final/data/hw4/cloud_icp_target0.csv')
    # Uncomment to use test data
    # rng = np.random.RandomState(0)
    # X = rng.random_sample((500, 3))

    # Perform point feature histogram method
    start = time.clock()
    sourceSig = pfh(pc_source, 10) # Modify number of neighbors
    targetSig = pfh(pc_target, 10)
    print("PFH ICP: ")
    icp(pc_source1, pc_target1, sourceSig, targetSig, 0.05, False)
    end = time.clock()
    print("PFH Time: ", end - start)

    start = time.clock()
    fsourceSig = fpfh(pc_source, 10)
    ftargetSig = fpfh(pc_target, 10)
    print("FPFH ICP: ")
    icp(pc_source1, pc_target1, fsourceSig, ftargetSig, 0.05, True)
    end = time.clock()
    print("FPFH Time: ", end - start)

    # pc_source1 = utils.load_pc('/home/kristian/Desktop/EECS 498/EECS498Final/data/hw4/cloud_icp_source.csv')
    # pc_target1 = utils.load_pc('/home/kristian/Desktop/EECS 498/EECS498Final/data/hw4/cloud_icp_target1.csv')
    # pc_source, errors = icp(pc_source, pc_target, sourceSig, targetSig, 0.05)

    # end = time.clock()
    # print("PFH Time: ", end - start)
    # start = time.clock()
    # icp2(pc_source, pc_target, 0.05)
    # end = time.clock()
    # print("Normal ICP Time: ", end - start)


    # Plot aligned point cloud and target point cloud
    # utils.view_pc_two([pc_source, pc_target], None, ['b', 'r'], ['o', '^'])
    # plt.axis([-0.15, 0.15, -0.15, 0.15])

    # Plot error vs. num iterations
    # fig = plt.figure()
    # ax = fig.add_subplot(1,1,1)
    # xdata = np.arange(0,len(errors))
    # ax.plot(xdata, errors, color='tab:blue')
    # ax.set_xlabel('num iterations')
    # ax.set_ylabel('error')
    # ax.set_title('Error v. Iterations for ICP')

    input("Press enter to end:")

if __name__ == '__main__':
    main()