/*
Copyright (c) 2020 ETH Zurich

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Author: Katrin Lasinger
*/

#define _USE_MATH_DEFINES

#include <numeric>
#include <vector>
#include <algorithm>
#include <omp.h>
#include "mex.h"
#include <math.h>
#include "FastExp64.h"

#include <Eigen/Dense>
using namespace Eigen;

// Code adapted from Johannes L. Schoeneberger
// Copyright (c) 2018, ETH Zurich and UNC Chapel Hill.
//https://github.com/colmap/colmap/blob/master/src/base/triangulation.cc
Eigen::Vector3d TriangulatePoint(const Eigen::Matrix<double, 3, 4>& proj_matrix1,
                                 const Eigen::Matrix<double, 3, 4>& proj_matrix2,
                                 const double point1x, const double point1y,
                                 const double point2x, const double point2y) {
  Eigen::Matrix4d A;

  A.row(0) = point1x * proj_matrix1.row(2) - proj_matrix1.row(0);
  A.row(1) = point1y * proj_matrix1.row(2) - proj_matrix1.row(1);
  A.row(2) = point2x * proj_matrix2.row(2) - proj_matrix2.row(0);
  A.row(3) = point2y * proj_matrix2.row(2) - proj_matrix2.row(1);

  Eigen::JacobiSVD<Eigen::Matrix4d> svd(A, Eigen::ComputeFullV);

  return svd.matrixV().col(3).hnormalized();
}

 Vector2d polynomialCameraForward( Matrix<double, 19, 1> a_x, Matrix<double, 19, 1> a_y, Vector3d pt3d ){
	double pX = pt3d[0];
	double pY = pt3d[1];
	double pZ = pt3d[2];

	Matrix<double, 19, 1> A;
	A << 1, pX, pY, pZ, pX*pX, pX*pY, pY*pY, pX*pZ, pY*pZ, pZ*pZ, pX*pX*pX, pX*pX*pY, pX*pY*pY, pY*pY*pY, pX*pX*pZ, pX*pY*pZ, pY*pY*pZ, pX*pZ*pZ, pY*pZ*pZ;

	Vector2d pt2d;
	pt2d[0] = A.dot(a_x);
	pt2d[1] = A.dot(a_y);

	return pt2d;

 }

template <class Vec19_>
Vector3d TriangulatePoint_poly(Vec19_ a_x0, Vec19_ a_y0, Vec19_ a_x1, Vec19_ a_y1, Vector2d p0, Vector2d p1)
{
	double thresh = 0.001;
	Vector3d pt3d(0,0,0); //seems to be sufficient enough to just initialize with 0

	for(int iterations = 0; iterations<1000; iterations++ )
	{

        // check new projection
        Vector2d pt2dEst_1 = polynomialCameraForward(a_x0,a_y0,pt3d);
        Vector2d pt2dEst_2 = polynomialCameraForward(a_x1,a_y1,pt3d);

        Vector2d diff1 = p0-pt2dEst_1;
        Vector2d diff2 = p1-pt2dEst_2;

        if (diff1.norm()<thresh && diff2.norm()<thresh){
            break;
		}

        double pX = pt3d[0];
		double pY = pt3d[1];
		double pZ = pt3d[2];

        Vec19_ A_dX, A_dY, A_dZ;
	    A_dX << 0, 1, 0, 0, 2*pX, pY, 0, pZ, 0, 0, 3*pX*pX, 2*pX*pY, pY*pY, 0, 2*pX*pZ, pY*pZ, 0, pZ*pZ, 0;
	    A_dY << 0, 0, 1, 0, 0, pX, 2*pY, 0, pZ, 0, 0, pX*pX, pX*2*pY, 3*pY*pY, 0, pX*pZ, 2*pY*pZ, 0, pZ*pZ;
	    A_dZ << 0, 0, 0, 1, 0, 0, 0, pX, pY, 2*pZ, 0, 0, 0, 0, pX*pX, pX*pY, pY*pY, pX*2*pZ, pY*2*pZ;

		MatrixXd A(4,3);

        A(0,0) = A_dX.dot(a_x0);
        A(1,0) = A_dX.dot(a_y0);
		A(2,0) = A_dX.dot(a_x1);
        A(3,0) = A_dX.dot(a_y1);
		A(0,1) = A_dY.dot(a_x0);
        A(1,1) = A_dY.dot(a_y0);
		A(2,1) = A_dY.dot(a_x1);
        A(3,1) = A_dY.dot(a_y1);
        A(0,2) = A_dZ.dot(a_x0);
        A(1,2) = A_dZ.dot(a_y0);
		A(2,2) = A_dZ.dot(a_x1);
        A(3,2) = A_dZ.dot(a_y1);

        Vector4d delta_x;
		delta_x << diff1(0), diff1(1), diff2(0), diff2(1);

        //delta_X = A\delta_x;
		Vector3d delta_X = A.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(delta_x);

        pt3d = pt3d+delta_X;

	}

	return pt3d;

}

