#
#
#      0===========================================================0
#      |              TP4 Surface Reconstruction                   |
#      0===========================================================0
#
#
# ------------------------------------------------------------------------------------------
#
#      Jean-Emmanuel DESCHAUD - 02/02/2018
#


# Import numpy package and name it "np"
import numpy as np

# Import functions from scikit-learn
from sklearn.neighbors import KDTree

# Import functions to read and write ply files
from ply import write_ply, read_ply

# Import time package
import time

from skimage import measure

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

import trimesh


# Hoppe surface reconstruction
def compute_hoppe(points,normals,scalar_field,grid_resolution,min_grid,size_voxel):

    # create the grid
    x_ = np.arange(grid_resolution + 1)
    y_ = x_.copy()
    z_ = x_.copy()
    x, y, z = np.meshgrid(x_, y_, z_, indexing='ij')
    grid = np.stack([x, y, z])
    grid = grid.reshape((3,len(x)**3)).T # reshape grid as Nx3 (necessary for tree search)
    grid = min_grid + size_voxel * grid

    # create kd tree to query more efficiently
    tree = KDTree(points)
    neighbors = tree.query(grid, return_distance=False)
    neighbors = neighbors.flatten()

    # compute hoppe function
    grid = np.sum(normals[neighbors]*(grid - points[neighbors]), axis=1)
    grid = grid.reshape((len(x), len(x), len(x))) # reshape grid as (res,res,res)

    scalar_field[:] = grid[:]
    return
    

# IMLS surface reconstruction
def compute_imls(points,normals,scalar_field,grid_resolution,min_grid,size_voxel,knn):

    # create the grid
    x_ = np.arange(grid_resolution + 1)
    y_ = x_.copy()
    z_ = x_.copy()
    x, y, z = np.meshgrid(x_, y_, z_, indexing='ij')
    grid = np.stack([x, y, z])
    grid = grid.reshape((3,len(x)**3)).T # reshape grid as Nx3 (necessary for tree search)
    grid = min_grid + size_voxel * grid

    # create kd tree to query more efficiently
    tree = KDTree(points)
    neighbors = tree.query(grid, k=knn, return_distance=False)

    # compute imls function
    h = 0.001
    diff = grid[:,None,:] - points[neighbors]
    theta = np.exp(-np.linalg.norm(diff, axis=2)**2/h**2)
    theta = np.clip(theta, 1e-8, np.inf)
    num = np.sum( np.sum(normals[neighbors] * diff, axis=2) * theta, axis=1 )
    denom = np.sum(theta, axis = 1)
    grid = num/denom
    grid = grid.reshape((len(x), len(x), len(x))) # reshape grid as (res,res,res)

    scalar_field[:] = grid[:]
    return
    

if __name__ == '__main__':

    t0 = time.time()
    
    # Path of the file
    file_path = '../data/bunny_normals.ply'

    # Load point cloud
    data = read_ply(file_path)

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

	# Compute the min and max of the data points
    min_grid = np.amin(points, axis=0)
    max_grid = np.amax(points, axis=0)
				
	# Increase the bounding box of data points by decreasing min_grid and inscreasing max_grid
    min_grid = min_grid - 0.10*(max_grid-min_grid)
    max_grid = max_grid + 0.10*(max_grid-min_grid)

	# grid_resolution is the number of voxels in the grid in x, y, z axis
    grid_resolution = 100
    size_voxel = np.array([(max_grid[0]-min_grid[0])/grid_resolution,(max_grid[1]-min_grid[1])/grid_resolution,(max_grid[2]-min_grid[2])/grid_resolution])
	
	# Create a volume grid to compute the scalar field for surface reconstruction
    scalar_field = np.zeros((grid_resolution+1,grid_resolution+1,grid_resolution+1),dtype = np.float32)

	# Compute the scalar field in the grid
    compute_hoppe(points,normals,scalar_field,grid_resolution,min_grid,size_voxel)
    # compute_imls(points,normals,scalar_field,grid_resolution,min_grid,size_voxel,30)

	# Compute the mesh from the scalar field based on marching cubes algorithm
    verts, faces, normals_tri, values_tri = measure.marching_cubes_lewiner(scalar_field, level=0.0, spacing=(size_voxel[0],size_voxel[1],size_voxel[2]))
	
    # Export the mesh in ply using trimesh lib
    mesh = trimesh.Trimesh(vertices = verts, faces = faces)
    mesh.export(file_obj='../bunny_mesh_hoppe_30.ply', file_type='ply')
	
    print("Total time for surface reconstruction : ", time.time()-t0)