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

#include "logging.h"
#include "plumb_line.h"
#include "lens_correction.h"

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

/*Mex Function Inputs:
 1)inImg        Desc: Input RGB-color image
                Type: NxMx3 matrix (double)
 
 2)coefs        Desc: Model parameters/coefficients (see Lens Distortion Correction Model).
                      The model parameters are assumed to have been computed based on 
                      normalized coordinates.
                Type: Kx1 column vector (double)
 
*/

/*Mex Function Outputs:
 1)outImg       Desc: Output RGB-color image
                Type: NxMx3 matrix (double)
 */

/*Background
 >Performs forward transfer of known distorted coordinate values to undistorted/corrected
  coordinates using the provided correction-model.
 
 */


#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 IMAGE_IN       prhs[0]
#define COEFS_IN       prhs[1]
#define SCALE_IN       prhs[2]
#define OFFSET_IN      prhs[3]
    
#define IMAGE_OUT      plhs[0]
    
    matrix_t *rgbImgIn;
    coef_t   *pCoefsIn;
    vector_t *pScaleIn;
    vector_t *pOffsetIn;
    
    matrix_t *rgbImgOut;
    
    matrix_t* pDisPts, *pCorrPts;
    
    const mwSize nCoords = 2;
    const matrix_t nanValue = mxGetNaN();
    
    const mwSize *pImgDims, *pCoefDims, *pScaleInDims, *pOffsetInDims;
    mwSize nImgDims, nCoefs;
    mwSize nRows, nCols, nComps, nDataPts, nScaleIn, nOffsetIn;
    
    index_t ptIdx;
    int rNew, cNew;
    double rNorm, cNorm;
    
    if(nrhs != 4) /* Check the number of input arguments */
    {
        LOG_ERROR("Wrong number of input arguments.");
    }
    else if(nlhs != 1) /* Check the number of output arguments */
    {
        LOG_ERROR("Too many output arguments.");
    }
    
    if( !IS_REAL_3D_FULL_DOUBLE(IMAGE_IN) )
    {
        LOG_ERROR("Input color image matrix must be a 3D double-precision matrix");
    }
    
    if( !IS_REAL_2D_FULL_DOUBLE(COEFS_IN))
    {
        LOG_ERROR("Coefficients must be 2D double-precision matrices");
    }
    
    nImgDims      = mxGetNumberOfDimensions(IMAGE_IN);
    pImgDims      = mxGetDimensions(IMAGE_IN);
    pCoefDims     = mxGetDimensions(COEFS_IN);
    pScaleInDims  = mxGetDimensions(SCALE_IN);
    pOffsetInDims = mxGetDimensions(OFFSET_IN);
    
    rgbImgIn      = mxGetPr(IMAGE_IN);
    pCoefsIn      = mxGetPr(COEFS_IN);
    pScaleIn      = mxGetPr(SCALE_IN);
    pOffsetIn     = mxGetPr(OFFSET_IN);
    
    nRows         = (uint32_t)pImgDims[0];
    nCols         = (uint32_t)pImgDims[1];
    nComps        = (uint32_t)pImgDims[2];
    nCoefs        = pCoefDims[0];
    nDataPts      = nRows*nCols;
    
    nCoefs = (uint32_t)pCoefDims[0];
    
    if(nCoefs != __std_ncoefs__)
    {
        LOG_ERROR("Invalid number of coefficients (" << nCoefs << "), expected: " << __std_ncoefs__ );
    }
    
    nScaleIn  = (uint32_t)pScaleInDims[0];
    nOffsetIn = (uint32_t)pOffsetInDims[0];
    
    if( nScaleIn != 2 || nOffsetIn != 2)
    {
        LOG_ERROR("Invalid number of scale or offset entries (" << nScaleIn << "," << nOffsetIn << "), expected: " << 2 );
    }
    else if( pScaleIn[0] <= 0 || pScaleIn[1] <= 0 || pOffsetIn[0] <= 0 || pOffsetIn[1] <= 0)
    {
        LOG_ERROR("Invalid scale or offset entries");
    }
    
    //all elements initialized to zero by default
    IMAGE_OUT = mxCreateNumericArray(nImgDims, pImgDims, mxDOUBLE_CLASS, mxREAL);
    rgbImgOut  = (matrix_t*)mxGetData(IMAGE_OUT);
    
    //initialize output image to all NaN's
    for(index_t r=0; r< nRows; ++r)
    {
        for(index_t c=0; c< nCols; ++c)
        {
            rgbImgOut[r+nRows*c+(nRows*nCols)*R_PLANE] = nanValue;
            rgbImgOut[r+nRows*c+(nRows*nCols)*G_PLANE] = nanValue;
            rgbImgOut[r+nRows*c+(nRows*nCols)*B_PLANE] = nanValue;
        }
    }
    
    
    pDisPts  = (matrix_t*)calloc(nDataPts*nCoords, sizeof(matrix_t));
    pCorrPts = (matrix_t*)calloc(nDataPts*nCoords, sizeof(matrix_t));
    
    //create point matrix of normalized coordinates
    ptIdx = 0;
    for(index_t r=0; r< nRows; ++r)
    {
        for(index_t c=0; c< nCols; ++c)
        {
            rNorm = ( (matrix_t)r )/((matrix_t)pScaleIn[1]) - pOffsetIn[1];
            cNorm = ( (matrix_t)c )/((matrix_t)pScaleIn[0]) - pOffsetIn[0];
            //rNorm = ( (matrix_t)r )/nRows - 1.0/2.0;
            //cNorm = ( (matrix_t)c )/nCols - 1.0/2.0;
            
            pDisPts[ptIdx] = cNorm;
            pDisPts[ptIdx + nDataPts] = rNorm;
            
            ptIdx++;
        }
    }
    
    //perform correction mapping
    __std_distortion_func__(pDisPts, pCorrPts, pCoefsIn, (uint32_t)nDataPts);
    
    //transfer pixel values from distorted to corrected coordinates
    ptIdx = 0;
    for(index_t r=0; r< nRows; ++r)
    {
        for(index_t c=0; c< nCols; ++c)
        {
            cNew = std::round( (pCorrPts[ptIdx] + pOffsetIn[0])*pScaleIn[0] );
            rNew = std::round( (pCorrPts[ptIdx + nDataPts] + pOffsetIn[1])*pScaleIn[1] );
            //cNew = std::round( pCorrPts[ptIdx]*nCols + (double)nCols/2.0 );
            //rNew = std::round( pCorrPts[ptIdx + nDataPts]*nRows + (double)nRows/2.0 );
            
            if( rNew < nRows && rNew >= 0 && cNew < nCols && cNew >= 0 )
            {
                rgbImgOut[rNew+nRows*cNew+(nRows*nCols)*R_PLANE] = rgbImgIn[r+nRows*c+(nRows*nCols)*R_PLANE];
                rgbImgOut[rNew+nRows*cNew+(nRows*nCols)*G_PLANE] = rgbImgIn[r+nRows*c+(nRows*nCols)*G_PLANE];
                rgbImgOut[rNew+nRows*cNew+(nRows*nCols)*B_PLANE] = rgbImgIn[r+nRows*c+(nRows*nCols)*B_PLANE];
            }
            
            ptIdx++;
        }
    }
    
}//end of mexFunction

#endif//XCODE_COMPILATION