Vector3d polynomialCameraImg2World(Matrix<double, 19, 1> a_x, Matrix<double, 19, 1> a_y, Vector2d pt2d, double depth)
{
	double thresh = 0.001;
	Vector3d pt3d(0,0,depth); //seems to be sufficient enough to just initialize with 0

	for(int iterations = 0; iterations<1000; iterations++ )
	{

        // check new projection
        Vector2d pt2dEst = polynomialCameraForward(a_x,a_y,pt3d);

        Vector2d diff = pt2d-pt2dEst;

        if (diff.norm()<thresh){
            break;
		}

        double pX = pt3d[0];
		double pY = pt3d[1];
		double pZ = pt3d[2];

        Matrix<double, 19, 1> A_dX, A_dY, A_dZ;
	    A_dX << 0, 1, 0, 0, 2*pX, pY, 0, pZ, 0, 0, 3*pX*pX, 2*pX*pY, pY*pY, 0, 2*pX*pZ, pY*pZ, 0, pZ*pZ, 0;
	    A_dY << 0, 0, 1, 0, 0, pX, 2*pY, 0, pZ, 0, 0, pX*pX, pX*2*pY, 3*pY*pY, 0, pX*pZ, 2*pY*pZ, 0, pZ*pZ;
	    //A_dZ << 0, 0, 0, 1, 0, 0, 0, pX, pY, 2*pZ, 0, 0, 0, 0, pX*pX, pX*pY, pY*pY, pX*2*pZ, pY*2*pZ;

		MatrixXd A(2,2);

        A(0,0) = A_dX.dot(a_x);
        A(1,0) = A_dX.dot(a_y);
		A(0,1) = A_dY.dot(a_x);
        A(1,1) = A_dY.dot(a_y);

        Vector2d delta_x;
		delta_x << diff(0), diff(1);

		Vector2d delta_X = A.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(delta_x);

        pt3d(0) = pt3d(0)+delta_X(0);
        pt3d(1) = pt3d(1)+delta_X(1);

	}

	return pt3d;
}

/// matlab calling - polynomial camera model
void triangulatePartPoly ( int nlhs, mxArray *plhs[],
                int nrhs, const mxArray *prhs[])
{

  typedef double Scalar;
  typedef Matrix<double, 19, 1> Vector19d;

  Scalar *output1;

  const mxArray *cell_array_ptr_in = prhs[0]; // per cam 2d array of xy coordinates+intensity of 2d part pos
  std::vector< Scalar*> part;
  const mxArray *cell_element_ptr;
  const mxArray *cell_element_ptr2;
  mwIndex jcell;

  int numcam = mxGetNumberOfElements(cell_array_ptr_in);

  part.resize(numcam);
  std::vector<int> numPart(numcam,0);

  for (jcell=0; jcell<numcam; jcell++) {
	  cell_element_ptr = mxGetCell(cell_array_ptr_in,jcell);

	  const mwSize* dims = mxGetDimensions(cell_element_ptr);

	  numPart[jcell] = dims[1];

	  part[jcell] = (Scalar*) mxGetPr(cell_element_ptr);
  }

  const mxArray *cell_array_ptr_in_a = prhs[1]; // per cam: 19x2 coefficients of camera model
  std::vector< Vector19d > a_x;
  std::vector< Vector19d > a_y;
  a_x.resize(numcam);
  a_y.resize(numcam);

  for (jcell=0; jcell<numcam; jcell++) {
	  cell_element_ptr = mxGetCell(cell_array_ptr_in_a,jcell);

	  Scalar* a_vec = (Scalar*) mxGetPr(cell_element_ptr);
	  for (int i=0;i<19;i++){
		a_x[jcell](i) = a_vec[i];
		a_y[jcell](i) = a_vec[i+19];
	  }

  }
  int N       = (int)  (*mxGetPr(prhs[2]));
  int M       = (int)  (*mxGetPr(prhs[3]));
  int L       = (int)  (*mxGetPr(prhs[4]));
  Scalar triangError = (Scalar)  (*mxGetPr(prhs[5]));
  Scalar triangErrorSq = triangError*triangError;

  Scalar integralWeight = 1.0/sqrt(2.0*M_PI);

  std::vector< Scalar > part3d; //list of 3d pts incl intensity (so 4d)

  std::vector<std::vector<Scalar> > local(omp_get_max_threads());
  std::vector<int> numTriangulatedPts(omp_get_max_threads(),0);
#pragma omp parallel //num_threads(1)
{
	int np = omp_get_num_threads();

	int currThreadNum = omp_get_thread_num();

  //iterate over 2d peaks of cam0
  //#pragma omp parallel for
#pragma omp for //schedule(static)
  for(int i=0;i<numPart[0];i++){

	Vector2d pt_ref;
	pt_ref << part[0][i*3],part[0][i*3+1];
	Scalar distSq = pow(floor(pt_ref(0)+0.5)-pt_ref(0),2)+pow(floor(pt_ref(1)+0.5)-pt_ref(1),2);
	Scalar int_ref = part[0][i*3+2]/integralWeight/ exp(-distSq); //sigma=1

    //first slice (at z=0)
	Vector3d pt3d_front = polynomialCameraImg2World( a_x[0], a_y[0], pt_ref, 0 );
    //last slice (at z=L-1)
    Vector3d pt3d_back = polynomialCameraImg2World( a_x[0], a_y[0], pt_ref, L-1 );

    //project 3d start and end pt to image
    Vector2d front_epi = polynomialCameraForward( a_x[1], a_y[1], pt3d_front );
	Vector2d back_epi = polynomialCameraForward( a_x[1], a_y[1], pt3d_back );

    Scalar maxBB_x = std::max(front_epi(0),back_epi(0)) + triangError;
    Scalar maxBB_y = std::max(front_epi(1),back_epi(1)) + triangError;
    Scalar minBB_x = std::min(front_epi(0),back_epi(0)) - triangError;
    Scalar minBB_y = std::min(front_epi(1),back_epi(1)) - triangError;

	Scalar lowerTerm = std::sqrt(std::pow(back_epi(0)-front_epi(0),2)+std::pow(back_epi(1)-front_epi(1),2));

    int numPtsOnLine=0;
    int numFound3dpts=0;
    int idxCurrPt = numTriangulatedPts[currThreadNum];

    //get distance of pts to epipolar line
    for (int j=1; j<numPart[1];j++){

		if(part[1][j*3]<minBB_x || part[1][j*3]>maxBB_x || part[1][j*3+1]<minBB_y || part[1][j*3+1]>maxBB_y)
			continue;

       // compute distance
       Scalar upperTerm = abs((back_epi(0)-front_epi(0))*(front_epi(1)-part[1][j*3+1])-(front_epi(0)-part[1][j*3])*(back_epi(1)-front_epi(1)));
       Scalar d = upperTerm/lowerTerm;

       if (d > triangError)
           continue;

       numPtsOnLine++;

	   Vector3d pt3d = TriangulatePoint_poly(a_x[0],a_y[0],a_x[1],a_y[1],pt_ref,Vector2d(part[1][j*3],part[1][j*3+1]));

       // can add a check if pt is actually within volume (since my
       // epipolar line is alway from z=0 to z=L-1 it could be that part
       // of the ray is outside of the volume to the sides)
       // possible check for x and y should be sufficient
	   // assuming that N is already size including padding (so minus also for upper bound)
       if( pt3d(0) < 0 || pt3d(0) > N-1 || pt3d(1) < 0 || pt3d(1) > M-1 || pt3d(2) < 0 || pt3d(2) > L-1)
           continue;

		bool match=true;

	   // loop over remaining cams
	   for (int c=2;c<numcam;c++){
		   if(match == false) //if no match in prev other cam
			   continue;
		   match = false;
		   Vector2d pt2d_proj = polynomialCameraForward( a_x[c], a_y[c], pt3d );

		   Scalar maxBB_x_oc = pt2d_proj(0) + triangError;
           Scalar maxBB_y_oc = pt2d_proj(1) + triangError;
           Scalar minBB_x_oc = pt2d_proj(0) - triangError;
           Scalar minBB_y_oc = pt2d_proj(1) - triangError;

		   for (int k=1; k<numPart[c];k++){

			   if(part[c][k*3]<minBB_x_oc || part[c][k*3]>maxBB_x_oc || part[c][k*3+1]<minBB_y_oc || part[c][k*3+1]>maxBB_y_oc)
				   continue;

			   if (pow(part[c][k*3]-pt2d_proj(0),2) + (pow(part[c][k*3+1]-pt2d_proj(1),2)) >triangErrorSq)
					continue;

			   match = true;
		   }
	   }

	   if(match){
		   numTriangulatedPts[currThreadNum]++;
		   numFound3dpts++;
		   local[currThreadNum].push_back(pt3d(0));
		   local[currThreadNum].push_back(pt3d(1));
		   local[currThreadNum].push_back(pt3d(2));
		   local[currThreadNum].push_back(int_ref);

	   }

	}

    //distribute intensity over found ref pts
    if(numFound3dpts>1){
		for(int ii=0;ii<numFound3dpts;ii++)
			local[currThreadNum][(idxCurrPt+ii)*4+3] = int_ref*(Scalar)4/(Scalar)(3+numFound3dpts);
	}
  }
}

for (int p = 0; p < omp_get_max_threads(); ++p){
	if (numTriangulatedPts[p] >0)
		part3d.insert(part3d.end(),local[p].begin(),local[p].end());
}

int numTriangulatedPtsAll = part3d.size()/4;

  //------------------------------------------------------------------------------
  //write back to matlab

  plhs[0] = mxCreateDoubleMatrix( 4, numTriangulatedPtsAll, mxREAL);
  output1  = mxGetPr(plhs[0]);
  for (int i=0;i<part3d.size();i++)
      output1[i] = part3d[i];

}

