/*
Copyright 2023. All rights reserved.
Robotics Research Group, University of Lleida.

Authors: Eduard Clotet.

The core of ComputeNormals is based on the LIBICP library provided by Andreas Geiger
(Karlsruhe Institute of Technology, Germany). LIBICP can be downladed from
the following link: 
https://github.com/MichiganTechRoboticsLab/SLAM-Scan-Matching/tree/master/ICP/libicp

Boost library is used to improve performance:
http://www.boost.org/users/download/

ComputeNormals is free software; you can redistribute it and/or modify it under the
terms of the GNU General Public License as published by the Free Software
Foundation; either version 2 of the License, or any later version.

ComputeNormals is distributed in the hope that it will be useful, but WITHOUT ANY
WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
PARTICULAR PURPOSE. See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with
ComputeNormals; if not, write to the Free Software Foundation, Inc., 51 Franklin
Street, Fifth Floor, Boston, MA 02110-1301, USA 
*/

#include "mex.h"
#include <iostream>
#include <string.h>
#define _USE_MATH_DEFINES
#include <math.h>
#include <boost/numeric/ublas/io.hpp>
#include "matrix.h"
#include "kdtree.h"

using namespace std;


#define __FILENAME__ (strrchr(__FILE__, '\\') ? strrchr(__FILE__, '\\') + 1 : __FILE__)


/***************************************************************
 * MEX help information
 ***************************************************************/
void displayHelp();

void computeNormals(const int32_t num_neighbors,const double flatness, const int32_t dim,kdtree::KDTree* M_tree, double* M_normal, int32_t*** neighborIndx);

void computeNormal (const kdtree::KDTreeResultVector &neighbors,double *M_normal,const double flatness, const int32_t dim, kdtree::KDTree* M_tree);

/***************************************************************
 * Main MEX code.
 ***************************************************************/
void mexFunction (int nlhs,mxArray *plhs[],int nrhs,const mxArray *prhs[]) {
    if (nrhs == 0 && nlhs == 0)
    {
        displayHelp();
        return;
    }
    
    // Validate mandatory input arguments
    if (!mxIsDouble(prhs[0]) || (mxGetM(prhs[0])!=2 && mxGetM(prhs[0])!=3))
        mexErrMsgTxt("Input M (model points) must be a double 2xN or 3xN matrix.");
    
    kdtree::KDTree*     M_tree;
    kdtree::KDTreeArray M_data;
    
    
    // Gather input data
    double *M                   = (double*)mxGetPr(prhs[0]);
    int32_t M_num               = (int32_t)mxGetN(prhs[0]);
    int32_t dim                 = (int32_t)mxGetM(prhs[0]);

    // Number of neighbor points used to compute the normals
    int numNeighbors;
    
    // Input error control:
    if (nrhs == 2)
    {
        if (!mxIsNumeric(prhs[1]))
            mexErrMsgTxt("Second input parameter must be a numeric value");
    
        numNeighbors = (int32_t)mxGetPr(prhs[1])[0];
        
        if (numNeighbors <= 0)
            mexErrMsgTxt("Second input parameter must greater than 1");
    }
    else
        numNeighbors = 10;
    
    
    //double  flatness              = *((double*)mxGetPr(prhs[2]));
    double  flatness              = 5.0;
    
    
    if (dim!=2 && dim!=3) {
        printf("ERROR: ComputeNormals works only for data of dimensionality 2 or 3\n");
        cout << "ERROR: ComputeNormals works only for data of dimensionality 2 or 3" << endl;
        M_tree = 0;
        return;
    }
    
    if (numNeighbors > M_num)
        numNeighbors = M_num;
    
    // copy model points to M_data
    M_data.resize(boost::extents[M_num][dim]);
    for (int32_t m=0; m<M_num; m++)
        for (int32_t n=0; n<dim; n++)
            M_data[m][n] = (float)M[m*dim+n];

    // build a kd tree from the model point cloud
    M_tree = new kdtree::KDTree(M_data);

    double *M_normal;
            
    M_normal = (double*)malloc(M_tree->N*dim*sizeof(double));
    
    /*********************************************************************/
    // Create a vector to store neighborhood
    // Define the initial size of neigbours list for the points to be recalculated
    int32_t** neighborsIndx = (int32_t**)malloc(M_num*sizeof(int32_t*));
    
    for (int32_t i = 0; i < M_num; i++) 
        neighborsIndx[i] = new int32_t[numNeighbors];
    
    
    computeNormals(numNeighbors, flatness, dim, M_tree, M_normal, &neighborsIndx);
    
    
    // Rearrange normals to be radial outwards
    for (int i = 0; i <M_num; i++)
    {
        double pt1x = M[i*2]+M_normal[i*2];
        double pt1y = M[i*2+1]+M_normal[i*2+1];
        
        double pt2x = M[i*2]-M_normal[i*2];
        double pt2y = M[i*2+1]-M_normal[i*2+1];      
        
        double hypot1 = pt1x*pt1x+pt1y*pt1y;
        double hypot2 = pt2x*pt2x+pt2y*pt2y;
        
        if (hypot1 < hypot2)
        {
            M_normal[i*2] = -M_normal[i*2];
            M_normal[i*2+1] = -M_normal[i*2+1];
        }
    }
    
    
    // Output outlier points.
    const int dims[] = {dim,M_num};
    plhs[0] = mxCreateNumericArray(2,dims,mxDOUBLE_CLASS,mxREAL);
    double *normals   = (double*)mxGetPr(plhs[0]);
    memcpy(normals,M_normal,M_num*sizeof(double)*dim);
        
    free(M_normal);
    
    // Output neighbors matrix.
    const int neighborDims[] = {numNeighbors,M_num};
    plhs[1] = mxCreateNumericArray(2,neighborDims,mxINT32_CLASS,mxREAL);
    int32_t *neigbhorsOut   = (int32_t*)mxGetPr(plhs[1]);

    for (int32_t i = 0; i < M_num; i++)
        memcpy(&neigbhorsOut[numNeighbors*i], &neighborsIndx[i][0], numNeighbors*sizeof(int32_t));
    
    
    for (int32_t i = 0; i < M_num; i++) 
        delete[] neighborsIndx[i];
        
    free(neighborsIndx);
}


