# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import numpy as np
from scipy import stats
import igl
import meshplot as mp


# +
# Utility function to generate a tet grid
# n is a 3-tuple with the number of cell in every direction
# mmin/mmax are the grid bounding box corners

def tet_grid(n, mmin, mmax):
    nx = n[0]
    ny = n[1]
    nz = n[2]
    
    delta = mmax-mmin
    
    deltax = delta[0]/(nx-1)
    deltay = delta[1]/(ny-1)
    deltaz = delta[2]/(nz-1)
    
    T = np.zeros(((nx-1)*(ny-1)*(nz-1)*6, 4), dtype=np.int64)
    V = np.zeros((nx*ny*nz, 3))

    mapping = -np.ones((nx, ny, nz), dtype=np.int64)


    index = 0
    for i in range(nx):
        for j in range(ny):
            for k in range(nz):
                mapping[i, j, k] = index
                V[index, :] = [i*deltax, j*deltay, k*deltaz]
                index += 1
    assert(index == V.shape[0])
    
    tets = np.array([
        [0,1,3,4],
        [5,2,6,7],
        [4,1,5,3],
        [4,3,7,5],
        [3,1,5,2],
        [2,3,7,5]
    ])
    
    index = 0
    for i in range(nx-1):
        for j in range(ny-1):
            for k in range(nz-1):
                indices = [
                    (i,   j,   k),
                    (i+1, j,   k),
                    (i+1, j+1, k),
                    (i,   j+1, k),

                    (i,   j,   k+1),
                    (i+1, j,   k+1),
                    (i+1, j+1, k+1),
                    (i,   j+1, k+1),
                ]
                
                for t in range(tets.shape[0]):
                    tmp = [mapping[indices[ii]] for ii in tets[t, :]]
                    T[index, :]=tmp
                    index += 1
                    
    assert(index == T.shape[0])
    
    V += mmin
    return V, T
# -

# # Reading point cloud

# +
pi, v = igl.read_triangle_mesh("data/cat.off")
# pi, v = igl.read_triangle_mesh("data/luigi.off")

pi /= 10
ni = igl.per_vertex_normals(pi, v)
mp.plot(pi, shading={"point_size": 8})


# -

# Computing bounding box of point cloud
def compute_bbox(pi):
    """
    computes bounding box for point cloud
    Input:
        pi : point cloud
    """
    bbmin = np.min(pi, axis = 0)
    bbmax = np.max(pi, axis = 0)
    
    return bbmin, bbmax


def create_spatial_grid(pi, bmin, bmax, n):
    '''
    This function creates the spatial grid for a point cloud
    Note : Only handles square grids and equally spaced
    Input:
        pi : point cloud
        bmin : min index of bounding box
        
    '''
    delta = np.round((bmax - bmin)/n, 2)
    grid = {}
    grid["bl"] = []
    grid["ind"] = []
    grid["ind_arr"] = []
    for x_i in range(n):
        for y_i in range(n):
            for z_i in range(n):
                grid["ind"].append(str([x_i, y_i, z_i]))
                grid["ind_arr"].append([x_i, y_i, z_i])
                grid[str([x_i, y_i, z_i])] = []
                grid["bl"].append(bbmin + np.array([x_i*delta[0], y_i*delta[1], z_i*delta[2]]))

    for i in range(len(pi)):
        x_i = int(min(np.floor((pi[i,0] - bmin[0])/delta[0]), float(n-1)))
        y_i = int(min(np.floor((pi[i,1] - bmin[1])/delta[1]), float(n-1)))
        z_i = int(min(np.floor((pi[i,2] - bmin[2])/delta[2]), float(n-1)))

        grid[str([x_i, y_i, z_i])].append(i)
              
    return grid



# +
# Setting up Constraints
def find_closest_point(point, points):
    '''
    This function returns the index of the closest point in points to point. 
    '''
    dist = np.linalg.norm(np.subtract(points, point), axis = 1)
    return np.argmin(dist)


