/*M///////////////////////////////////////////////////////////////////////////////////////
//
//  IMPORTANT: READ BEFORE DOWNLOADING, COPYING, INSTALLING OR USING.
//
//  By downloading, copying, installing or using the software you agree to this license.
//  If you do not agree to this license, do not download, install,
//  copy or use the software.
//
//
//                        Intel License Agreement
//                For Open Source Computer Vision Library
//
// Copyright (C) 2000, Intel Corporation, all rights reserved.
// Third party copyrights are property of their respective owners.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
//   * Redistribution's of source code must retain the above copyright notice,
//     this list of conditions and the following disclaimer.
//
//   * Redistribution's in binary form must reproduce the above copyright notice,
//     this list of conditions and the following disclaimer in the documentation
//     and/or other materials provided with the distribution.
//
//   * The name of Intel Corporation may not be used to endorse or promote products
//     derived from this software without specific prior written permission.
//
// This software is provided by the copyright holders and contributors "as is" and
// any express or implied warranties, including, but not limited to, the implied
// warranties of merchantability and fitness for a particular purpose are disclaimed.
// In no event shall the Intel Corporation or contributors be liable for any direct,
// indirect, incidental, special, exemplary, or consequential damages
// (including, but not limited to, procurement of substitute goods or services;
// loss of use, data, or profits; or business interruption) however caused
// and on any theory of liability, whether in contract, strict liability,
// or tort (including negligence or otherwise) arising in any way out of
// the use of this software, even if advised of the possibility of such damage.
//
//M*/

#include "_cvaux.h"
#include <stdio.h>

#undef quad

#if _MSC_VER >= 1200
#pragma warning( disable: 4701 )
#endif

CvCalibFilter::CvCalibFilter()
{
    /* etalon data */
    etalonType = CV_CALIB_ETALON_USER;
    etalonParamCount = 0;
    etalonParams = 0;
    etalonPointCount = 0;
    etalonPoints = 0;

    /* camera data */
    cameraCount = 1;

    memset( points, 0, sizeof(points));
    memset( undistMap, 0, sizeof(undistMap));
    undistImg = 0;
    memset( latestCounts, 0, sizeof(latestCounts));
    memset( latestPoints, 0, sizeof(latestPoints));
    memset( &stereo, 0, sizeof(stereo) );
    maxPoints = 0;
    framesTotal = 15;
    framesAccepted = 0;
    isCalibrated = false;

    imgSize = cvSize(0,0);
    grayImg = 0;
    tempImg = 0;
    storage = 0;

    memset( rectMap, 0, sizeof(rectMap));
}


CvCalibFilter::~CvCalibFilter()
{
    SetCameraCount(0);
    cvFree( &etalonParams );
    cvFree( &etalonPoints );
    cvReleaseMat( &grayImg );
    cvReleaseMat( &tempImg );
    cvReleaseMat( &undistImg );
    cvReleaseMemStorage( &storage );
}


bool CvCalibFilter::SetEtalon( CvCalibEtalonType type, double* params,
                               int pointCount, CvPoint2D32f* points )
{
    int i, arrSize;

    Stop();

    for( i = 0; i < MAX_CAMERAS; i++ )
        cvFree( latestPoints + i );

    if( type == CV_CALIB_ETALON_USER || type != etalonType )
    {
        cvFree( &etalonParams );
    }

    etalonType = type;

    switch( etalonType )
    {
    case CV_CALIB_ETALON_CHESSBOARD:
        etalonParamCount = 3;
        if( !params || cvRound(params[0]) != params[0] || params[0] < 3 ||
            cvRound(params[1]) != params[1] || params[1] < 3 || params[2] <= 0 )
        {
            assert(0);
            return false;
        }

        pointCount = cvRound((params[0] - 1)*(params[1] - 1));
        break;

    case CV_CALIB_ETALON_USER:
        etalonParamCount = 0;

        if( !points || pointCount < 4 )
        {
            assert(0);
            return false;
        }
        break;

    default:
        assert(0);
        return false;
    }

    if( etalonParamCount > 0 )
    {
        arrSize = etalonParamCount * sizeof(etalonParams[0]);
        etalonParams = (double*)cvAlloc( arrSize );
    }

    arrSize = pointCount * sizeof(etalonPoints[0]);

    if( etalonPointCount != pointCount )
    {
        cvFree( &etalonPoints );
        etalonPointCount = pointCount;
        etalonPoints = (CvPoint2D32f*)cvAlloc( arrSize );
    }

    switch( etalonType )
    {
    case CV_CALIB_ETALON_CHESSBOARD:
        {
            int etalonWidth = cvRound( params[0] ) - 1;
            int etalonHeight = cvRound( params[1] ) - 1;
            int x, y, k = 0;

            etalonParams[0] = etalonWidth;
            etalonParams[1] = etalonHeight;
            etalonParams[2] = params[2];

            for( y = 0; y < etalonHeight; y++ )
                for( x = 0; x < etalonWidth; x++ )
                {
                    etalonPoints[k++] = cvPoint2D32f( (etalonWidth - 1 - x)*params[2],
                                                      y*params[2] );
                }
        }
        break;

    case CV_CALIB_ETALON_USER:
        memcpy( etalonParams, params, arrSize );
        memcpy( etalonPoints, points, arrSize );
        break;

    default:
        assert(0);
        return false;
    }

    return true;
}


