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

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

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

/*Mex Function Inputs:
 1)pts          Desc: Points matrix (ex. [x1,y1; x2,y2; ... xN,yN]). 
                      Ordered hiearchically, by group (image) and by line-segment within group
                Type: Nx2 matrix (double)
 
 2)lineSizeMat  Desc: Line size matrix. (ex. value at index (i,j) corresponds to # of data pts
                      belonging to line j of group i.
                Type: GxL matrix (double)
 
 3)weights      Desc: Weights for each point.
                Type: Nx1 column vector (double)
 
 4)coefs        Desc: Model parameters/coefficients
                Type: Kx1 column vector (double)
 
 5)maxIter      Desc: Maximum iteration count
                Type: positive integer
 
 6)methodType   Desc: Lens distortion correction plumb-line method type
                Type: MethodType Enumeration (double)
 
 */

/*Cost Function:
 
 sum(all groups: g){ sum(all lines in group: l) { sum(all pts in line: i) { w(i,l,g)*( alpha(g)*x_u(i,l,g) + beta(g)*y_u(i,l,g) - gamma(g,l) )^2 } } }
 
 */
/* Lens Distortion Correction Model:
 
 >see function "correction_model"
 
 */

/*Mex Function Outputs:
 1)residual  Desc: Reisdual of linear solve for x = [b's; c's; gamma's] s.t residual = A*x
             Type: Nx1 column-vector (double)
 
 2)alphas    Desc: Orientation parameter (i.e. sin(theta)) for each regression line group
             Type: Gx1 column-vector (double)
 
 3)betas     Desc: Orientation parameter (i.e. cos(theta)) for each regression line group
             Type: Gx1 column-vector (double)
 
 3)gammas    Desc: Offset value for each regression line
             Type: Qx1 column-vector (double)
 
 4)corrPts   Desc: Un-distorted (corrected) point matrix (ex. [x1',y1'; x2',y2'; ... xN',yN']).
             Type: Nx2 matrix (double)
 
 */

/*Background:
 >Solves un-distortion function using Levenberg-Marquardt
 >see: Wang, Jianhua, et al. "A new calibration model of camera lens distortion."
       Pattern Recognition 41.2 (2008): 607-615. 
 */


#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 PTS_IN         prhs[0]
#define LINESIZEMAT_IN prhs[1]
#define WEIGHTS_IN     prhs[2]
#define COEFS_IN       prhs[3]
#define MAXITER_IN     prhs[4]
#define METHODTYPE_IN  prhs[5]
    
