#
#
#      0=============================0
#      |    TP4 Point Descriptors    |
#      0=============================0
#
#
# ------------------------------------------------------------------------------------------
#
#      Script of the practical session
#
# ------------------------------------------------------------------------------------------
#
#      Hugues THOMAS - 13/12/2017
#


import numpy as np
from matplotlib import pyplot as plt
from sklearn.neighbors import KDTree
from ply import write_ply, read_ply
import time

import matplotlib.pyplot as plt

# ------------------------------------------------------------------------------------------
#
#           Functions
#       \***************/
#
#



def PCA(points):

    # compute centroid
    centroid = points.mean(axis=0)
    # compute centered PC
    points = points - centroid[None, :]
    # compute covariance matrix
    cov = 1/points.shape[0] * points.T @ points  
    
    # return eigenvalues and eigenvectors
    return np.linalg.eigh(cov)


def compute_local_PCA(query_points, cloud_points, radius, k=-1):

    # This function needs to compute PCA on the neighborhoods of all query_points in cloud_points

    # create a kdtree for more efficient queries
    tree = KDTree(cloud_points)

    # query the neighbors for each point in query_points
    if k<=0: # radius method
        neighbors_list = tree.query_radius(query_points, radius)
    else: # k-nearest-neighbors method 
        neighbors_list = tree.query(query_points, k, return_distance=False)

    all_eigenvalues = []
    all_eigenvectors = []

    # lengths = []
    for neighbors in neighbors_list:
        # lengths.append(len(neighbors))
        eigvals, eigvects = PCA(cloud_points[neighbors])
        all_eigenvalues.append(eigvals)
        all_eigenvectors.append(eigvects)

    # plt.hist(lengths, 100, density=True, facecolor='g', alpha=0.75)
    # plt.show()

    all_eigenvalues = np.stack(all_eigenvalues)
    all_eigenvectors = np.stack(all_eigenvectors)

    return all_eigenvalues, all_eigenvectors


def compute_features(query_points, cloud_points, radius):

    # compute eigenvalues and eigenvectors
    eigvals, eigvects = compute_local_PCA(query_points, cloud_points, radius)

    # compute the features
    eps = 1e-8
    normals = eigvects[:, :, 0]
    verticality = 2 * np.arcsin(np.abs(normals[:, 2])) / np.pi
    linearity = 1 - eigvals[:, 1] / (eigvals[:, 2]+eps)
    planarity = (eigvals[:, 1] - eigvals[:, 0]) / (eigvals[:, 2]+eps)
    sphericity = eigvals[:, 0] / (eigvals[:, 2]+eps)

    return verticality, linearity, planarity, sphericity


# ------------------------------------------------------------------------------------------
#
#           Main
#       \**********/
#
# 
#   Here you can define the instructions that are called when you execute this file
#

if __name__ == '__main__':

    # PCA verification
    # ****************
    if False:

        # Load cloud as a [N x 3] matrix
        cloud_path = '../data/Lille_street_small.ply'
        cloud_ply = read_ply(cloud_path)
        cloud = np.vstack((cloud_ply['x'], cloud_ply['y'], cloud_ply['z'])).T

        # Compute PCA on the whole cloud
        eigenvalues, eigenvectors = PCA(cloud)

        # Print your result
        print(eigenvalues)

        # Expected values :
        #
        #   [lambda_3; lambda_2; lambda_1] = [ 5.25050177 21.7893201  89.58924003]
        #
        #   (the convention is always lambda_1 >= lambda_2 >= lambda_3)
        #

		
    # Normal computation
    # ******************
    if False:

        # Load cloud as a [N x 3] matrix
        cloud_path = '../data/Lille_street_small.ply'
        cloud_ply = read_ply(cloud_path)
        cloud = np.vstack((cloud_ply['x'], cloud_ply['y'], cloud_ply['z'])).T

        # Compute PCA on the whole cloud
        all_eigenvalues, all_eigenvectors = compute_local_PCA(cloud, cloud, 0.50)
        # all_eigenvalues, all_eigenvectors = compute_local_PCA(cloud, cloud, 0.50, k=30)
        normals = all_eigenvectors[:, :, 0]

        # Save cloud with normals
        write_ply('../Lille_street_small_normals.ply', (cloud, normals), ['x', 'y', 'z', 'nx', 'ny', 'nz'])


    # Features computation
    # ******************
    if True:

        # Load cloud as a [N x 3] matrix
        cloud_path = '../data/Lille_street_small.ply'
        cloud_ply = read_ply(cloud_path)
        cloud = np.vstack((cloud_ply['x'], cloud_ply['y'], cloud_ply['z'])).T

        # Compute features on the whole cloud
        verticality, linearity, planarity, sphericity = compute_features(cloud, cloud, 0.50)

        # Save cloud with features
        write_ply('../Lille_street_small_features.ply', [cloud, verticality, linearity, planarity, sphericity], ['x', 'y', 'z', 'v', 'l', 'p', 's'])