CvCalibEtalonType
CvCalibFilter::GetEtalon( int* paramCount, const double** params,
                          int* pointCount, const CvPoint2D32f** points ) const
{
    if( paramCount )
        *paramCount = etalonParamCount;

    if( params )
        *params = etalonParams;

    if( pointCount )
        *pointCount = etalonPointCount;

    if( points )
        *points = etalonPoints;

    return etalonType;
}


void CvCalibFilter::SetCameraCount( int count )
{
    Stop();
    
    if( count != cameraCount )
    {
        for( int i = 0; i < cameraCount; i++ )
        {
            cvFree( points + i );
            cvFree( latestPoints + i );
            cvReleaseMat( &undistMap[i][0] );
            cvReleaseMat( &undistMap[i][1] );
            cvReleaseMat( &rectMap[i][0] );
            cvReleaseMat( &rectMap[i][1] );
        }

        memset( latestCounts, 0, sizeof(latestPoints) );
        maxPoints = 0;
        cameraCount = count;
    }
}

   
bool CvCalibFilter::SetFrames( int frames )
{
    if( frames < 5 )
    {
        assert(0);
        return false;
    }
    
    framesTotal = frames;
    return true;
}


void CvCalibFilter::Stop( bool calibrate )
{
    int i, j;
    isCalibrated = false;

    // deallocate undistortion maps
    for( i = 0; i < cameraCount; i++ )
    {
        cvReleaseMat( &undistMap[i][0] );
        cvReleaseMat( &undistMap[i][1] );
        cvReleaseMat( &rectMap[i][0] );
        cvReleaseMat( &rectMap[i][1] );
    }

    if( calibrate && framesAccepted > 0 )
    {
        int n = framesAccepted;
        CvPoint3D32f* buffer =
            (CvPoint3D32f*)cvAlloc( n * etalonPointCount * sizeof(buffer[0]));
        CvMat mat;
        float* rotMatr = (float*)cvAlloc( n * 9 * sizeof(rotMatr[0]));
        float* transVect = (float*)cvAlloc( n * 3 * sizeof(transVect[0]));
        int* counts = (int*)cvAlloc( n * sizeof(counts[0]));

        cvInitMatHeader( &mat, 1, sizeof(CvCamera)/sizeof(float), CV_32FC1, 0 );
        memset( cameraParams, 0, cameraCount * sizeof(cameraParams[0]));

        for( i = 0; i < framesAccepted; i++ )
        {
            counts[i] = etalonPointCount;
            for( j = 0; j < etalonPointCount; j++ )
                buffer[i * etalonPointCount + j] = cvPoint3D32f( etalonPoints[j].x,
                                                                 etalonPoints[j].y, 0 );
        }

        for( i = 0; i < cameraCount; i++ )
        {
            cvCalibrateCamera( framesAccepted, counts,
                               imgSize, points[i], buffer,
                               cameraParams[i].distortion,
                               cameraParams[i].matrix,
                               transVect, rotMatr, 0 );

            cameraParams[i].imgSize[0] = (float)imgSize.width;
            cameraParams[i].imgSize[1] = (float)imgSize.height;
            
//            cameraParams[i].focalLength[0] = cameraParams[i].matrix[0];
//            cameraParams[i].focalLength[1] = cameraParams[i].matrix[4];

//            cameraParams[i].principalPoint[0] = cameraParams[i].matrix[2];
//            cameraParams[i].principalPoint[1] = cameraParams[i].matrix[5];

            memcpy( cameraParams[i].rotMatr, rotMatr, 9 * sizeof(rotMatr[0]));
            memcpy( cameraParams[i].transVect, transVect, 3 * sizeof(transVect[0]));

            mat.data.ptr = (uchar*)(cameraParams + i);
            
            /* check resultant camera parameters: if there are some INF's or NAN's,
               stop and reset results */
            if( !cvCheckArr( &mat, CV_CHECK_RANGE | CV_CHECK_QUIET, -10000, 10000 ))
                break;
        }


        isCalibrated = i == cameraCount;

        {/* calibrate stereo cameras */
            if( cameraCount == 2 )
            {
                stereo.camera[0] = &cameraParams[0];
                stereo.camera[1] = &cameraParams[1];

                icvStereoCalibration( framesAccepted, counts,
                                   imgSize,
                                   points[0],points[1],
                                   buffer,
                                   &stereo);

                for( i = 0; i < 9; i++ )
                {
                    stereo.fundMatr[i] = stereo.fundMatr[i];
                }
                
            }

        }

        cvFree( &buffer );
        cvFree( &counts );
        cvFree( &rotMatr );
        cvFree( &transVect );
    }

    framesAccepted = 0;
}