/// matlab calling
void triangulatePart ( int nlhs, mxArray *plhs[],
                int nrhs, const mxArray *prhs[])
{

  typedef double Scalar;

  Scalar *output1;

  const mxArray *cell_array_ptr_in = prhs[0]; // per cam 2d array of xy coordinates+intensity of 2d part pos
  std::vector< Scalar*> part;
  const mxArray *cell_element_ptr;
  const mxArray *cell_element_ptr2;
  mwIndex jcell;

  int numcam = mxGetNumberOfElements(cell_array_ptr_in);

  part.resize(numcam);
  std::vector<int> numPart(numcam,0);

  for (jcell=0; jcell<numcam; jcell++) {
	  cell_element_ptr = mxGetCell(cell_array_ptr_in,jcell);

	  const mwSize* dims = mxGetDimensions(cell_element_ptr);

	  numPart[jcell] = dims[1];

	  part[jcell] = (Scalar*) mxGetPr(cell_element_ptr);
  }

  const mxArray *cell_array_ptr_in_P = prhs[1]; // per cam 2d array of xy coordinates+intensity of 2d part pos
  const mxArray *cell_array_ptr_in_C = prhs[2];
  std::vector< Matrix<Scalar, 3, 4> > P;
  std::vector< Vector3d > C;
  P.resize(numcam);
  C.resize(numcam);

  for (jcell=0; jcell<numcam; jcell++) {
	  cell_element_ptr = mxGetCell(cell_array_ptr_in_P,jcell);
	  cell_element_ptr2 = mxGetCell(cell_array_ptr_in_C,jcell);

	  Scalar* p_vec = (Scalar*) mxGetPr(cell_element_ptr);
	  for(int j=0;j<4;j++)
		for (int i=0;i<3;i++)
			  P[jcell](i,j) = p_vec[i+j*3];

	  Scalar* c_vec = (Scalar*) mxGetPr(cell_element_ptr2);
		for (int i=0;i<3;i++)
			  C[jcell](i) = c_vec[i];
  }
  int N       = (int)  (*mxGetPr(prhs[3]));
  int M       = (int)  (*mxGetPr(prhs[4]));
  int L       = (int)  (*mxGetPr(prhs[5]));
  Scalar triangError = (Scalar)  (*mxGetPr(prhs[6]));
  Scalar triangErrorSq = triangError*triangError;

  Scalar test = P[0](1,2);
  Scalar integralWeight = 1.0/sqrt(2.0*M_PI);

  std::vector< Scalar > part3d; //list of 3d pts incl intensity (so 4d)

  Matrix3d M_inv = P[0].block<3,3>(0,0).inverse();
  Matrix3d M2_inv = P[1].block<3,3>(0,0).inverse();
  Vector3d Pcol = P[0].col(3);

  std::vector<std::vector<Scalar> > local(omp_get_max_threads());
  std::vector<int> numTriangulatedPts(omp_get_max_threads(),0);
#pragma omp parallel //num_threads(1)
{

	int np = omp_get_num_threads();

	int currThreadNum = omp_get_thread_num();

  //iterate over 2d peaks of cam0
  //#pragma omp parallel for
#pragma omp for //schedule(static)
  for(int i=0;i<numPart[0];i++){

	Vector3d pt_ref;
	pt_ref << part[0][i*3],part[0][i*3+1],1;
	Scalar distSq = pow(floor(pt_ref(0)+0.5)-pt_ref(0),2)+pow(floor(pt_ref(1)+0.5)-pt_ref(1),2);
	Scalar int_ref = part[0][i*3+2]/integralWeight/ exp(-distSq); //sigma=1

	Vector3d pt;
	pt = M_inv * (pt_ref - Pcol);

    Vector3d x1=C[0];
    Vector3d x2=pt;
    Vector3d x21 = x2-x1;
    Scalar dvN = -x21(2);

    //first slice (at z=0)
    Scalar dx1Z = x1(2);
    Scalar t = dx1Z/dvN;
    Vector4d pt3d_front;
	pt3d_front << x1 + t*x21,1;
    //last slice (at z=L-1)
    dx1Z = x1(2)-(L-1);
    t = dx1Z/dvN;
    Vector4d pt3d_back;
	pt3d_back << x1 + t*x21,1;

    //project 3d start and end pt to image
    Vector3d front_epi = P[1]*pt3d_front;
    front_epi = front_epi/front_epi(2);
    Vector3d back_epi = P[1]*pt3d_back;
    back_epi = back_epi/back_epi(2);

    Scalar maxBB_x = std::max(front_epi(0),back_epi(0)) + triangError;
    Scalar maxBB_y = std::max(front_epi(1),back_epi(1)) + triangError;
    Scalar minBB_x = std::min(front_epi(0),back_epi(0)) - triangError;
    Scalar minBB_y = std::min(front_epi(1),back_epi(1)) - triangError;

	Scalar lowerTerm = std::sqrt(std::pow(back_epi(0)-front_epi(0),2)+std::pow(back_epi(1)-front_epi(1),2));

    int numPtsOnLine=0;
    int numFound3dpts=0;
    int idxCurrPt = numTriangulatedPts[currThreadNum];

    //get distance of pts to epipolar line
    for (int j=1; j<numPart[1];j++){

		// if outside bounding box ignore
		if(part[1][j*3]<minBB_x || part[1][j*3]>maxBB_x || part[1][j*3+1]<minBB_y || part[1][j*3+1]>maxBB_y)
			continue;

       // compute distance
       Scalar upperTerm = abs((back_epi(0)-front_epi(0))*(front_epi(1)-part[1][j*3+1])-(front_epi(0)-part[1][j*3])*(back_epi(1)-front_epi(1)));
       Scalar d = upperTerm/lowerTerm;

       if (d > triangError)
           continue;

       numPtsOnLine++;

	   Vector3d pt3d = TriangulatePoint(P[0],P[1],pt_ref(0),pt_ref(1),part[1][j*3],part[1][j*3+1]);

       // can add a check if pt is actually within volume (since my
       // epipolar line is alway from z=0 to z=L-1 it could be that part
       // of the ray is outside of the volume to the sides)
       // possible check for x and y should be sufficient
	   // assuming that N is already size including padding (so minus also for upper bound)
       if( pt3d(0) < 0 || pt3d(0) > N-1 || pt3d(1) < 0 || pt3d(1) > M-1 || pt3d(2) < 0 || pt3d(2) > L-1)
           continue;

		Vector4d pt3d_hom;
		pt3d_hom << pt3d,1;

		bool match=true;

	   // loop over remaining cams
	   for (int c=2;c<numcam;c++){
		   if(match == false) //if no match in prev other cam
			   continue;
		   match = false;
		   Vector3d pt2d_proj = P[c]*pt3d_hom;
		   pt2d_proj = pt2d_proj/pt2d_proj(2);

		   Scalar maxBB_x_oc = pt2d_proj(0) + triangError;
           Scalar maxBB_y_oc = pt2d_proj(1) + triangError;
           Scalar minBB_x_oc = pt2d_proj(0) - triangError;
           Scalar minBB_y_oc = pt2d_proj(1) - triangError;

		   for (int k=1; k<numPart[c];k++){
			   //Vector3d pt_cam_other;
			   //pt_cam_other << part[c][k*3],part[c][k*3+1],1;

			   //if(pt_cam_other(0)<minBB_x_oc || pt_cam_other(0)>maxBB_x_oc || pt_cam_other(1)<minBB_y_oc || pt_cam_other(1)>maxBB_y_oc)
			//	   continue;

			  // if (pow(pt_cam_other(0)-pt2d_proj(0),2) + (pow(pt_cam_other(1)-pt2d_proj(1),2)) >triangErrorSq)
				//	continue;

			   if(part[c][k*3]<minBB_x_oc || part[c][k*3]>maxBB_x_oc || part[c][k*3+1]<minBB_y_oc || part[c][k*3+1]>maxBB_y_oc)
				   continue;

			   if (pow(part[c][k*3]-pt2d_proj(0),2) + (pow(part[c][k*3+1]-pt2d_proj(1),2)) >triangErrorSq)
					continue;

			   match = true;
		   }
	   }

	   if(match){
		   numTriangulatedPts[currThreadNum]++;
		   numFound3dpts++;
		   local[currThreadNum].push_back(pt3d(0));
		   local[currThreadNum].push_back(pt3d(1));
		   local[currThreadNum].push_back(pt3d(2));
		   local[currThreadNum].push_back(int_ref);
	   }

	}

    //distribute intensity over found ref pts
    if(numFound3dpts>1){
		for(int ii=0;ii<numFound3dpts;ii++)
			local[currThreadNum][(idxCurrPt+ii)*4+3] = int_ref*(Scalar)4/(Scalar)(3+numFound3dpts);
	}

  }

}

for (int p = 0; p < omp_get_max_threads(); ++p){
	if (numTriangulatedPts[p] >0)
		part3d.insert(part3d.end(),local[p].begin(),local[p].end());
}

int numTriangulatedPtsAll = part3d.size()/4;

  //input: cell array of 2d pts per cam incl intensity, P (and others?)
  //output: triangulated 3d pts (get 0:n part per 2d pt)

  //------------------------------------------------------------------------------
  //write back to matlab

  plhs[0] = mxCreateDoubleMatrix( 4, numTriangulatedPtsAll, mxREAL);
  output1  = mxGetPr(plhs[0]);
  for (int i=0;i<part3d.size();i++)
      output1[i] = part3d[i];

}


