from pyntcloud import PyntCloud
import glob
import sys
import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
# from scipy.spatial import KDTree
from sklearn.neighbors import KDTree




def set_proper_aspect_ratio(ax):
    extents = np.array([getattr(ax, 'get_{}lim'.format(dim))() for dim in 'xyz'])
    sz = extents[:,1] - extents[:,0]
    centers = np.mean(extents, axis=1)
    maxsize = max(abs(sz))
    r = maxsize/2
    for ctr, dim in zip(centers, 'xyz'):
        getattr(ax, 'set_{}lim'.format(dim))(ctr - r, ctr + r)

clouds = glob.glob(os.path.join(sys.argv[1], '*.ply'))

#clouds = clouds[::100]

azim = 45
cloud_num = 0
fig = plt.figure(figsize=(10, 10))
plt.style.use('dark_background')

SMALL_SIZE = 18
MEDIUM_SIZE = 24
BIGGER_SIZE = 36

plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title

nn_stats = np.zeros((len(clouds), 3))
for ind, cloud in enumerate(clouds):
    # print(cloud)
    p = PyntCloud.from_file(cloud)

    if p.xyz.shape[0] > 2:

        kdtree = KDTree(p.xyz, leaf_size=2)
        nndist, nn_points = kdtree.query(p.xyz, 2)
        print(np.median(nndist[:, 1]))

    else:
        print('no points :(')

    nn_stats[ind, :] = np.array([ind, p.xyz.shape[0], np.median(nndist[:, 1])])

    ptp = p.xyz.ptp()

    ax = plt.axes(projection='3d')
    ax.view_init(elev=20, azim=azim % 360)
    fig.set_facecolor('black')
    ax.set_facecolor('black')
    # First remove fill
    ax.xaxis.pane.fill = False
    ax.yaxis.pane.fill = False
    ax.zaxis.pane.fill = False
    #azim = azim+0.25
    cloud_num += 1

    ax.scatter(
        p.xyz[:, 0],
        p.xyz[:, 1],
        p.xyz[:, 2],
        marker="D",
        facecolors=np.repeat([[255, 125, 0]], p.xyz.shape[0], axis=0) / 255,
        zdir="z",
        depthshade=True,
        s=ptp / 10)

    #set_proper_aspect_ratio(ax)
    ax.set_xlabel('  (mm)  ',labelpad=20)
    ax.set_ylabel('  (mm)  ',labelpad=20)
    ax.set_zlabel('  (mm)  ',labelpad=20)
    ax.set_xlim((-150, 150))
    ax.set_ylim((-150, 150))
    ax.set_zlim((0, 300))
    #plt.show()
    plt.savefig('E:\\EyeRISTransectExport\\cloud_movies_8_17\\' + os.path.basename(cloud)[:-4] + '.png', format='png')

np.savetxt(sys.argv[1].rstrip('\\') +'_nn_stats.csv', nn_stats)