/*
 * Copyright 2023. All rights reserved.
 * Robotics Research Group, University of Lleida.
 *
 * Authors: Eduard Clotet.
 *
 * This MEX file merges two point clouds (S and T) after aplying the given 
 * transformation to the coordinates of T.
 *
 * PCMerge 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.
 *
 * PCMerge 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
 * PCMerge; 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

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

void generateTR(const double * posIn, const bool isMatrix, const int32_t dim, double **R, double **t);

void rotTransPoints(double *P,const int32_t P_num,double **R,double **t, const int32_t dim);
/***************************************************************
 * Rotates and translates the points in P according to R|t.
 * The rotated points are stored in P.
 *
 * @param P set of points to be rotated (input/output parameter).
 * @param P_num number of points cointained in P.
 * @param R Rotation matrix.
 * @param t Translation matrix.
 * @param dim number of dimensions of P.
 ***************************************************************/


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

/***************************************************************
 * Main MEX code
 ***************************************************************/
void mexFunction (int nlhs,mxArray *plhs[],int nrhs,const mxArray *prhs[]) {
    bool emptyM = false;
    bool isMatrix = false;
    if (nrhs == 0 && nlhs == 0)
    {
        displayHelp();
        return;
    }

    // Validate mandatory input arguments
    if (nrhs != 3)
        mexErrMsgTxt("Error: INCORRECT NUMBER OF INPUT ARGUMENTS");
    
    
    if (mxGetN(prhs[0]) == 0)
        emptyM = true;
        
    if (!emptyM && (!mxIsDouble(prhs[0]) || (mxGetM(prhs[0])!=2 && mxGetM(prhs[0])!=3)))
        mexErrMsgTxt("Input M (model points) must be a double 2xN or 3xN matrix.");
    if (!mxIsDouble(prhs[1]) || (mxGetM(prhs[1])!=2 && mxGetM(prhs[1])!=3))
        mexErrMsgTxt("Input T (template points) must be a double 2xN or or 3xN matrix.");
    if (!emptyM && mxGetM(prhs[0])!=mxGetM(prhs[1]) && mxGetN(prhs[0]) != 0)
        mexErrMsgTxt("Input M and T must have same number of rows (=dimensionality).");
    if (!mxIsDouble(prhs[2]))
        mexErrMsgTxt("Input position must be a 1x3 (2D) or 1x6 (3D) double vector.");
    else
    {
        if (mxGetM(prhs[1]) == 2)
        {
            if (mxGetN(prhs[2])!=3 && mxGetN(prhs[2])*mxGetM(prhs[2]) != 9)
                mexErrMsgTxt("Invalid position matrix/vector.");
            
            if (mxGetN(prhs[2])*mxGetM(prhs[2]) == 9)
                isMatrix = true;
        }
        else if (mxGetM(prhs[1]) == 3)
        {
            if (mxGetN(prhs[2])!=6 && mxGetN(prhs[2])*mxGetM(prhs[2]) != 16)
                mexErrMsgTxt("Invalid position matrix/vector.");
            
            if (mxGetN(prhs[2])*mxGetM(prhs[2]) == 16)
                isMatrix = true;
        }
    }
    
    // Gather input data
    double *M                   = (double*)mxGetPr(prhs[0]);
    int32_t M_num               = (int32_t)mxGetN(prhs[0]);
    double *T                   = (double*)mxGetPr(prhs[1]);
    int32_t T_num               = (int32_t)mxGetN(prhs[1]);
    int32_t dim                 = (int32_t)mxGetM(prhs[1]);
    double *posIn               = (double*)mxGetPr(prhs[2]);
    
    double *initPos;
    int numInPoints = mxGetN(prhs[2])*mxGetM(prhs[2]);
    
    initPos = new double[numInPoints];
    memcpy(initPos, posIn, numInPoints*sizeof(double));
    
    double ** r = new double*[dim];
    for(int i = 0; i < dim; ++i)
        r[i] = new double[dim];
    
    
    double ** tt = new double*[dim];
    for(int i = 0; i < dim; ++i)
        tt[i] = new double[1];
    
    generateTR(initPos, isMatrix, dim, r, tt);
        
    // Create a copy of M
    double *W;
    W = new double[M_num*dim];
    memcpy(W,M,M_num*dim*sizeof(double));
    
    
    // Create a copy of T
    double *TCpy = new double[T_num*dim];
    memcpy(TCpy,T,T_num*dim*sizeof(double));
    
    // Rotate and translate T 
    rotTransPoints(TCpy, T_num,r,tt, dim);
    
    // Create output parameters
    int totalPoints = M_num+T_num;
    const int mapDims[] = {dim,totalPoints};
    plhs[0] = mxCreateNumericArray(2,mapDims,mxDOUBLE_CLASS,mxREAL);
    double *comMap   = (double*)mxGetPr(plhs[0]);

    // Load 
    memcpy(comMap,M,M_num*sizeof(double)*dim);
    memcpy(&comMap[M_num*dim],TCpy,T_num*sizeof(double)*dim);
    
    // clear allocated memory
    delete[] TCpy;
    delete[] W;
    delete[] initPos;
    
    for(int i = 0; i < dim; ++i)
    {
        delete[] r[i];
        delete[] tt[i];
    }
    delete[] r;
    delete[] tt;
}

