//
//  main.cpp
//  standard_lens_calibration
//
//  Created by Mehdi Stapleton on 2015-06-12.
//  Copyright (c) 2015 Mehdi Stapleton. All rights reserved.
//

#include "logging.h"
#include "plumb_line.h"
#include "convergence_criteria.h"
#include "lsqnonlin.h"

//---------------------------------------------------------------
//                          DESCRIPTION
//---------------------------------------------------------------

/*Mex Function Inputs:
 1)pDisPts      Desc: Normalized distorted image coordinates [x1,y1; x2,y2; ... xN,yN]
                Type: Nx2 matrix (double)
 
 2)lineSizeMat  Desc: The line-segment size matrix. Entry (i,j) tabulates the 
                      number of points in the ith group jth line segment.
                Type: GxK matrix (double)
 
 3) coefs       Desc: Initial estimate of lens distortion model parameters.
                Type: 1xQ vector (double)
 
 4) tol         Desc: Tolerance for change in model paramters. Measures the convergence of 
                       optimization routine.
                Type: Scalar (>0)
 
 5) maxIter     Desc: Maximum number of iterations for optimization routine.
                Tyoe: Scalar (>0)
*/

/*Mex Function Outputs:
 1)coefs        Desc: Lens distortion correction model parameters.
                Type: Qx1 vector (double)
 */

/*Background:
 >Solves for the correct (un-distorted/original) image coordinates which
   gave rise to the input distorted coordinates. The input image coordinates
   are assumed to originate from a camera with the lens distortion model of the
   form:
 
     x_d = H_x(x, y, k);
     y_d = H_y(x, y, k);
     
     x_d = distorted ;  x = undistorted
     y_d = distorted ;  y = undistorted

   The levenberg-marquardt optimization routine is applied to the model
   parameters to minimize cost functions which penalize the
   non-'straightness' of the distorted image coordinate lines.
 */


#ifndef XCODE_COMPILATION

//-------------------------------------------------------------------------
//              MEX File Generator Function Entry-Point
//-------------------------------------------------------------------------

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
    /* Macros for the ouput and input arguments */
#define DISPTS_IN       prhs[0]
#define LINESIZEMAT_IN  prhs[1]
#define COEFS_IN        prhs[2]
#define TOL_IN          prhs[3]
#define MAXITER_IN      prhs[4]
    