bool CvCalibFilter::FindEtalon( IplImage** imgs )
{
    return FindEtalon( (CvMat**)imgs );
}


bool CvCalibFilter::FindEtalon( CvMat** mats )
{
    bool result = true;

    if( !mats || etalonPointCount == 0 )
    {
        assert(0);
        result = false;
    }

    if( result )
    {
        int i, tempPointCount0 = etalonPointCount*2;

        for( i = 0; i < cameraCount; i++ )
        {
            if( !latestPoints[i] )
                latestPoints[i] = (CvPoint2D32f*)
                    cvAlloc( tempPointCount0*2*sizeof(latestPoints[0]));
        }

        for( i = 0; i < cameraCount; i++ )
        {
            CvSize size;
            int tempPointCount = tempPointCount0;
            bool found = false;

            if( !CV_IS_MAT(mats[i]) && !CV_IS_IMAGE(mats[i]))
            {
                assert(0);
                break;
            }

            size = cvGetSize(mats[i]);

            if( size.width != imgSize.width || size.height != imgSize.height )
            {
                imgSize = size;
            }

            if( !grayImg || grayImg->width != imgSize.width ||
                grayImg->height != imgSize.height )
            {
                cvReleaseMat( &grayImg );
                cvReleaseMat( &tempImg );
                grayImg = cvCreateMat( imgSize.height, imgSize.width, CV_8UC1 );
                tempImg = cvCreateMat( imgSize.height, imgSize.width, CV_8UC1 );
            }

            if( !storage )
                storage = cvCreateMemStorage();

            switch( etalonType )
            {
            case CV_CALIB_ETALON_CHESSBOARD:
                if( CV_MAT_CN(cvGetElemType(mats[i])) == 1 )
                    cvCopy( mats[i], grayImg );
                else
                    cvCvtColor( mats[i], grayImg, CV_BGR2GRAY );
                found = cvFindChessBoardCornerGuesses( grayImg, tempImg, storage,
                                                       cvSize( cvRound(etalonParams[0]),
                                                       cvRound(etalonParams[1])),
                                                       latestPoints[i], &tempPointCount ) != 0;
                if( found )
                    cvFindCornerSubPix( grayImg, latestPoints[i], tempPointCount,
                                        cvSize(5,5), cvSize(-1,-1),
                                        cvTermCriteria(CV_TERMCRIT_ITER|CV_TERMCRIT_EPS,10,0.1));
                break;
            default:
                assert(0);
                result = false;
                break;
            }

            latestCounts[i] = found ? tempPointCount : -tempPointCount;
            result = result && found;
        }
    }

    if( storage )
        cvClearMemStorage( storage );

    return result;
}


