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

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

//---------------------------------------------------------------
//                  LENS DISTORTION FUNCTIONS
//---------------------------------------------------------------

void standard_correction_model(matrix_t* const pDisPts, matrix_t* const pCorrPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    matrix_t x, y, rSq;
    
    for(index_t i=0; i< nDataPts; ++i)
    {
        x = pDisPts[i];
        y = pDisPts[i + nDataPts];
        
        rSq = pow(x - pCoefsIn[STD_COEF_CX], 2.0) + pow(y - pCoefsIn[STD_COEF_CY], 2.0);
        
        pCorrPts[i]            = x*( pCoefsIn[STD_COEF_K1]*rSq + pCoefsIn[STD_COEF_K2]*pow(rSq, 2.0) ) + pCoefsIn[STD_COEF_Q1]*( pow(x, 2.0) + pow(y, 2.0) ) + x*( pCoefsIn[STD_COEF_P1]*x + pCoefsIn[STD_COEF_P2]*y ) + x;
        pCorrPts[i + nDataPts] = y*( pCoefsIn[STD_COEF_K1]*rSq + pCoefsIn[STD_COEF_K2]*pow(rSq, 2.0) ) + pCoefsIn[STD_COEF_Q2]*( pow(x, 2.0) + pow(y, 2.0) ) + y*( pCoefsIn[STD_COEF_P1]*x + pCoefsIn[STD_COEF_P2]*y ) + y;
        
    }
    
}


void standard_fov_correction_model(matrix_t* const pDisPts, matrix_t* const pCorrPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    matrix_t xd, yd, ru, rdSq, rd, rMult;
    
    for(index_t i=0; i< nDataPts; ++i)
    {
        xd = pDisPts[i];
        yd = pDisPts[i + nDataPts];
        
        rdSq = pow(xd - pCoefsIn[FOV_COEF_CX], 2.0) + pow(yd - pCoefsIn[FOV_COEF_CY], 2.0);
        
        rd = sqrt(rdSq);
        ru = tan( rd*pCoefsIn[FOV_COEF_W] )/( 2*tan(pCoefsIn[FOV_COEF_W]/2) );
        
        if(isnan(ru) || isinf(ru) || rd == 0 )
        {
            rMult = 1.0;
        }
        else
        {
            rMult = (ru/rd);
        }
        
        pCorrPts[i]            = rMult*( 1.0 + pCoefsIn[FOV_COEF_K2]*pow(rdSq, 2.0) + pCoefsIn[FOV_COEF_K3]*pow(rdSq, 3.0) )*xd;
        pCorrPts[i + nDataPts] = rMult*( 1.0 + pCoefsIn[FOV_COEF_K2]*pow(rdSq, 2.0) + pCoefsIn[FOV_COEF_K3]*pow(rdSq, 3.0) )*yd;
        
    }
    
}


void standard_rad_correction_model(matrix_t* const pDisPts, matrix_t* const pCorrPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    matrix_t xd, yd, rd, Lr;
    
    for(index_t i=0; i< nDataPts; ++i)
    {
        xd = pDisPts[i];
        yd = pDisPts[i + nDataPts];
        
        rd = sqrt( pow(xd - pCoefsIn[RAD_COEF_CX], 2.0) + pow(yd - pCoefsIn[RAD_COEF_CY], 2.0) );
        
        Lr = pCoefsIn[RAD_COEF_K0] + pCoefsIn[RAD_COEF_K1]*rd + pCoefsIn[RAD_COEF_K2]*pow(rd, 2.0) + pCoefsIn[RAD_COEF_K3]*pow(rd, 3.0) + pCoefsIn[RAD_COEF_K4]*pow(rd, 4.0);
        
        pCorrPts[i] = Lr*(xd - pCoefsIn[RAD_COEF_CX]) + pCoefsIn[RAD_COEF_CX];
        pCorrPts[i + nDataPts] = Lr*(yd - pCoefsIn[RAD_COEF_CY]) + pCoefsIn[RAD_COEF_CY];
        
    }
}


void blind_rad_correction_model(matrix_t* const pDisPts, matrix_t* const pCorrPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    standard_rad_correction_model(pCorrPts, pDisPts, pCoefsIn, nDataPts);
}