def create_constraints(pi, ni, eps_init, find_closest_point):
    '''
    This function creates the constraint matrix A, b for every point in the point cloud
    Input:
        pi : points in the point cloud
        ni: normals for each point in the point cloud
        eps_init : epsillon distance to be moved in and out wrt each point
        find_closest_point : function that computes closest point
    '''
    n = len(pi)
    f = [] # distance constraints (0, eps, -eps)   
    p = [] # array of points which include pi, pi+, pi-
    
    for i in range(n):
        p.append(pi[i])
        f.append(0)
        # computing pi+
        eps = eps_init
        while True:
            # normalizing the normal
            n_ni = ni[i]/np.linalg.norm(ni[i])
            if find_closest_point(pi[i] + eps*n_ni , pi) != i:
                eps /= 2.0
            else:
                p.append(pi[i] + eps*ni[i])
                f.append(eps)
                break
        # computing pi -
        eps = eps_init
        while True:
            if find_closest_point(pi[i] - eps*n_ni , pi) != i:
                eps /= 2.0
            else:
                p.append(pi[i] - eps*ni[i])
                f.append(-eps)
                break
        
    p = np.asarray(p)
    f = np.asarray(f)
            
    return p, f


# +
# Setting up Constraints
def find_closest_point_acc(point, pi, sp_grid):
    '''
    This function returns the index of the closest point in points to point using an
    accelerated data structure
    Input:
        point : point of interest 
        pi : point cloud
        sp_grid : spatial grid structure information
    '''
    dist = np.linalg.norm(np.subtract(sp_grid["bl"], point), axis = 1)
    # Taking top 2 closest grids just to be safe as in some cases the
    # distance of the bottom left corner of the adjacent grid is closer than then the correct grid
    # to be chosen
    k = 30
    ind = np.argsort(dist)
    indices = np.concatenate((sp_grid[sp_grid["ind"][ind[0]]], sp_grid[sp_grid["ind"][ind[1]]]))
    for l in range(2,k):
        indices = np.concatenate((indices, sp_grid[sp_grid["ind"][ind[l]]]))
    indices = indices.astype(int)
    dist = np.linalg.norm(np.subtract(pi[indices], point), axis = 1)
    ind = np.argmin(dist)
    return indices[ind]

def closest_points_acc(point, points, h, sp_grid):
    dist = np.linalg.norm(np.subtract(sp_grid["bl"], point), axis = 1)
    # Taking top 2 closest grids just to be safe as in some cases the
    # distance of the bottom left corner of the adjacent grid is closer than then the correct grid
    # to be chosen
    k = 25#15#6
    ind = np.argsort(dist)
    indices = np.concatenate((sp_grid[sp_grid["ind"][ind[0]]], sp_grid[sp_grid["ind"][ind[1]]]))
    for l in range(2,k):
        indices = np.concatenate((indices, sp_grid[sp_grid["ind"][ind[l]]]))
    indices = indices.astype(int)
    dist = np.linalg.norm(np.subtract(points[indices], point), axis = 1)
    ind = np.argwhere(dist < h)
    return indices[ind]


# -

def align_point_cloud(pi):
    '''
    This function rotates the point cloud such that it is aligned to the default
    (x,y,z) axis (1,0,0), (0,1,0), (0,0,1)
    Input:
        pi : point cloud to be aligned
    '''
    val, mat = np.linalg.eig(np.cov(pi.T))
    mean = pi.mean(axis=0)
    
    R = np.identity(4)
    R[0:3,0:3] = np.linalg.inv(mat)
    
    T = np.identity(4)
    T[:,3][0:3] = mean
    T_inv = -T
    T_inv[3,3] = 1.0
    
    pi_tr = np.hstack((pi, np.ones((len(pi),1))))
    
    p_aligned = (np.matmul(np.matmul(np.matmul(T, R),T_inv),pi_tr.T)).T
    
    return p_aligned[:,0:3]


# +
# # Computing bounding box
bbmin, bbmax = compute_bbox(pi)
bbox_diag = np.linalg.norm(bbmax - bbmin)
grid = create_spatial_grid(pi, bbmin, bbmax, 4)

#computing wedland radius
eps_init = 0.01*bbox_diag
# creating constraint matrices for moving least square
find_closest_point_accelerated = lambda point, points : find_closest_point_acc(point, points, grid)
# this computes constraints using no acceleration 
p, f = create_constraints(pi, ni, eps_init, find_closest_point)
# this computes the constraints array using acceleration
# p, f = create_constraints(pi, ni, eps_init, find_closest_point_accelerated)

# p = align_point_cloud(p)
# Part one plot
tmp = np.tile([[0,0,1], [1,0,0], [0,1,0]], (int(len(p)/3.0),1))
mp.plot(p, c = tmp, shading={"point_size": 6})


