//
//  main.cpp
//  blind_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 parameters
                Type: Qx1 vector (double)
 
 4)tol          Desc: Stop condition tolerance for fixed point iteration and coefficients
                Type: positive scalar 
 
 5)maxIter      Desc: Maximum iteration count for fixed point iteration and entire algorithm
                Type: positive integer
 
*/

/*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
 
   An alternate-minimization scheme is employed in order to solve for the
   correct (x,y) coordinates as well as the model parameters k.
 
 */


#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;
    
    vector_t *pWeights;
    vector_t *pAlphas, *pBetas, *pGammas;
    matrix_t *pCurrCorrPts, *pNextCorrPts, *pClosestPts;
    coef_t   *pCurrCoefs, *pNextCoefs;
    
    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)__blind_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);
    
    pCurrCorrPts = (matrix_t*)calloc(nPtEls, sizeof(matrix_t));
    pNextCorrPts = (matrix_t*)calloc(nPtEls, sizeof(matrix_t));
    pClosestPts  = (matrix_t*)calloc(nPtEls, sizeof(matrix_t));
    pCurrCoefs   = (coef_t*)calloc(__blind_ncoefs__, sizeof(coef_t));
    pNextCoefs   = (coef_t*)calloc(__blind_ncoefs__, sizeof(coef_t));
    
    pAlphas      = (vector_t*)malloc(sizeof(vector_t)*nGroups);
    pBetas       = (vector_t*)malloc(sizeof(vector_t)*nGroups);
    pGammas      = (vector_t*)malloc(sizeof(vector_t)*nLines);
    
    /*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 != __blind_ncoefs__) /*invalid coefficients*/
    {
        LOG_ERROR("Invalid number of initial coefficients (" << nCoefsIn << ")");
    }
    else /*valid initial coefficient estimate*/
    {
        for(index_t i= 0; i< __blind_ncoefs__; ++i)
        {
            pNextCoefs[i] = pCoefsIn[i];
        }
    }
    
    
    /*blind iterations until relative change in coefficients is below threshold*/
    
    const uint32_t nLmParams    = (uint32_t)(__blind_ncoefs__-nLensDisCtrParams);
    const uint32_t nLmResiduals = (uint32_t)(nPtEls);
    
    const double coefRelNormTol   = *pTolIn;
    const double ptsRelNormTol    = *pTolIn;
    
    const real lmTol = sqrt(__cminpack_func__(dpmpar)(1));
    int  *pIwa, info, lwa, ldfjac;
    real *pRes, *pLineRes, *pJac, *pWa, resNorm;
    lsqdata_t lsqData;
    
    lsqData.nPts     = (int)nDataPts;
    lsqData.nCoords  = (int)nCoords;
    lsqData.pDisCtr  = (real*)(pNextCoefs + nLmParams);
    lsqData.pDisPts  = (real*)pDisPtsIn;
    lsqData.pCorrPts = (real*)pClosestPts;
    
    if(nLmParams > __blind_ncoefs__)
    {
        LOG_ERROR("LM parameters exceed number of coefficients");
    }
    
    lwa      = nLmResiduals*(__blind_ncoefs__ + 1) + 5*__blind_ncoefs__;
    ldfjac   = nLmResiduals;
    pRes     = (real*)calloc(nLmResiduals, sizeof(real));
    pLineRes = (real*)calloc(nDataPts, sizeof(real));
    pJac     = (real*)malloc(nLmResiduals*__blind_ncoefs__*sizeof(real));
    pWa      = (real*)calloc(lwa, sizeof(real));
    pIwa     = (int*)calloc(__blind_ncoefs__, sizeof(int));
    
    uint32_t bIter  = 0;
    uint32_t fpIter = 0;
    
    memcpy( pNextCorrPts, pDisPtsIn, sizeof(matrix_t)*nPtEls);
    
    while( bIter == 0 || ( bIter <  nMaxIters && relative_norm_delta(pCurrCoefs, pNextCoefs, (uint32_t)__blind_ncoefs__) > coefRelNormTol ) )
    {
        memcpy(pCurrCoefs, pNextCoefs, sizeof(coef_t)*__blind_ncoefs__);
        
        /*fixed point method to solve for next corrected coordinate estimate*/
        fpIter = 0;
        while( fpIter == 0 || ( fpIter < nMaxIters && relative_norm_delta(pCurrCorrPts, pNextCorrPts, (uint32_t)nPtEls) > ptsRelNormTol ) )
        {
            memcpy( pCurrCorrPts, pNextCorrPts, sizeof(matrix_t)*nPtEls);
            
            __fixed_point_operator__(pDisPtsIn, pCurrCorrPts, pNextCorrPts, pCurrCoefs, (uint32_t)nDataPts);
            
            fpIter++;
        }
        
        /*compute regression line parameters for all line segments*/
        harp_best_fit_lines(pAlphas, pBetas, pGammas, pNextCorrPts, pLineSizeMatIn, (uint32_t)nDataPts, (uint32_t)nGroups, (uint32_t)nMaxLinesPerGroup);
        
        /*compute residual error between 'corrected' coordinate and regression line*/
        compute_residual(pLineRes, pNextCorrPts, pLineSizeMatIn, pAlphas, pBetas, pGammas, pWeights, nDataPts, nGroups, nLines, nMaxLinesPerGroup);
        
        resNorm = __cminpack_func__(enorm)((int)(nDataPts), pLineRes);
        
        pErrorOut[bIter] = (vector_t)resNorm;
        
        /*project next corrected coordinate estimate onto feasible set of straight lines  (constraint on group lines being parallel)*/
        closest_line_points(pNextCorrPts, pClosestPts, pLineSizeMatIn, pAlphas, pBetas, pGammas, (uint32_t)nGroups, (uint32_t)nMaxLinesPerGroup, (uint32_t)nDataPts);
        
        memcpy(pNextCorrPts, pClosestPts, sizeof(matrix_t)*nPtEls);
        
        /*solve for optimal lens distortion parameters (excluding distortion center) which minimize regression line residuals*/
        memcpy(pNextCoefs, pCurrCoefs, sizeof(coef_t)*__blind_ncoefs__);
        
        info = __cminpack_func__(lmdif1)(lsq_line_cost, &lsqData, (int)nLmResiduals, (int)nLmParams, (real*)pNextCoefs, pRes, /*pJac, ldfjac,*/ lmTol, pIwa, pWa, lwa);
        
        resNorm = __cminpack_func__(enorm)((int)(nLmResiduals), pRes);
        
        /*log iteration details
        LOG_DEBUG("==== Iteration " << bIter << " ====");
        LOG_DEBUG("Info code returned by lmdif1: " << info);
        LOG_DEBUG("Residual l2 norm: " << resNorm);
        */
        
        bIter++;
        
    }
    
    
    /*solve for distortion center*/
    info = __cminpack_func__(lmdif1)(lsq_line_cost, &lsqData, (int)nLmResiduals, (int)__blind_ncoefs__, (real*)pNextCoefs, pRes, /*pJac, ldfjac,*/ lmTol, pIwa, pWa, lwa);
    
    resNorm = __cminpack_func__(enorm)((int)(nLmResiduals), pRes);
    
    LOG_DEBUG("==== Final Iteration ====");
    LOG_DEBUG("Info code returned by lmdif1: " << info);
    LOG_DEBUG("Residual l2 norm: " << resNorm);
    
    
    /*output finalized coefficients and line points*/
    memcpy(pCoefsOut, pNextCoefs, sizeof(coef_t)*__blind_ncoefs__);
    
    memcpy(pCorrPtsOut, pClosestPts, sizeof(matrix_t)*nPtEls);
    
    
    /*clean up*/
    free (pIwa);
    free (pWa);
    free (pRes);
    free (pJac);
    
    free (pClosestPts);
    free (pCurrCoefs);
    free (pNextCoefs);
    free (pCurrCorrPts);
    free (pNextCorrPts);
    
    free (pAlphas);
    free (pBetas);
    free (pGammas);
    
    
}//end of mexFunction

#endif//XCODE_COMPILATION