void blind_rad_fixed_point_operator(matrix_t* const pDisPts, matrix_t* const pCorrPts, matrix_t* const pNextPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    matrix_t xd, yd, xu, yu, Lru;
    matrix_t ru;
    
    /* nonlinear fixed point method for vector-valued function. gauss-seidel variant updates variables each inner calculation.*/
    
    for(index_t i=0; i< nDataPts; ++i)
    {
        xd = pDisPts[i];
        yd = pDisPts[i + nDataPts];
        
        xu = pCorrPts[i];
        yu = pCorrPts[i + nDataPts];
        
        ru = sqrt( pow(xu - pCoefsIn[RAD_COEF_CX], 2.0) + pow(yu - pCoefsIn[RAD_COEF_CY], 2.0) );
        
        Lru = pCoefsIn[RAD_COEF_K0] + pCoefsIn[RAD_COEF_K1]*ru + pCoefsIn[RAD_COEF_K2]*pow(ru, 2.0) + pCoefsIn[RAD_COEF_K3]*pow(ru, 3.0) + pCoefsIn[RAD_COEF_K4]*pow(ru, 4.0);
        
        if( Lru == 0 )
        {
            LOG_ERROR("blind_rad_fixed_point_operator: Lru is zero-valued");
        }
        
        pNextPts[i] = (xd - pCoefsIn[RAD_COEF_CX])/Lru + pCoefsIn[RAD_COEF_CX];
        
        /* recompute Lru based on new value of xu */
        xu   = pNextPts[i];
        ru   = sqrt( pow(xu - pCoefsIn[RAD_COEF_CX], 2.0) + pow(yu - pCoefsIn[RAD_COEF_CY], 2.0) );
        Lru  = pCoefsIn[RAD_COEF_K0] + pCoefsIn[RAD_COEF_K1]*ru + pCoefsIn[RAD_COEF_K2]*pow(ru, 2.0) + pCoefsIn[RAD_COEF_K3]*pow(ru, 3.0) + pCoefsIn[RAD_COEF_K4]*pow(ru, 4.0);
        
        if( Lru == 0 )
        {
            LOG_ERROR("blind_rad_fixed_point_operator: Lru is zero-valued");
        }
        
        pNextPts[i + nDataPts] = (yd - pCoefsIn[RAD_COEF_CY])/Lru + pCoefsIn[RAD_COEF_CY];
        
         /* check for NaN values */
         if( (pNextPts[i] != pNextPts[i]) || (pNextPts[i + nDataPts] != pNextPts[i + nDataPts]) )
         {
             LOG_ERROR("blind_std_fixed_point_operator: NextPts are NaN-valued");
         }
        
    }
    
}


void blind_fov_correction_model(matrix_t* const pDisPts, matrix_t* const pCorrPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    matrix_t xu, yu, ru, ruSq, rd, rMult;
    
    for(index_t i=0; i< nDataPts; ++i)
    {
        xu = pCorrPts[i];
        yu = pCorrPts[i + nDataPts];
        
        ruSq = pow(xu - pCoefsIn[FOV_COEF_CX], 2.0) + pow(yu - pCoefsIn[FOV_COEF_CY], 2.0);
        
        ru = sqrt(ruSq);
        rd = atan( 2*ru*tan(pCoefsIn[FOV_COEF_W]/2) )/pCoefsIn[FOV_COEF_W];            
        
        if(isnan(rd) || isinf(rd) || ru == 0 )
        {
            rMult = 1.0;
        }
        else
        {
            rMult = (rd/ru);
        }
            
        pDisPts[i]            = rMult*( 1.0 + pCoefsIn[FOV_COEF_K2]*pow(ruSq, 2.0) + pCoefsIn[FOV_COEF_K3]*pow(ruSq, 3.0) )*xu;
        pDisPts[i + nDataPts] = rMult*( 1.0 + pCoefsIn[FOV_COEF_K2]*pow(ruSq, 2.0) + pCoefsIn[FOV_COEF_K3]*pow(ruSq, 3.0) )*yu;
    
    }
    
}