#include "icp.h"
#include <Eigen/Core>
#include <Eigen/Cholesky>
#include <Eigen/Geometry>

namespace cuda_icp{
Eigen::Matrix4d TransformVector6dToMatrix4d(const Eigen::Matrix<double, 6, 1> &input) {
    Eigen::Matrix4d output;
    output.setIdentity();
    output.block<3, 3>(0, 0) =
            (Eigen::AngleAxisd(input(2), Eigen::Vector3d::UnitZ()) *
             Eigen::AngleAxisd(input(1), Eigen::Vector3d::UnitY()) *
             Eigen::AngleAxisd(input(0), Eigen::Vector3d::UnitX()))
                    .matrix();
    output.block<3, 1>(0, 3) = input.block<3, 1>(3, 0);
    return output;
}

Mat4x4f eigen_to_custom(const Eigen::Matrix4f& extrinsic){
    Mat4x4f result;
    for(uint32_t i=0; i<4; i++){
        for(uint32_t j=0; j<4; j++){
            result[i][j] = extrinsic(i, j);
        }
    }
    return result;
}

Mat4x4f eigen_slover_666(float *A, float *b)
{
    Eigen::Matrix<float, 6, 6> A_eigen(A);
    Eigen::Matrix<float, 6, 1> b_eigen(b);
    // ICP point to plane may be unstable, refer to
    // https://www.cs.princeton.edu/~smr/papers/icpstability.pdf
    // add a term ||x|| to make update reasonably small:
    // f = ||(Rp + T - q) * n|| + penalty * ||X||   ==>
    // (ATA + Identity * penalty) * X = B
    Eigen::Matrix6d iden = Eigen::Matrix6d::Identity();
    double penalty = 0.01;
    Eigen::Matrix6d ATA_with_pen = A_eigen.cast<double>() + penalty*iden;

    const Eigen::Matrix<double, 6, 1> update = ATA_with_pen.ldlt().solve(b_eigen.cast<double>());
    Eigen::Matrix4d extrinsic = TransformVector6dToMatrix4d(update);
    return eigen_to_custom(extrinsic.cast<float>());
}

void transform_pcd(std::vector<Vec3f>& model_pcd, Mat4x4f& trans){

#pragma omp parallel for
    for(uint32_t i=0; i < model_pcd.size(); i++){
        Vec3f& pcd = model_pcd[i];
        float new_x = trans[0][0]*pcd.x + trans[0][1]*pcd.y + trans[0][2]*pcd.z + trans[0][3];
        float new_y = trans[1][0]*pcd.x + trans[1][1]*pcd.y + trans[1][2]*pcd.z + trans[1][3];
        float new_z = trans[2][0]*pcd.x + trans[2][1]*pcd.y + trans[2][2]*pcd.z + trans[2][3];
        pcd.x = new_x;
        pcd.y = new_y;
        pcd.z = new_z;
    }
}

template<class T>
void cpu_exclusive_scan_serial(T* start, uint32_t N){
    T cache = start[0];
    start[0] = 0;
    for (uint32_t i = 1; i < N; i++)
    {
        T temp = cache + start[i-1];
        cache = start[i];
        start[i] = temp;
    }
}

template<class T>
std::vector<Vec3f> depth2cloud_cpu(T* depth, uint32_t width, uint32_t height, Mat3x3f& K,
                                uint32_t stride, uint32_t tl_x, uint32_t tl_y)
{
    std::vector<uint32_t> mask(width*height/stride/stride, 0);

#pragma omp parallel for collapse(2)
    for(uint32_t x=0; x<width/stride; x++){
        for(uint32_t y=0; y<height/stride; y++){
            if(depth[x*stride + y*stride*width] > 0) mask[x + y*width] = 1;
        }
    }

    // scan to find map: depth idx --> cloud idx
    uint32_t mask_back_temp = mask.back();

    // without cuda this can't be used
#ifdef CUDA_ON
//    thrust::exclusive_scan(thrust::host, mask.begin(), mask.end(), mask.begin(), 0); // in-place scan
    cpu_exclusive_scan_serial(mask.data(), mask.size()); // serial version is better in cpu maybe
#else
    cpu_exclusive_scan_serial(mask.data(), mask.size());
#endif
    uint32_t total_pcd_num = mask.back() + mask_back_temp;

    std::vector<Vec3f> cloud(total_pcd_num);

#pragma omp parallel for collapse(2)
    for(uint32_t x=0; x<width/stride; x++){
        for(uint32_t y=0; y<height/stride; y++){

            uint32_t idx_depth = x*stride + y*stride*width;
            uint32_t idx_mask = x + y*width;

            if(depth[idx_depth] <= 0) continue;

            float z_pcd = depth[idx_depth]/1000.0f;
            float x_pcd = (x + tl_x - K[0][2])/K[0][0]*z_pcd;
            float y_pcd = (y + tl_y - K[1][2])/K[1][1]*z_pcd;

            cloud[mask[idx_mask]] = {x_pcd, y_pcd, z_pcd};
        }
    }
    return cloud;
}

template std::vector<Vec3f> depth2cloud_cpu(int32_t* depth, uint32_t width, uint32_t height, Mat3x3f& K,
                                uint32_t stride, uint32_t tl_x, uint32_t tl_y);
template std::vector<Vec3f> depth2cloud_cpu(uint16_t* depth, uint32_t width, uint32_t height, Mat3x3f& K,
                                            uint32_t stride, uint32_t tl_x, uint32_t tl_y);

template<class Scene>
RegistrationResult ICP_Point2Plane_cpu(std::vector<Vec3f> &model_pcd, const Scene scene,
                                       const ICPConvergenceCriteria criteria)
{
    RegistrationResult result;
    RegistrationResult backup;

    std::vector<float> A_host(36, 0);
    std::vector<float> b_host(6, 0);
    thrust__pcd2Ab<Scene> trasnformer(scene);

    // use one extra turn
    for(uint32_t iter=0; iter<=criteria.max_iteration_; iter++){

        Vec29f reducer;

#pragma omp declare reduction( + : Vec29f : omp_out += omp_in) \
                       initializer (omp_priv = Vec29f::Zero())

#pragma omp parallel for reduction(+: reducer)
        for(size_t pcd_iter=0; pcd_iter<model_pcd.size(); pcd_iter++){
            Vec29f result = trasnformer(model_pcd[pcd_iter]);
            reducer += result;
        }

        Vec29f& Ab_tight = reducer;

        backup = result;

        float& count = Ab_tight[28];
        float& total_error = Ab_tight[27];
        if(count == 0) return result;  // avoid divid 0

        result.fitness_ = float(count) / model_pcd.size();
        result.inlier_rmse_ = std::sqrt(total_error / count);

        // last extra iter, just compute fitness & mse
        if(iter == criteria.max_iteration_) return result;

        if(std::abs(result.fitness_ - backup.fitness_) < criteria.relative_fitness_ &&
           std::abs(result.inlier_rmse_ - backup.inlier_rmse_) < criteria.relative_rmse_){
            return result;
        }

        for(int i=0; i<6; i++) b_host[i] = Ab_tight[21 + i];

        int shift = 0;
        for(int y=0; y<6; y++){
            for(int x=y; x<6; x++){
                A_host[x + y*6] = Ab_tight[shift];
                A_host[y + x*6] = Ab_tight[shift];
                shift++;
            }
        }

        Mat4x4f extrinsic = eigen_slover_666(A_host.data(), b_host.data());

        transform_pcd(model_pcd, extrinsic);
        result.transformation_ = extrinsic * result.transformation_;
    }

    // never arrive here
    return result;
}

template RegistrationResult ICP_Point2Plane_cpu(std::vector<Vec3f> &model_pcd, const Scene_projective scene,
const ICPConvergenceCriteria criteria);
template RegistrationResult ICP_Point2Plane_cpu(std::vector<Vec3f> &model_pcd, const Scene_nn scene,
const ICPConvergenceCriteria criteria);

/// !!!!!!!!!!!!!!!!!!!!! legacy
// just for test and comparation
template<class Scene>
RegistrationResult ICP_Point2Plane_cpu_global_memory_version(std::vector<Vec3f> &model_pcd, const Scene scene,
                                       const ICPConvergenceCriteria criteria)
{
    RegistrationResult result;
    RegistrationResult backup;

    // buffer can make pcd handling indenpendent
    // may waste memory, but make it easy to parallel
    Eigen::Matrix<float, Eigen::Dynamic, 6> A_buffer(model_pcd.size(), 6); A_buffer.setZero();
    Eigen::Matrix<float, Eigen::Dynamic, 1> b_buffer(model_pcd.size(), 1); b_buffer.setZero();

    std::vector<uint32_t> valid_buffer(model_pcd.size(), 0);

    // use one extra turn
    for(uint32_t iter=0; iter<=criteria.max_iteration_; iter++){

#pragma omp parallel for
        for(uint32_t i = 0; i<model_pcd.size(); i++){
            const auto& src_pcd = model_pcd[i];

            Vec3f dst_pcd, dst_normal; bool valid;
            scene.query(src_pcd, dst_pcd, dst_normal, valid);
            if(valid){

                // dot
                b_buffer(i) = (dst_pcd - src_pcd).x * dst_normal.x +
                              (dst_pcd - src_pcd).y * dst_normal.y +
                              (dst_pcd - src_pcd).z * dst_normal.z;

                // cross
                A_buffer(i, 0) = dst_normal.z*src_pcd.y - dst_normal.y*src_pcd.z;
                A_buffer(i, 1) = dst_normal.x*src_pcd.z - dst_normal.z*src_pcd.x;
                A_buffer(i, 2) = dst_normal.y*src_pcd.x - dst_normal.x*src_pcd.y;

                A_buffer(i, 3) = dst_normal.x;
                A_buffer(i, 4) = dst_normal.y;
                A_buffer(i, 5) = dst_normal.z;

                valid_buffer[i] = 1;
            }else{
                b_buffer(i) = 0;

                A_buffer(i, 0) = 0;
                A_buffer(i, 1) = 0;
                A_buffer(i, 2) = 0;
                A_buffer(i, 3) = 0;
                A_buffer(i, 4) = 0;
                A_buffer(i, 5) = 0;

                valid_buffer[i] = 0;
            }
            // else: invalid is 0 in A & b, ATA ATb means adding 0,
            // so don't need to consider valid_buffer, just multi matrix
        }

        uint32_t count = 0;
        float total_error = 0;
#pragma omp parallel for reduction(+:count, total_error)
        for(uint32_t i=0; i<model_pcd.size(); i++){
            count += valid_buffer[i];
            total_error += (b_buffer(i)*b_buffer(i));
        }

        backup = result;

        if(count == 0) return result;  // avoid divid 0

        result.fitness_ = float(count) / model_pcd.size();
        result.inlier_rmse_ = std::sqrt(total_error / count);

//        {
//            std::cout << " --- cpu --- " << iter << " --- cpu ---" << std::endl;
//            std::cout << "total error: " << total_error << std::endl;
//            std::cout << "result.fitness_: " << result.fitness_ << std::endl;
//            std::cout << "result.inlier_rmse_: " << result.inlier_rmse_ << std::endl;
//            std::cout << " --- cpu --- " << iter << " --- cpu ---" << std::endl << std::endl;
//        }

        // last extra iter, just compute fitness & mse
        if(iter == criteria.max_iteration_) return result;

        if(std::abs(result.fitness_ - backup.fitness_) < criteria.relative_fitness_ &&
           std::abs(result.inlier_rmse_ - backup.inlier_rmse_) < criteria.relative_rmse_){
            return result;
        }

        Eigen::Matrix<float, 6, 6> A = A_buffer.transpose()*A_buffer;
        Eigen::Matrix<float, 6, 1> b = A_buffer.transpose()*b_buffer;

//        std::cout << "~~~~~~~~A~~~~~~" << std::endl;
//        std::cout << A;
//        std::cout << "\n~~~~~~~~~~~~~~\n" << std::endl;

//        std::cout << "~~~~~~~~b~~~~~~" << std::endl;
//        std::cout << b;
//        std::cout << "\n~~~~~~~~~~~~~~\n" << std::endl;

        Mat4x4f extrinsic = eigen_slover_666(A.data(), b.data());

//        std::cout << "~~extrinsic~~~~" << std::endl;
//        std::cout << extrinsic;
//        std::cout << "\n~~~~~~~~~~~~~~\n" << std::endl;

        transform_pcd(model_pcd, extrinsic);
        result.transformation_ = extrinsic * result.transformation_;
    }

    // never arrive here
    return result;
}

template RegistrationResult ICP_Point2Plane_cpu_global_memory_version(std::vector<Vec3f> &model_pcd, const Scene_projective scene,
const ICPConvergenceCriteria criteria);
template RegistrationResult ICP_Point2Plane_cpu_global_memory_version(std::vector<Vec3f> &model_pcd, const Scene_nn scene,
const ICPConvergenceCriteria criteria);
}