#define COEFS_OUT       plhs[0]
#define CORRPTS_OUT     plhs[1]
#define ERROR_OUT       plhs[2]
    
    
    matrix_t *pDisPtsIn, *pLineSizeMatIn, *pCoefsIn;
    scalar_t *pTolIn, *pMaxIterIn;
    coef_t   *pCoefsOut;
    matrix_t *pCorrPtsOut;
    vector_t *pErrorOut;

    coef_t   *pNextCoefs;
    vector_t *pWeights;
    
    const mwSize *pDisPtsDims, *pLineSizeMatDims, *pCoefsInDims;
    mwSize pCoefDims[1];
    mwSize pErrorDims[1];
    uint32_t nCoords, nDataPts, nPtEls, nGroups, nMaxLinesPerGroup, nLines, nCoefsIn, nMaxIters;
    
    const vector_t nanValue = mxGetNaN();
    
    if(nrhs != 5) /* Check the number of input arguments */
    {
        LOG_ERROR("Wrong number of input arguments.");
    }
    else if(nlhs != 3) /* Check the number of output arguments */
    {
        LOG_ERROR("Too many output arguments.");
    }
    
    if( !IS_REAL_2D_FULL_DOUBLE(DISPTS_IN) )
    {
        LOG_ERROR("Distorted coordinates matrix must be a 2D double-precision matrix");
    }
    
    if( !IS_REAL_2D_FULL_DOUBLE(LINESIZEMAT_IN) )
    {
        LOG_ERROR("Line Size matrix must be a 2D double-precision matrix");
    }
    
    if( !IS_REAL_2D_FULL_DOUBLE(COEFS_IN) )
    {
        LOG_ERROR("Initial lens distortion model parameters/coefficients matrix must be a 2D double-precision matrix");
    }
    
    if( !IS_REAL_SCALAR(TOL_IN) )
    {
        LOG_ERROR("Tolerance must be a double-precision scalar");
    }
    
    if( !IS_REAL_SCALAR(MAXITER_IN) )
    {
        LOG_ERROR("Maximum iteration count must be a double-precision scalar");
    }
    
    pDisPtsDims      = mxGetDimensions(DISPTS_IN);
    pLineSizeMatDims = mxGetDimensions(LINESIZEMAT_IN);
    pCoefsInDims     = mxGetDimensions(COEFS_IN);
    
    pDisPtsIn      = mxGetPr(DISPTS_IN);
    pLineSizeMatIn = mxGetPr(LINESIZEMAT_IN);
    pCoefsIn       = mxGetPr(COEFS_IN);
    pTolIn         = mxGetPr(TOL_IN);
    pMaxIterIn     = mxGetPr(MAXITER_IN);
    
    nDataPts          = (uint32_t)pDisPtsDims[0];
    nCoords           = (uint32_t)pDisPtsDims[1];
    nGroups           = (uint32_t)pLineSizeMatDims[0];
    nMaxLinesPerGroup = (uint32_t)pLineSizeMatDims[1];
    nCoefsIn          = (uint32_t)pCoefsInDims[0];
    nMaxIters         = (uint32_t)pMaxIterIn[0];
    
    nPtEls = nDataPts*nCoords;
    
    nLines = 0;
    for(index_t g=0; g<nGroups; ++g)
    {
        for(index_t l=0; l<nMaxLinesPerGroup; ++l)
        {
            if(pLineSizeMatIn[g + l*nGroups])
            {
                nLines++;
            }
            else
            {
                break;
            }
        }
    }
    
    pCoefDims[0]     = (mwSize)__std_ncoefs__;
    pErrorDims[0]    = (mwSize)nMaxIters;
    
    /*set all weights to one*/
    pWeights = (vector_t*)malloc(sizeof(vector_t)*nDataPts);
    for( index_t i =0; i< nDataPts; ++i)
    {
        pWeights[i] = 1.0;
    }
    
    /*all elements initialized to zero by default*/
    COEFS_OUT   = mxCreateNumericArray(1, pCoefDims, mxDOUBLE_CLASS, mxREAL);
    CORRPTS_OUT = mxCreateNumericArray(2, pDisPtsDims, mxDOUBLE_CLASS, mxREAL);
    ERROR_OUT   = mxCreateNumericArray(1, pErrorDims, mxDOUBLE_CLASS, mxREAL);
    
    pCoefsOut   = (coef_t*)mxGetData(COEFS_OUT);
    pCorrPtsOut = (matrix_t*)mxGetData(CORRPTS_OUT);
    pErrorOut   = (vector_t*)mxGetData(ERROR_OUT);
    
    pNextCoefs  = (coef_t*)calloc(__std_ncoefs__, sizeof(coef_t));
    
    /*initialize error curve to all NaN's*/
    for(index_t i=0; i<nMaxIters; ++i)
    {
        pErrorOut[i] = nanValue;
    }
    
    /*initialize lens distortion model parameters/coefficients*/
    if(pCoefsIn == NULL) /*input empty array*/
    {
        /*default: all zero*/
    }
    else if(nCoefsIn != __std_ncoefs__) /*invalid coefficients*/
    {
        LOG_ERROR("Invalid number of initial coefficients (" << nCoefsIn << ")");
    }
    else /*valid initial coefficient estimate*/
    {
        for(index_t i= 0; i< __std_ncoefs__; ++i)
        {
            pNextCoefs[i] = pCoefsIn[i];
        }
    }
    
    /*minimize orthogonal distance measure (perpendicular offsets from best fit lines) and slope-difference terms (for smoothness)*/
    const real lmTol = *pTolIn;
    int  *pIwa, info, lwa, ldfjac;
    real *pRes, *pJac, *pWa, resNorm;
    lsqdata_t lsqData;
    
    const int nLmParams    = (uint32_t)(__std_ncoefs__);
    const int nLmResiduals = (int)(nDataPts);
    
    lsqData.nPts     = (int)nDataPts;
    lsqData.nCoords  = (int)nCoords;
    lsqData.nGroups  = (int)nGroups;
    lsqData.nLines   = (int)nLines;
    lsqData.nMaxLinesPerGroup = (int)nMaxLinesPerGroup;
    
    lsqData.pDisPts      = pDisPtsIn;
    lsqData.pWeights     = pWeights;
    lsqData.pLineSizeMat = pLineSizeMatIn;
    
    lwa    = nLmResiduals*(__std_ncoefs__ + 1) + 5*__std_ncoefs__;
    pRes   = (real*)calloc(nLmResiduals, sizeof(real));
    pWa    = (real*)calloc(lwa, sizeof(real));
    pIwa   = (int*)calloc(__std_ncoefs__, sizeof(int));
    //ldfjac = nLmResiduals;
    //pJac   = (real*)malloc(nLmResiduals*__std_ncoefs__*sizeof(real));
    
    info = __cminpack_func__(lmdif1)(lsq_line_cost, &lsqData, nLmResiduals, nLmParams, (real*)pNextCoefs, pRes, /*pJac, ldfjac,*/ lmTol, pIwa, pWa, lwa);
    
    resNorm = __cminpack_func__(enorm)(nLmResiduals, pRes);
    
    pErrorOut[0] = (vector_t)resNorm;
    
    LOG_DEBUG("Info code returned by lmdif1: " << info);
    LOG_DEBUG("Residual l2 norm: " << resNorm);
    
    /*compute the corrected image coordinates (see Lens Distortion Correction Model)*/
    __std_distortion_func__(pDisPtsIn, pCorrPtsOut, pNextCoefs, (uint32_t)nDataPts);
    
    /*output finalized coefficients and line points*/
    memcpy(pCoefsOut, pNextCoefs, sizeof(coef_t)*__std_ncoefs__);
    
    /*clean up*/
    free (pIwa);
    free (pWa);
    free (pRes);
    //free (pJac);
    
    free (pNextCoefs);
    free (pWeights);
    
    
}//end of mexFunction

#endif//XCODE_COMPILATION