bool CvCalibFilter::Push( const CvPoint2D32f** pts )
{
    bool result = true;
    int i, newMaxPoints = etalonPointCount*(MAX(framesAccepted,framesTotal) + 1);

    isCalibrated = false;

    if( !pts )
    {
        for( i = 0; i < cameraCount; i++ )
            if( latestCounts[i] <= 0 )
                return false;
        pts = (const CvPoint2D32f**)latestPoints;
    }

    for( i = 0; i < cameraCount; i++ )
    {
        if( !pts[i] )
        {
            assert(0);
            break;
        }

        if( maxPoints < newMaxPoints )
        {
            CvPoint2D32f* prev = points[i];
            cvFree( points + i );
            points[i] = (CvPoint2D32f*)cvAlloc( newMaxPoints * sizeof(prev[0]));
            memcpy( points[i], prev, maxPoints * sizeof(prev[0]));
        }

        memcpy( points[i] + framesAccepted*etalonPointCount, pts[i],
                etalonPointCount*sizeof(points[0][0]));
    }

    if( maxPoints < newMaxPoints )
        maxPoints = newMaxPoints;

    result = i == cameraCount;

    if( ++framesAccepted >= framesTotal )
        Stop( true );
    return result;
}


bool CvCalibFilter::GetLatestPoints( int idx, CvPoint2D32f** pts,
                                     int* count, bool* found )
{
    int n;
    
    if( (unsigned)idx >= (unsigned)cameraCount ||
        !pts || !count || !found )
    {
        assert(0);
        return false;
    }
    
    n = latestCounts[idx];
    
    *found = n > 0;
    *count = abs(n);
    *pts = latestPoints[idx];

    return true;
}


void CvCalibFilter::DrawPoints( IplImage** dst )
{
    DrawPoints( (CvMat**)dst );
}


void CvCalibFilter::DrawPoints( CvMat** dstarr )
{
    int i, j;

    if( !dstarr )
    {
        assert(0);
        return;
    }

    if( latestCounts )
    {
        for( i = 0; i < cameraCount; i++ )
        {
            if( dstarr[i] && latestCounts[i] )
            {
                CvMat dst_stub, *dst;
                int count = 0;
                bool found = false;
                CvPoint2D32f* pts = 0;

                GetLatestPoints( i, &pts, &count, &found );

                dst = cvGetMat( dstarr[i], &dst_stub );

                static const CvScalar line_colors[] =
                {
                    {{0,0,255}},
                    {{0,128,255}},
                    {{0,200,200}},
                    {{0,255,0}},
                    {{200,200,0}},
                    {{255,0,0}},
                    {{255,0,255}}
                };

                const int colorCount = sizeof(line_colors)/sizeof(line_colors[0]);
                const int r = 4;
                CvScalar color = line_colors[0];
                CvPoint prev_pt = { 0, 0};

                for( j = 0; j < count; j++ )
                {
                    CvPoint pt;
                    pt.x = cvRound(pts[j].x);
                    pt.y = cvRound(pts[j].y);

                    if( found )
                    {
                        if( etalonType == CV_CALIB_ETALON_CHESSBOARD )
                            color = line_colors[(j/cvRound(etalonParams[0]))%colorCount];
                        else
                            color = CV_RGB(0,255,0);

                        if( j != 0 )
                            cvLine( dst, prev_pt, pt, color, 1, CV_AA );
                    }

                    cvLine( dst, cvPoint( pt.x - r, pt.y - r ),
                            cvPoint( pt.x + r, pt.y + r ), color, 1, CV_AA );

                    cvLine( dst, cvPoint( pt.x - r, pt.y + r),
                            cvPoint( pt.x + r, pt.y - r), color, 1, CV_AA );

                    cvCircle( dst, pt, r+1, color, 1, CV_AA );

                    prev_pt = pt;
                }
            }
        }
    }
}


/* Get total number of frames and already accepted pair of frames */
int CvCalibFilter::GetFrameCount( int* total ) const
{
    if( total )
        *total = framesTotal;

    return framesAccepted;
}


/* Get camera parameters for specified camera. If camera is not calibrated
   the function returns 0 */