void blind_fov_fixed_point_operator(matrix_t* const pDisPts, matrix_t* const pCorrPts, matrix_t* const pNextPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    matrix_t xd, yd, xu, yu;
    matrix_t ruSq, ru, rd, rMult;
    
    for(index_t i=0; i< nDataPts; ++i)
    {
        xd = pDisPts[i];
        yd = pDisPts[i + nDataPts];
        
        xu = pCorrPts[i];
        yu = pCorrPts[i + nDataPts];
        
        ruSq = pow(xu - pCoefsIn[FOV_COEF_CX], 2.0) + pow(yu - pCoefsIn[FOV_COEF_CY], 2.0);

        ru = sqrt(ruSq);
        rd = atan( 2*ru*tan(pCoefsIn[FOV_COEF_W]/2) )/pCoefsIn[FOV_COEF_W];
        
        if( isnan(rd) || isinf(rd) || rd == 0 )
        {
            rMult = 1.0;
        }
        else
        {
            rMult = (rd/ru);
        }
        
        pNextPts[i]            = xd/rMult - ( pCoefsIn[FOV_COEF_K2]*pow(ruSq, 2.0) + pCoefsIn[FOV_COEF_K3]*pow(ruSq, 3.0) )*xu;
        pNextPts[i + nDataPts] = yd/rMult - ( pCoefsIn[FOV_COEF_K2]*pow(ruSq, 2.0) + pCoefsIn[FOV_COEF_K3]*pow(ruSq, 3.0) )*yu;
        
    }
    
}


void blind_std_correction_model(matrix_t* const pDisPts, matrix_t* const pCorrPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    standard_correction_model(pCorrPts, pDisPts, pCoefsIn, nDataPts);
}


void blind_std_fixed_point_operator(matrix_t* const pDisPts, matrix_t* const pCorrPts, matrix_t* const pNextPts, const coef_t* pCoefsIn, uint32_t nDataPts)
{
    matrix_t xDis, yDis, xCorr, yCorr;
    matrix_t rCorrSq;
    
    /* nonlinear fixed point method for vector-valued function. gauss-seidel variant updates variables each inner calculation.*/
    
    for(index_t i=0; i< nDataPts; ++i)
    {
        xDis = pDisPts[i];
        yDis = pDisPts[i + nDataPts];
        
        xCorr = pCorrPts[i];
        yCorr = pCorrPts[i + nDataPts];
        
        rCorrSq = pow(xCorr - pCoefsIn[STD_COEF_CX], 2.0)  + pow(yCorr - pCoefsIn[STD_COEF_CY], 2.0);
        
        pNextPts[i] = xDis - xCorr*( pCoefsIn[STD_COEF_K1]*rCorrSq + pCoefsIn[STD_COEF_K2]*pow(rCorrSq, 2.0) ) - pCoefsIn[STD_COEF_Q1]*( pow(xCorr, 2.0) + pow(yCorr, 2.0) ) - xCorr*( pCoefsIn[STD_COEF_P1]*xCorr + pCoefsIn[STD_COEF_P2]*yCorr );
        
        xCorr = pNextPts[i];
        rCorrSq = pow(xCorr - pCoefsIn[STD_COEF_CX], 2.0) + pow(yCorr - pCoefsIn[STD_COEF_CY], 2.0);
        
        pNextPts[i + nDataPts] = yDis - yCorr*( pCoefsIn[STD_COEF_K1]*rCorrSq + pCoefsIn[STD_COEF_K2]*pow(rCorrSq, 2.0) ) - pCoefsIn[STD_COEF_Q2]*( pow(xCorr, 2.0) + pow(yCorr, 2.0) ) - yCorr*( pCoefsIn[STD_COEF_P1]*xCorr + pCoefsIn[STD_COEF_P2]*yCorr );
        
        /* check for NaN values */
        /*
        if( (pNextPts[i] != pNextPts[i]) || (pNextPts[i + nDataPts] != pNextPts[i + nDataPts]) )
        {
            LOG_WARNING("blind_std_fixed_point_operator: NextPts are NaN-valued");
        }
        */
        
    }
    
}


//---------------------------------------------------------------
//                  LINE-FITTING FUNCTIONS
//---------------------------------------------------------------