#ifndef LDLT_H
#define LDLT_H
#include <Eigen/Dense>
#include "mtao/type_utils.h"

namespace mtao::solvers::cholesky {

template <typename Matrix>
struct DenseLDLT_MIC0
{
    typedef typename Matrix::Scalar Scalar;
    template <typename A, typename B, typename C>
    inline Scalar tripleVectorProduct(const A & a, const B & b, const C & c){
        return a.cwiseProduct(b).dot(c);
    }
    inline Scalar tripleProduct(const Matrix & a, uint i, uint j)
    {
        return tripleVectorProduct(a.row(i).head(j),a.row(j).head(j),a.diagonal().head(j));
    }
    DenseLDLT_MIC0(const Matrix & A)
    {
        LD=A.template triangularView<Eigen::Lower>();
        int i,j;
        for(i=0; i<A.rows(); ++i)
        {
            for(j=0; j<i; ++j)
                if(std::abs(LD(j,j))>0.0001)
                {
                    LD(i,j)-=tripleProduct(LD,i,j);
                    LD(i,j)/=LD(j,j);
                }
                else
                {
                    LD(i,j)=0;
                }
            LD(i,i) -= tripleProduct(LD,i,i);
        }

    }
    template <typename Vector>
    void solve(const Vector & b, Vector & x)
    {
        x = LD.template triangularView<Eigen::UnitLower>().solve(b);
        x.noalias() = x.cwiseQuotient(LD.diagonal());//safe beacuse it's a dot
        LD.template triangularView<Eigen::UnitLower>().transpose().solveInPlace(x);
    }
    Matrix getA()
    {
        Matrix A = LD.template triangularView<Eigen::UnitLower>().transpose();
        A = LD.diagonal().asDiagonal() * A;
        A = LD.template triangularView<Eigen::UnitLower>() * A;

        return A;
    }
private:
    Matrix LD;
};

template <typename Matrix, typename Vector>
struct SparseLDLT_MIC0
{
    typedef typename Matrix::Scalar Scalar;
    SparseLDLT_MIC0() {}
    SparseLDLT_MIC0(const Matrix & A, const Vector& v): SparseLDLT_MIC0(A) {}
    SparseLDLT_MIC0(const Matrix & A)
    {
        // L=tril(A);
        L=A.template triangularView<Eigen::StrictlyLower>();//Don't copy the diagonal
        for(int i=0; i<L.rows(); ++i)
        {
            if(L.coeff(i,i)!=0)
                L.coeffRef(i,i)=0;
        }
        Dinv=D=A.diagonal();

        // for k=1:size(L,2)
        for(int k=0; k<A.rows(); ++k)//k is the column that we're infecting the remaining columns with
        {//L(:,k)

            //Solidify the current column values
            //==================================
            if(D(k)==0) continue;
            if(Dinv(k)<0.25*D(k))//If D has shrunk too much since it started
                Dinv(k)=1/D(k);
            else
                Dinv(k)=1/Dinv(k);
//            L.innerVector(k) *= Dinv(k);
            L.innerVector(k) = L.innerVector(k) *  Dinv(k);

            //Add k terms to all of the following columns
            //===========================================
            for(typename Matrix::InnerIterator it(L,k); it; ++it)// -L(i,k)*D(k)*L(j,k)
            {
                int j = it.row();//j>k
                if(j<=k) continue;
                Scalar missing=0;
                Scalar multiplier=it.value();//L(j,k)*D(k)

                typename Matrix::InnerIterator k_it(L,k);
                typename Matrix::InnerIterator j_it(L,j);
                //move down teh column of L(:,k) to collect missing elements in the match with A(:,j)
                //i=k_it.row()

                while (k_it && k_it.row()<j){//L(i,k)
                    while(j_it)//L(i,j) occasionally
                    {
                        if(j_it.row() < k_it.row())
                            ++j_it;
                        else if(j_it.row() == k_it.row())//L(i,k) are L(i,j) are nonzero
                            break;
                        else
                        {
                            missing += k_it.value();//L(i,k) will fill something not in L(i,j)
                            break;
                        }
                    }
                    ++k_it;
                }

                if(k_it && j_it.row() == j)
                {
                    Dinv(j) -= it.value() * multiplier;
                }

                typename Matrix::InnerIterator j_it2(L,j);
                while(k_it && j_it2)
                {
                    if(j_it2.row() < k_it.row())
                        ++j_it2;
                    else if(j_it2.row() == k_it.row())//L(i,k) and L(i,j) are both nonzero, -=L(i,k)*L(j,k)*D(k)
                    {
                        j_it2.valueRef() -= multiplier * k_it.value() ;//k_it.value()=L(i,k)
                        ++j_it2;
                        ++k_it;
                    }
                    else
                    {
                        missing+=k_it.value();
                        ++k_it;
                    }
                }

                while(k_it)
                {
                    missing+=k_it.value();
                    ++k_it;
                }
                Dinv(j)-=0.97*missing*multiplier;
            }
        }

        /*
           std::cout << L << std::endl;
           */

    }
    void solve(const Vector & b, Vector & x)
    {
        x = L.template triangularView<Eigen::UnitLower>().solve(b);
        x.noalias() = x.cwiseProduct(Dinv);//safe beacuse it's a dot
        L.transpose().template triangularView<Eigen::UnitUpper>().solveInPlace(x);
    }
    Matrix getA()
    {
        Matrix
                A = L.template triangularView<Eigen::UnitLower>();
        A = A * D.asDiagonal();
        A = A * L.template triangularView<Eigen::UnitLower>().transpose();

        return A;
    }
private:
    Matrix L;
    Vector D,Dinv;
};

template <typename MatrixType, typename VectorType>
using LDLT_MIC0 = std::conditional_t<
std::is_base_of_v<Eigen::SparseMatrixBase<MatrixType>,MatrixType>,
    SparseLDLT_MIC0<MatrixType,VectorType>,
    DenseLDLT_MIC0<MatrixType>>;

}