const CvCamera* CvCalibFilter::GetCameraParams( int idx ) const
{
    if( (unsigned)idx >= (unsigned)cameraCount )
    {
        assert(0);
        return 0;
    }
    
    return isCalibrated ? cameraParams + idx : 0;
}


/* Get camera parameters for specified camera. If camera is not calibrated
   the function returns 0 */
const CvStereoCamera* CvCalibFilter::GetStereoParams() const
{
    if( !(isCalibrated && cameraCount == 2) )
    {
        assert(0);
        return 0;
    }
    
    return &stereo;
}


/* Sets camera parameters for all cameras */
bool CvCalibFilter::SetCameraParams( CvCamera* params )
{
    CvMat mat;
    int arrSize;
    
    Stop();
    
    if( !params )
    {
        assert(0);
        return false;
    }

    arrSize = cameraCount * sizeof(params[0]);

    cvInitMatHeader( &mat, 1, cameraCount * (arrSize/sizeof(float)),
                     CV_32FC1, params );
    cvCheckArr( &mat, CV_CHECK_RANGE, -10000, 10000 );

    memcpy( cameraParams, params, arrSize );
    isCalibrated = true;

    return true;
}


bool CvCalibFilter::SaveCameraParams( const char* filename )
{
    if( isCalibrated )
    {
        int i, j;
        
        FILE* f = fopen( filename, "w" );

        if( !f ) return false;

        fprintf( f, "%d\n\n", cameraCount );

        for( i = 0; i < cameraCount; i++ )
        {
            for( j = 0; j < (int)(sizeof(cameraParams[i])/sizeof(float)); j++ )
            {
                fprintf( f, "%15.10f ", ((float*)(cameraParams + i))[j] );
            }
            fprintf( f, "\n\n" );
        }

        /* Save stereo params */

        /* Save quad */
        for( i = 0; i < 2; i++ )
        {
            for( j = 0; j < 4; j++ )
            {
                fprintf(f, "%15.10f ", stereo.quad[i][j].x );
                fprintf(f, "%15.10f ", stereo.quad[i][j].y );
            }
            fprintf(f, "\n");
        }

        /* Save coeffs */
        for( i = 0; i < 2; i++ )
        {
            for( j = 0; j < 9; j++ )
            {
                fprintf(f, "%15.10lf ", stereo.coeffs[i][j/3][j%3] );
            }
            fprintf(f, "\n");
        }


        fclose(f);
        return true;
    }

    return true;
}


bool CvCalibFilter::LoadCameraParams( const char* filename )
{
    int i, j;
    int d = 0;
    FILE* f = fopen( filename, "r" );

    isCalibrated = false;

    if( !f ) return false;

    if( fscanf( f, "%d", &d ) != 1 || d <= 0 || d > 10 )
        return false;

    SetCameraCount( d );
    
    for( i = 0; i < cameraCount; i++ )
    {
        for( j = 0; j < (int)(sizeof(cameraParams[i])/sizeof(float)); j++ )
        {
            fscanf( f, "%f", &((float*)(cameraParams + i))[j] );
        }
    }


    /* Load stereo params */

    /* load quad */
    for( i = 0; i < 2; i++ )
    {
        for( j = 0; j < 4; j++ )
        {
            fscanf(f, "%f ", &(stereo.quad[i][j].x) );
            fscanf(f, "%f ", &(stereo.quad[i][j].y) );
        }
    }

    /* Load coeffs */
    for( i = 0; i < 2; i++ )
    {
        for( j = 0; j < 9; j++ )
        {
            fscanf(f, "%lf ", &(stereo.coeffs[i][j/3][j%3]) );
        }
    }
    
    
    fclose(f);

    stereo.warpSize = cvSize( cvRound(cameraParams[0].imgSize[0]), cvRound(cameraParams[0].imgSize[1]));

    isCalibrated = true;
    
    return true;
}


bool CvCalibFilter::Rectify( IplImage** srcarr, IplImage** dstarr )
{
    return Rectify( (CvMat**)srcarr, (CvMat**)dstarr );
}