void generateTR(const double * posIn, const bool isMatrix, const int32_t dim, double **R, double **t)
{
    if (isMatrix)
    {
        if (dim==2) {
            for (int32_t i=0; i<2; i++) {
                for (int32_t j=0; j<2; j++)
                    R[i][j] = posIn[j*3+i];
                t[i][0] = posIn[2*3+i];
            }
            
            // input initial transformation (dim=3)
        } else {
            for (int32_t i=0; i<3; i++) {
                for (int32_t j=0; j<3; j++)
                    R[i][j] = posIn[j*4+i];
                t[i][0] = posIn[3*4+i];
            }
        }
    }
    else
    {
        if (dim == 2)
        {
            R[0][0] = cos(posIn[2]);
            R[0][1] = -sin(posIn[2]);
            R[1][0] = sin(posIn[2]);
            R[1][1] = cos(posIn[2]);            
        }
        else
        {
            double a = posIn[3];
            double b = posIn[4];
            double g = posIn[5];
            R[0][0] = cos(b)*cos(a);
            R[0][1] = sin(g)*sin(b)*cos(a) - cos(g)*sin(a);
            R[0][2] = cos(g)*sin(b)*cos(a) + sin(g)*sin(a);
            R[1][0] = cos(b)*sin(a);
            R[1][1] = sin(g)*sin(b)*sin(a) + cos(g)*cos(a);
            R[1][2] = cos(g)*sin(b)*sin(a) - sin(g)*cos(a);
            R[2][0] = -sin(b);
            R[2][1] = sin(g)*cos(b);
            R[2][2] = cos(g)*cos(b);
        }  
        for (int i = 0; i < dim; i++)
            t[i][0] = posIn[i];    
    }      
}

void rotTransPoints(double *P,const int32_t P_num,double **R,double **t, const int32_t dim)
{
    double r00 = R[0][0]; double r01 = R[0][1];
    double r10 = R[1][0]; double r11 = R[1][1];
    double t0  = t[0][0]; double t1  = t[1][0];
    
    double x;
    double y;
    
    if (dim==2)
    {
        for (int32_t i=0; i<P_num; i++)
        {
            x = r00*P[i*2+0] + r01*P[i*2+1] + t0;
            y = r10*P[i*2+0] + r11*P[i*2+1] + t1;
            // transform point according to R|t
            P[i*2+0] = x;
            P[i*2+1] = y;
        }
    }
    else
    {
        // extract matrix and translation vector
        double r02 = R[0][2];
        double r12 = R[1][2];
        double r20 = R[2][0]; double r21 = R[2][1]; double r22 = R[2][2];
        double t2  = t[2][0];
        double z;
        // check for all points if they are inliers
        for (int32_t i=0; i<P_num; i++)
        {
            x = r00*P[i*3+0] + r01*P[i*3+1] + r02*P[i*3+2] + t0;
            y = r10*P[i*3+0] + r11*P[i*3+1] + r12*P[i*3+2] + t1;
            z = r20*P[i*3+0] + r21*P[i*3+1] + r22*P[i*3+2] + t2;
            // transform point according to R|t
            P[i*3+0] = x;
            P[i*3+1] = y;
            P[i*3+2] = z;
        }
    }
}
/* Displays MEX info. */
void displayHelp()
{
    printf("%s Merges two point clouds (O and A), applying the given rotation/translation to A.\n\n",__FILENAME__);
    printf("This MEX can be executed using the following call:\n");
    printf("\tMP = PCMerge(O, A, ep)\n\n");
    printf("\tInput parameters:\n");
    printf("\t\t-> S (Static PC): A 2xN or 3xN array containing the coordinates of the points that will not be transformed, where S(i,:) contains the coordinates of the i'th dimension.\n");
    printf("\t\t-> T (Transform PC): A 2xN or 3xN array containing the coordinates of the points that will be transformed and merged. Where T(i,:) contains the coordinates the i'th dimension.\n");
    printf("\t\t-> ep (expected position): a position vector or a tranformation matrix containing describing the transformation that must be applied to A.\n");
    printf("\t\t\t2D matching:\n");
    printf("\t\t\t\t- As a position vector:\n");
    printf("\t\t\t\t\t- ep(1:2): x and y coordinates.\n");
    printf("\t\t\t\t\t- ep(3): orientation (in radians).\n");
    printf("\t\t\t\t- As a transformation matrix:\n");
    printf("\t\t\t\t\t- 3x3 double matrix containing the transformation matrix.\n");
    printf("\t\t\t3D matching:\n");
    printf("\t\t\t\t- As a position vector:\n");
    printf("\t\t\t\t- ep(1:3): x, y and z coordinates.\n");
    printf("\t\t\t\t- ep(4:6): yaw, pitch and roll angles (in radians).\n");
    printf("\t\t\t\t- As a transformation matrix:\n");
    printf("\t\t\t\t\t- 4x4 double matrix containing the transformation matrix.\n");    
    printf("\tOutput parameters:\n");
    printf("\t\t-> MP (Merged Points): 2xN or 3xN double array containing the concatenation of S and T.\n");
}