# +
# MLS for cat
def closest_points(point, points, h):
    '''
    computes the closest points to point
    '''
    dist = np.linalg.norm(np.subtract(points, point), axis = 1)
    return np.argwhere(dist < h)

def mls(point, pi, f, h, basis, wt, closest_points):
    '''
    computes the local approximation of function f and then checks
    if the given point is outside or inside. 
    Input:
        point : point on the grid to be evaluated
        pi : point cloud (with eps and -eps)
        f : constraints (value of f(pi[i]))
        h : wendland radius
        basis : basis function 
        wt : weight function (wedland)
        closest_points : function that computes the closest point within a radius
    '''
    # finding closest points to given point
    cl_p = closest_points(point, pi, h)
    A = []
    W = np.zeros((len(cl_p), len(cl_p)))
    b = f[cl_p] 
    # add exception later
    if len(cl_p) < len(basis(p[0])):
        return 10000
    else:
        for i in range(len(cl_p)):
            W[i,i] = wt(np.linalg.norm(point - pi[int(cl_p[i])]), h)
            A.append(basis(pi[int(cl_p[i])]))

        A = np.matrix(A)
        W = np.matrix(W)
        b = np.matrix(b)
        try:
            x = np.linalg.solve(A.T*W*A, A.T*W*b)
        except:
            print(A.T*W*A)
        
        return float(np.dot(x.T, basis(point)))


# +
# different basis functions
def poly_1(x):
    """
    1 st degree polynomial basis
    """
    val = np.zeros(4)
    val[0] = 1.0
    val[1:] = x
    return val

def poly_2(x):
    """
    2 st degree polynomial basis
    """
    val = np.zeros(10)
    val[0] = 1.0
    val[1:4] = x
    val[4:7] = [x[0]**2, x[1]**2, x[2]**2]
    val[7:] = [x[0]*x[1], x[1]*x[2], x[0]*x[2]]
    
    return val

# weight functions

def wedland_func(r, h):
    """
    wedland weight function
    """
    return ((1 - r/h)**4)*(4*(r/h) + 1)
    


# +
# Computing grid and mls for every point on the grid to check if point lies inside or outside
import time
n = 20 #20#15 #20
h = 35 #40#30 #35 

#computing new spatial grid for the points with constraints
bbmin, bbmax = compute_bbox(p)
bbox_diag = np.linalg.norm(bbmax - bbmin)
grid_new = create_spatial_grid(p, bbmin, bbmax, 10)

x, T = tet_grid((n, n, n), bbmin - 0.1 * bbox_diag, bbmax + 0.1 * bbox_diag)

ind = np.zeros(len(x))
fx = np.zeros(len(x))

closest_points_accelerated = lambda point, points, h: closest_points_acc(point, points, h, grid_new)

a = time.time()
for i in range(len(x)):
    # computes the mls approximation with acceleration
    fx[i] = mls(x[i].copy(), p, f, h, poly_2, wedland_func, closest_points_accelerated)
    # computes the mls approximation without acceleration
#     fx[i] = mls(x[i].copy(), p, f, h, poly_2, wedland_func, closest_points)

    if  fx[i] > 0:
        ind[i] = 1
    else:
        ind[i] = -1 
b = time.time()
b-a

# +
# normalizing the bounding box to make it compatible for plotting
x_norm = np.subtract(x, bbmin)
x_norm[:,0] /= bbmax[0] - bbmin[0]
x_norm[:,1] /= bbmax[1] - bbmin[1]
x_norm[:,2] /= bbmax[2] - bbmin[2]

mp.plot(x_norm, c=ind, shading={"point_size": 0.1,"width": 800, "height": 800})
# -

sv, sf, _, _ = igl.marching_tets(x_norm, T, fx, 0)
mp.plot(sv, sf, shading={"wireframe": True})


# +
# filtering to find the largest connected component
def filter_mls(sf):
    """
    Filters the mesh and returns the largest connected component of the reconstructed mesh
    Input:
        sf : array containing faces after reconstructing the point cloud
    """
    tmp1 = igl.face_components(sf)
    ind = stats.mode(tmp1)[0]
    f_sf = []
    for i in range(len(sf)):
        if tmp1[i] == ind:
            f_sf.append(sf[i])
    
    return np.array(f_sf)

mp.plot(sv, filter_mls(sf) ,shading={"wireframe": True})
# -



StarcoderdataJupyterScriptsDedupFiltered