//double* computeNormals(const int32_t num_neighbors,const double flatness, kdtree::KDTree* M_tree, const double dim) {
void computeNormals(const int32_t num_neighbors,const double flatness, const int32_t dim, kdtree::KDTree* M_tree, double* M_normal,int32_t*** neighborIndx) {
    kdtree::KDTreeResultVector neighbors;
    for (int32_t i=0; i<M_tree->N; i++) {
        M_tree->n_nearest_around_point(i,0,num_neighbors,neighbors);
        
        /*****************************************************************/
        // Store neighbors list
        /*****************************************************************/
        int32_t nIndx = 0;
        for (int32_t i2 = 0; i2 < neighbors.size(); i2++) 
        {
            (*neighborIndx)[i][nIndx] = neighbors[i2].idx+1;
            nIndx++;
        }
        /*****************************************************************/
                    
        
        if (dim==2) computeNormal(neighbors,M_normal+i*2,flatness,dim, M_tree);
        else        computeNormal(neighbors,M_normal+i*3,flatness,dim, M_tree);
    }
}


void computeNormal (const kdtree::KDTreeResultVector &neighbors,double *M_normal,const double flatness, const int32_t dim, kdtree::KDTree* M_tree) {
    // dimensionality 2
    if (dim==2) {
        // extract neighbors
        Matrix P(neighbors.size(),2);
        Matrix mu(1,2);

        for (int32_t i=0; i<neighbors.size(); i++) {
            double x = M_tree->the_data[neighbors[i].idx][0];
            double y = M_tree->the_data[neighbors[i].idx][1];
            P.val[i][0] = x;
            P.val[i][1] = y;
            mu.val[0][0] += x;
            mu.val[0][1] += y;
        }
        
        // zero mean
        mu       = mu/(double)neighbors.size();
        Matrix Q = P - Matrix::ones(neighbors.size(),1)*mu;
        
        // principal component analysis
        Matrix H = ~Q*Q;
        Matrix U,W,V;
        H.svd(U,W,V);
        
        // normal
        M_normal[0] = U.val[0][1];
        M_normal[1] = U.val[1][1];
        
    // dimensionality 3
    } else {
        
        // extract neighbors
        Matrix P(neighbors.size(),3);
        Matrix mu(1,3);
        for (int32_t i=0; i<neighbors.size(); i++) {
            double x = M_tree->the_data[neighbors[i].idx][0];
            double y = M_tree->the_data[neighbors[i].idx][1];
            double z = M_tree->the_data[neighbors[i].idx][2];
            P.val[i][0] = x;
            P.val[i][1] = y;
            P.val[i][2] = z;
            mu.val[0][0] += x;
            mu.val[0][1] += y;
            mu.val[0][2] += z;
        }
        
        // zero mean
        mu       = mu/(double)neighbors.size();
        Matrix Q = P - Matrix::ones(neighbors.size(),1)*mu;
        
        // principal component analysis
        Matrix H = ~Q*Q;
        Matrix U,W,V;
        H.svd(U,W,V);
        
        // normal
        M_normal[0] = U.val[0][2];
        M_normal[1] = U.val[1][2];
        M_normal[2] = U.val[2][2];
    }
}


/* Displays MEX info. */
void displayHelp()
{
    printf("%s (V 1)\n",__FILENAME__);
    printf("%s computes the normal vectors of a 2D or 3D point cloud\n",__FILENAME__);
    printf("The boost libary (boost_1_61_0) is used. This library is open \n");
    printf("source and its code can be downloaded from the official webstie:\n");
    printf("http://www.boost.org/users/download/.\n");
    printf("\n");
    printf("This library is based on the LIBICP library provided by Andreas Geiger.\n");
    printf("The original version of LIBICP can be downloaded from the following link:\n");
    printf("https://github.com/MichiganTechRoboticsLab/SLAM-Scan-Matching/tree/master/ICP/libicp\n\n");
    printf("\n");
    printf("This MEX can be executed using the following calls:\n");
    printf("\tnormals = ComputeNormals(M)\n");
    printf("\tnormals = ComputeNormals(M, numNeighbors)\n");
    printf("\t[normals, neighbors] = ComputeNormals(M)\n");
    printf("\t[normals, neighbors] = ComputeNormals(M, numNeighbors)\n\n");
    printf("\tInput parameters:\n");
    printf("\t\t-> M: A 2xN or 3xN vector containing the points of the Point Cloud, where M(i,:) contains the coordinates of the i'th dimension.\n");
    printf("\t\t-> numNeighbors (optional): numeric value indicating the number of neighbor points used to calculate the normals (default 10).\n");
    printf("\tOutput parameters:\n");
    printf("\t\t-> normals: A 2xN or 3xN vector containing the normals of the Point Cloud, where nromals(i,:) contains the coordinates of the i'th dimension.\n");
    printf("\t\t-> neighbors (optional): matrix containing the index of the neighbor points used to calculate the normals of each point, where neighbors(:,i) contains\n\t\t\t the N neighbors from the i'th point\n");
    
}