bool CvCalibFilter::Rectify( CvMat** srcarr, CvMat** dstarr )
{
    int i;

    if( !srcarr || !dstarr )
    {
        assert(0);
        return false;
    }

    if( isCalibrated && cameraCount == 2 )
    {
        for( i = 0; i < cameraCount; i++ )
        {
            if( srcarr[i] && dstarr[i] )
            {
                IplImage src_stub, *src;
                IplImage dst_stub, *dst;

                src = cvGetImage( srcarr[i], &src_stub );
                dst = cvGetImage( dstarr[i], &dst_stub );

                if( src->imageData == dst->imageData )
                {
                    if( !undistImg ||
                        undistImg->width != src->width ||
                        undistImg->height != src->height ||
                        CV_MAT_CN(undistImg->type) != src->nChannels )
                    {
                        cvReleaseMat( &undistImg );
                        undistImg = cvCreateMat( src->height, src->width,
                                                 CV_8U + (src->nChannels-1)*8 );
                    }
                    cvCopy( src, undistImg );
                    src = cvGetImage( undistImg, &src_stub );
                }

                cvZero( dst );

                if( !rectMap[i][0] || rectMap[i][0]->width != src->width ||
                    rectMap[i][0]->height != src->height )
                {
                    cvReleaseMat( &rectMap[i][0] );
                    cvReleaseMat( &rectMap[i][1] );
                    rectMap[i][0] = cvCreateMat(stereo.warpSize.height,stereo.warpSize.width,CV_32FC1);
                    rectMap[i][1] = cvCreateMat(stereo.warpSize.height,stereo.warpSize.width,CV_32FC1);
                    cvComputePerspectiveMap(stereo.coeffs[i], rectMap[i][0], rectMap[i][1]);
                }
                cvRemap( src, dst, rectMap[i][0], rectMap[i][1] );
            }
        }
    }
    else
    {
        for( i = 0; i < cameraCount; i++ )
        {
            if( srcarr[i] != dstarr[i] )
                cvCopy( srcarr[i], dstarr[i] );
        }
    }

    return true;
}

bool CvCalibFilter::Undistort( IplImage** srcarr, IplImage** dstarr )
{
    return Undistort( (CvMat**)srcarr, (CvMat**)dstarr );
}


bool CvCalibFilter::Undistort( CvMat** srcarr, CvMat** dstarr )
{
    int i;

    if( !srcarr || !dstarr )
    {
        assert(0);
        return false;
    }

    if( isCalibrated )
    {
        for( i = 0; i < cameraCount; i++ )
        {
            if( srcarr[i] && dstarr[i] )
            {
                CvMat src_stub, *src;
                CvMat dst_stub, *dst;

                src = cvGetMat( srcarr[i], &src_stub );
                dst = cvGetMat( dstarr[i], &dst_stub );

                if( src->data.ptr == dst->data.ptr )
                {
                    if( !undistImg || undistImg->width != src->width ||
                        undistImg->height != src->height ||
                        CV_ARE_TYPES_EQ( undistImg, src ))
                    {
                        cvReleaseMat( &undistImg );
                        undistImg = cvCreateMat( src->height, src->width, src->type );
                    }

                    cvCopy( src, undistImg );
                    src = undistImg;
                }

            #if 1
                {
                CvMat A = cvMat( 3, 3, CV_32FC1, cameraParams[i].matrix );
                CvMat k = cvMat( 1, 4, CV_32FC1, cameraParams[i].distortion );

                if( !undistMap[i][0] || undistMap[i][0]->width != src->width ||
                     undistMap[i][0]->height != src->height )
                {
                    cvReleaseMat( &undistMap[i][0] );
                    cvReleaseMat( &undistMap[i][1] );
                    undistMap[i][0] = cvCreateMat( src->height, src->width, CV_32FC1 );
                    undistMap[i][1] = cvCreateMat( src->height, src->width, CV_32FC1 );
                    cvInitUndistortMap( &A, &k, undistMap[i][0], undistMap[i][1] );
                }

                cvRemap( src, dst, undistMap[i][0], undistMap[i][1] );
            #else
                cvUndistort2( src, dst, &A, &k );
            #endif
                }
            }
        }
    }
    else
    {
        for( i = 0; i < cameraCount; i++ )
        {
            if( srcarr[i] != dstarr[i] )
                cvCopy( srcarr[i], dstarr[i] );
        }
    }


    return true;
}