#define RESIDUAL_OUT plhs[0]
#define ALPHAS_OUT   plhs[1]
#define BETAS_OUT    plhs[2]
#define GAMMAS_OUT   plhs[3]
#define CORRPTS_OUT  plhs[4]
    
    matrix_t *ptsIn;
    matrix_t *lineSizeMatIn;
    vector_t *weightsIn;
    coef_t   *coefsIn;
    scalar_t *pMethodTypeIn;
    scalar_t *pMaxIterIn;
    
    matrix_t *resOut;
    coef_t   *alphasOut;
    coef_t   *betasOut;
    coef_t   *gammasOut;
    matrix_t *corrPtsOut;
    
    count_t    maxIter;
    MethodType methodType;
    
    const mwSize *pPtsDims, *pLineSizeMatDims, *pWeightDims, *pCoefsDims;
    mwSize nCoefs, nDataPts, nGroups, nLines, nMaxLinesPerGroup;
    
    const mwSize nCoords = 2;
    
    mwSize resDims[1], alphasDims[1], betasDims[1], gammasDims[1], corrPtsDims[2];
    
    if(nrhs != 6) /* Check the number of input arguments */
    {
        LOG_ERROR("Wrong number of input arguments.");
    }
    else if(nlhs != 5) /* Check the number of output arguments */
    {
        LOG_ERROR("Too many output arguments.");
    }
    
    if( !IS_REAL_2D_FULL_DOUBLE(PTS_IN) )
    {
        LOG_ERROR("Points 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(WEIGHTS_IN) )
    {
        LOG_ERROR("Weights column-vector must be 2D double-precision matrix");
    }
    
    if( !IS_REAL_2D_FULL_DOUBLE(COEFS_IN) )
    {
        LOG_ERROR("Coefficient vector must be 2D double-precision matrices");
    }
    
    if( !IS_REAL_SCALAR(MAXITER_IN) )
    {
        LOG_ERROR("Maximum iteration count must be a double-precision scalar");
    }
    
    if( !IS_REAL_SCALAR(METHODTYPE_IN) )
    {
        LOG_ERROR("Method type must be a double-precision scalar");
    }
    
    pPtsDims         = mxGetDimensions(PTS_IN);
    pLineSizeMatDims = mxGetDimensions(LINESIZEMAT_IN);
    pWeightDims      = mxGetDimensions(WEIGHTS_IN);
    pCoefsDims       = mxGetDimensions(COEFS_IN);
    
    nGroups           = pLineSizeMatDims[0];
    nMaxLinesPerGroup = pLineSizeMatDims[1];
    nDataPts          = pPtsDims[0];
    nCoefs            = pCoefsDims[0];
    
    if( pWeightDims[0] != nDataPts)
    {
        LOG_ERROR("Expected length " << pWeightDims[0] << " of weight column-vector to be equal to number of data points " << nDataPts);
    }
    
    if( pPtsDims[1] != nCoords)
    {
        LOG_ERROR("Expected length points matrix to have coordinate pairs");
    }
    
    if( nCoefs != __std_ncoefs__ )
    {
        LOG_ERROR("Invalid number of coefficients provided (exp: " << __std_ncoefs__ << ", act: " << nCoefs << ")");
    }
    
    ptsIn          = mxGetPr(PTS_IN);
    lineSizeMatIn  = mxGetPr(LINESIZEMAT_IN);
    weightsIn      = mxGetPr(WEIGHTS_IN);
    coefsIn        = mxGetPr(COEFS_IN);
    pMethodTypeIn  = mxGetPr(METHODTYPE_IN);
    pMaxIterIn     = mxGetPr(MAXITER_IN);
    
    maxIter    = (count_t)*pMaxIterIn;
    methodType = (MethodType)*pMethodTypeIn;
    
    //compute total number of lines
    nLines = 0;
    for(index_t g=0; g < nGroups; ++g)
    {
        for(index_t l=0; l < nMaxLinesPerGroup; ++l)
        {
            if(lineSizeMatIn[g + l*nGroups] > 0)
            {
                nLines++;
            }
            else
            {
                //'zero-pt' lines right-pad the rows
                break;
            }
        }
    }
    
    if( methodType >= METHOD_NTYPES )
    {
        LOG_ERROR("Invalid method type specified");
    }
    
    if( maxIter <= 0 )
    {
        LOG_ERROR("Invalid maximum iteration count (<=0)");
    }
    
    resDims[0]     = nDataPts;
    alphasDims[0]  = nGroups;
    betasDims[0]   = nGroups;
    gammasDims[0]  = nLines;
    corrPtsDims[0] = nDataPts;
    corrPtsDims[1] = nCoords;
    
    //all elements initialized to zero by default
    RESIDUAL_OUT = mxCreateNumericArray(1, resDims, mxDOUBLE_CLASS, mxREAL);
    ALPHAS_OUT   = mxCreateNumericArray(1, alphasDims, mxDOUBLE_CLASS, mxREAL);
    BETAS_OUT    = mxCreateNumericArray(1, betasDims, mxDOUBLE_CLASS, mxREAL);
    GAMMAS_OUT   = mxCreateNumericArray(1, gammasDims, mxDOUBLE_CLASS, mxREAL);
    CORRPTS_OUT  = mxCreateNumericArray(2, corrPtsDims, mxDOUBLE_CLASS, mxREAL);
    
    resOut     = (matrix_t*)mxGetData(RESIDUAL_OUT);
    alphasOut  = (matrix_t*)mxGetData(ALPHAS_OUT);
    betasOut   = (matrix_t*)mxGetData(BETAS_OUT);
    gammasOut  = (matrix_t*)mxGetData(GAMMAS_OUT);
    corrPtsOut = (matrix_t*)mxGetData(CORRPTS_OUT);
    
    /* Compute the corrected image coordinates (see Lens Distortion Correction Model) */
    if( methodType == METHOD_STD )
    {
        __std_distortion_func__(ptsIn, corrPtsOut, coefsIn, (uint32_t)nDataPts);
    }
    else if( methodType == METHOD_BLIND )
    {
        count_t fpIter = 0;
        matrix_t *pCurrCorrPts, *pNextCorrPts;
        
        pCurrCorrPts = (matrix_t*)malloc(nDataPts*nCoords*sizeof(matrix_t));
        pNextCorrPts = (matrix_t*)malloc(nDataPts*nCoords*sizeof(matrix_t));
        memcpy( pNextCorrPts, ptsIn, nDataPts*nCoords*sizeof(matrix_t));
        
        while( fpIter < maxIter )
        {
            memcpy( pCurrCorrPts, pNextCorrPts, nDataPts*nCoords*sizeof(matrix_t));
            
            __fixed_point_operator__(ptsIn, pCurrCorrPts, pNextCorrPts, coefsIn, (uint32_t)nDataPts);
            
            fpIter++;
        }
        
        memcpy( corrPtsOut, pNextCorrPts, nDataPts*nCoords*sizeof(matrix_t));
        free (pNextCorrPts);
        free (pCurrCorrPts);
    }
    else
    {
        LOG_ERROR("Unknown method type!");
    }
     
    /* Compute the best-fit regression lines (using perpendicular offsets) */
    harp_best_fit_lines(alphasOut, betasOut, gammasOut, corrPtsOut, lineSizeMatIn, (uint32_t)nDataPts, (uint32_t)nGroups, (uint32_t)nMaxLinesPerGroup);
    
    /* Compute perpendicular residuals from regression lines (see Cost Function) */
    compute_residual(resOut, corrPtsOut, lineSizeMatIn, alphasOut, betasOut, gammasOut, weightsIn, (uint32_t)nDataPts, (uint32_t)nGroups, (uint32_t)nLines, (uint32_t)nMaxLinesPerGroup);
    
    
}//end of mexFunction

#endif//XCODE_COMPILATION