#endif


///////////////////////////////////////////////////////////////////////////////
// series.hpp                                                                //
//                                                                           //
//  Copyright 2010 Erwann Rogard. Distributed under the Boost                //
//  Software License, Version 1.0. (See accompanying file                    //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)         //
///////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_STATISTICS_DETAIL_NON_PARAMETRIC_EMPIRICAL_DISTRIBUTION_KOLMOGOROV_SMIRNOV_STATISTIC_SERIES_HPP_ER_2010
#define BOOST_STATISTICS_DETAIL_NON_PARAMETRIC_EMPIRICAL_DISTRIBUTION_KOLMOGOROV_SMIRNOV_STATISTIC_SERIES_HPP_ER_2010
#include <cmath>
#include <boost/numeric/conversion/converter.hpp>
#include <boost/accumulators/framework/accumulator_set.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/statistics/detail/non_parametric/empirical_distribution/kolmogorov_smirnov_statistic/value.hpp>

namespace boost{
namespace statistics{
namespace detail{
namespace empirical_distribution{
namespace kolmogorov_smirnov_statistic{

template<typename T1>
struct series_data
{
    typedef std::size_t size_type;
    series_data(){}
    size_type count; T1 value;
};

template<typename T>
std::ostream&
operator<<(std::ostream& os,const series_data<T>& data)
{
    return os << '('<< data.count << ',' << data.value << ')';
}

// Generates a series_data whose value is the Kolmogorov-Smirnov statistic
// for distribution dist, at each of the sample sizes
//    offset + {base^i : i = first_p,...,last_p };
// Each sample is drawn from gen().
template<typename D,typename G,typename It>
void series(
    const D& dist,
    G& gen,
    long offset,    // 1
    long base,      // 10
    int first_p,   // 0
    int last_p,    // 5
    It iter
){
    namespace ac = boost::accumulators;
    namespace ed = boost::statistics::detail::empirical_distribution;
	namespace ks = ed::kolmogorov_smirnov_statistic;
    typedef typename D::value_type val_;
    typedef std::size_t size_;
    typedef ks::tag::value<val_> tag_;
    typedef series_data<val_> series_data_;
    typedef typename G::result_type sample_;
    typedef ac::stats<tag_> stats_;
    typedef ac::accumulator_set<sample_,stats_> acc_;
    acc_ acc;

    typedef boost::numeric::converter<long double,long> int_float_;
    typedef boost::numeric::converter<long, long double> float_int_;

    long m = float_int_::convert(
        std::pow( int_float_::convert( base ), first_p )
    );
    long new_count;
    series_data_ series_data;
    series_data.count = 0;

    for(long p = first_p; p < last_p; p++){
        new_count = ( offset + m );
        for( long i = series_data.count; i < new_count; i++ ){ acc( gen() ); }
        series_data.count = new_count;
        series_data.value = ks::extract::value<val_>( acc, dist );
        ( *iter ) = series_data;
        m *= base;
    }

}

}// kolmogorov_smirnov_statistic
}// empirical_distribution
}// detail
}// statistics
}// boost

#endif