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

def brute_force_spherical(queries, supports, radius):

    neighborhoods = [[] for _ in range(queries.shape[0])]

    for ii in range(queries.shape[0]):
        dists = np.linalg.norm(queries[ii,:]-supports, axis=1)
        neighborhoods[ii] = np.where(dists<radius)

    return neighborhoods


def brute_force_KNN(queries, supports, k):

    neighborhoods = [[] for _ in range(queries.shape[0])]

    for ii in range(queries.shape[0]):
        dists = np.linalg.norm(queries[ii,:]-supports, axis=1)
        neighborhoods[ii] = np.argsort(dists)[:k]

    return neighborhoods


if __name__ == '__main__':

    # Load point cloud
    # ****************
    #
    #   Load the file '../data/indoor_scan.ply'
    #   (See read_ply function)
    #

    # Path of the file
    file_path = '../data/indoor_scan.ply'

    # Load point cloud
    data = read_ply(file_path)

    # Concatenate data
    points = np.vstack((data['x'], data['y'], data['z'])).T

    # Brute force neighborhoods
    # *************************
    #

    BRUTE = False 

    # If statement to skip this part if you want
    if BRUTE:

        # Define the search parameters
        neighbors_num = 100
        radius = 0.2
        num_queries = 10

        # Pick random queries
        random_indices = np.random.choice(points.shape[0], num_queries, replace=False)
        queries = points[random_indices, :]

        # Search spherical
        t0 = time.time()
        neighborhoods = brute_force_spherical(queries, points, radius)
        t1 = time.time()

        # Search KNN      
        neighborhoods = brute_force_KNN(queries, points, neighbors_num)
        t2 = time.time()

        # Print timing results
        print('{:d} spherical neighborhoods computed in {:.3f} seconds'.format(num_queries, t1 - t0))
        print('{:d} KNN computed in {:.3f} seconds'.format(num_queries, t2 - t1))

        # Time to compute all neighborhoods in the cloud
        total_spherical_time = points.shape[0] * (t1 - t0) / num_queries
        total_KNN_time = points.shape[0] * (t2 - t1) / num_queries
        print('Computing spherical neighborhoods on whole cloud : {:.0f} hours'.format(total_spherical_time / 3600))
        print('Computing KNN on whole cloud : {:.0f} hours'.format(total_KNN_time / 3600))

    # KDTree neighborhoods
    # ********************
    #

    # If statement to skip this part if wanted
    else:
        print(points.shape[0])
        radiuses = np.arange(0.0, 0.3, 0.01)
        times = []
        for radius in radiuses:

            # Define the search parameters
            num_queries = 1000

            # Pick random queries
            random_indices = np.random.choice(points.shape[0], num_queries, replace=False)
            queries = points[random_indices, :]

            # Create KDTree
            kdtree = KDTree(points, leaf_size=100)

            # Search spherical
            t0 = time.time()
            neighborhoods = kdtree.query_radius(queries, r=radius)
            t1 = time.time()

            # Print timing results
            print('{:d} spherical neighborhoods computed in {:.3f} seconds'.format(num_queries, t1 - t0))

            # Time to compute all neighborhoods in the cloud
            total_spherical_time = points.shape[0] * (t1 - t0) / num_queries
            print('Computing spherical neighborhoods on whole cloud : {:.0f} hours'.format(total_spherical_time / 3600))

            times.append(t1-t0)
        plt.plot(radiuses, times)
        plt.show()