void harp_best_fit_lines(vector_t* const pAlphas, vector_t* const pBetas, vector_t* const pGammas, const matrix_t* ptsIn, const matrix_t* lineSizeMatIn, uint32_t nDataPts, uint32_t nGroups, uint32_t nMaxLinesPerGroup)
{
    matrix_t x, y, xSqSum, ySqSum, xySum;
    matrix_t* xMean;
    matrix_t* yMean;
    
    xMean = (matrix_t*)calloc(nMaxLinesPerGroup, sizeof(matrix_t));
    yMean = (matrix_t*)calloc(nMaxLinesPerGroup, sizeof(matrix_t));
    
    matrix_t Axx, Axy, Ayy;
    coef_t theta;
    
    uint32_t nLinePts;
    index_t meanIdx;
    
    index_t ptIdx = 0;
    index_t prevLineIdx = 0;
    index_t currLineIdx = 0;
    
    for(index_t g=0; g < nGroups; ++g)
    {
        Axx = 0;
        Ayy = 0;
        Axy = 0;
        
        for(index_t l=0; l < nMaxLinesPerGroup; ++l)
        {
            nLinePts = lineSizeMatIn[g + l*nGroups];
            
            if(nLinePts <= 0)
            {
                break;
            }
            
            xMean[l] = 0;
            yMean[l] = 0;
            
            xSqSum = 0;
            ySqSum = 0;
            xySum  = 0;
            
            for(index_t i=0; i < nLinePts; ++i)
            {
                x = ptsIn[ptIdx];
                y = ptsIn[ptIdx + nDataPts];
                
                xMean[l] += x;
                yMean[l] += y;
                
                xSqSum += pow(x, 2.0);
                ySqSum += pow(y, 2.0);
                xySum  += x*y;
                
                ptIdx++;
                
            }
            
            assert(nLinePts > 0);
            
            xMean[l] /= (matrix_t)nLinePts;
            yMean[l] /= (matrix_t)nLinePts;
            
            Axx += xSqSum - (matrix_t)nLinePts*pow(xMean[l], 2.0);
            Ayy += ySqSum - (matrix_t)nLinePts*pow(yMean[l], 2.0);
            Axy += xySum -  (matrix_t)nLinePts*xMean[l]*yMean[l];
            
            currLineIdx++;
            
        }
        
        theta      = 0.5*atan2(-2.0*Axy, Axx - Ayy);
        pAlphas[g] = sin(theta);
        pBetas[g]  = cos(theta);
        
        meanIdx = 0;
        while(prevLineIdx < currLineIdx)
        {
            pGammas[prevLineIdx] = pAlphas[g]*xMean[meanIdx] + pBetas[g]*yMean[meanIdx];
            prevLineIdx++;
            meanIdx++;
        }
        
    }
    
    free (xMean);
    free (yMean);
    
}

void closest_line_points(const matrix_t* pPts, matrix_t* const pClosestPts, const matrix_t* pLineSizeMat, vector_t* const pAlphas, vector_t* const pBetas, vector_t* const pGammas, uint32_t nGroups, uint32_t nMaxLinesPerGroup, uint32_t nPts)
{
    uint32_t lineIdx = 0;
    uint32_t ptIdx   = 0;
    
    uint32_t nLinePts;
    matrix_t distToLine;
    
    for(index_t g=0; g<nGroups; ++g)
    {
        for(index_t l=0; l<nMaxLinesPerGroup; ++l)
        {
            nLinePts = pLineSizeMat[g + l*nGroups];
            
            if( nLinePts <= 0)
            {
                break;
            }
            
            for(index_t i=0; i<nLinePts; ++i)
            {
                distToLine = pAlphas[g]*pPts[ptIdx] + pBetas[g]*pPts[ptIdx + nPts] - pGammas[lineIdx];
                
                pClosestPts[ptIdx]        = pPts[ptIdx] - distToLine*pAlphas[g];
                pClosestPts[ptIdx + nPts] = pPts[ptIdx + nPts] - distToLine*pBetas[g];
                
                ptIdx++;
                
            }
            
            lineIdx++;
        }
    }
    
}


void compute_residual(matrix_t* const resOut, const matrix_t* ptsIn, const matrix_t* lineSizeMatIn, const coef_t* pAlphas, const coef_t* pBetas, const coef_t* pGammas, const vector_t* weightsIn, uint32_t nDataPts, uint32_t nGroups, uint32_t nLines, uint32_t nMaxLinesPerGroup)
{
    
    uint32_t nLinePts;
    index_t ptIdx   = 0;
    index_t lineIdx = 0;
    
    for(index_t g = 0; g < nGroups; ++g)
    {
        for(index_t l = 0;l < nMaxLinesPerGroup; ++l)
        {
            nLinePts = lineSizeMatIn[g + l*nGroups];
            
            if(nLinePts > 0)
            {
                for(index_t i = 0; i < nLinePts; ++i)
                {
                    resOut[ptIdx] = pAlphas[g]*ptsIn[ptIdx] + pBetas[g]*ptsIn[ptIdx + nDataPts] - pGammas[lineIdx];
                    
                    resOut[ptIdx] *= sqrt(weightsIn[ptIdx]);
                    
                    ptIdx++;
                }
            }
            else
            {
                break;
            }
            
            lineIdx++;
        }
    }
    
    assert( lineIdx == nLines+1 );
    assert( ptIdx == nDataPts+1 );
    
}



