#include "Registration/SurfRegistration/SurfRegistration.h"

#include <cv.h>

#include <cxcore.h>

#include <highgui.h>

using namespace cv;

#define BROKEN 1

//Code copied from cvsurf.cpp

/* Wavelet size at first layer of first octave. */ 

const int HAAR_SIZE0 = 9;    

/* Wavelet size increment between layers. This should be an even number, 

 such that the wavelet sizes in an octave are either all even or all odd.

 This ensures that when looking for the neighbours of a sample, the layers

 above and below are aligned correctly. */

const int HAAR_SIZE_INC = 6;

static int getPointOctave(const CvSURFPoint& kpt, const CvSURFParams& params)

{

    int octave = 0, layer = 0, best_octave = 0;

    float min_diff = FLT_MAX;

    for( octave = 1; octave < params.nOctaves; octave++ )

        for( layer = 0; layer < params.nOctaveLayers; layer++ )

        {

            float diff = std::abs(kpt.size - (float)((HAAR_SIZE0 + HAAR_SIZE_INC*layer) << octave));

            if( min_diff > diff )

            {

                min_diff = diff;

                best_octave = octave;

                if( min_diff == 0 )

                    return best_octave;

            }

        }

    return best_octave;

}

//End of Code copied from cvsurf.cpp

SurfRegistration::SurfRegistration() {

        surf_class=cv::SURF(0.1);

};

SurfRegistration::~SurfRegistration() {};

bool SurfRegistration::GetHomography(IplImage *img1_gray, IplImage *img2_gray, 

                                                                         CvMat *homography)

{

        assert(homography->cols==3 && homography->rows==3); //assert type too?

        Mat going_out;

        Mat no_info;

        vector<KeyPoint> img1_points, img2_points;

        vector<float> img1_descriptors, img2_descriptors;

        Seq<CvSURFPoint> img1_points_seq, img2_points_seq;

        CvSeq* img1_descriptors_seq = 0;

        CvSeq* img2_descriptors_seq = 0;

        CvSURFParams params = cvSURFParams(500,0);

        CvMemStorage* storage = cvCreateMemStorage(0);

        Mat match_indices, match_distances;

        vector<Point2f> img1_inliers, img2_inliers;

        cv::flann::AutotunedIndexParams my_params;

        my_params.build_weight=1;

        my_params.memory_weight=.1;

        my_params.sample_fraction=.5;

        my_params.target_precision=0.9;

        Mat tmp5=Mat(img1_gray);

#ifdef BROKEN

        surf_class(tmp5,no_info,img1_points,img1_descriptors);

        assert( _CrtCheckMemory()  );

        return true; // it dies in the return when cleaning up img1_descriptors

#endif

        cvExtractSURF (img1_gray,0,&img1_points_seq.seq,&img1_descriptors_seq,storage,params);

        cvExtractSURF (img2_gray,0,&img2_points_seq.seq,&img2_descriptors_seq,storage,params);

        //surf_class(Mat(img2_gray),no_info,img2_points,img2_descriptors);

        //This code was essentially copied from operator() in cvsurf.cpp

    img1_descriptors.resize(img1_descriptors_seq ? img1_descriptors_seq->total*img1_descriptors_seq->elem_size/sizeof(float) : 0);

    if(img1_descriptors.size() != 0)

        cvCvtSeqToArray(img1_descriptors_seq, &img1_descriptors[0]);

    img2_descriptors.resize(img2_descriptors_seq ? img2_descriptors_seq->total*img2_descriptors_seq->elem_size/sizeof(float) : 0);

    if(img2_descriptors.size() != 0)

        cvCvtSeqToArray(img2_descriptors_seq, &img2_descriptors[0]);

        {

        Seq<CvSURFPoint>::iterator it = img1_points_seq.begin();

        size_t i, n = img1_points_seq.size();

        img1_points.resize(n);

        for( i = 0; i < n; i++, ++it )

        {

            const CvSURFPoint& kpt = *it;

            img1_points[i] = KeyPoint(kpt.pt, (float)kpt.size, kpt.dir,

                                    kpt.hessian, getPointOctave(kpt, params));

        }

    }

        {

        Seq<CvSURFPoint>::iterator it = img2_points_seq.begin();

        size_t i, n = img2_points_seq.size();

        img2_points.resize(n);

        for( i = 0; i < n; i++, ++it )

        {

            const CvSURFPoint& kpt = *it;

            img2_points[i] = KeyPoint(kpt.pt, (float)kpt.size, kpt.dir,

                                    kpt.hessian, getPointOctave(kpt, params));

        }

    }

        //end of copied code...

        Mat tmp=Mat(img1_descriptors).reshape(1,img1_descriptors.size()/surf_class.descriptorSize());

        cv::flann::Index flann_class(tmp,

                //Mat(img1_descriptors).reshape(1,img1_descriptors.size()/surf_class.descriptorSize()),

                my_params);

        match_indices.create(img2_points.size(),2,CV_32S);

        match_distances.create(img2_points.size(),2,CV_32F);

        Mat tmp2=Mat(img2_descriptors).reshape(1,img2_descriptors.size()/surf_class.descriptorSize());

        flann_class.knnSearch(tmp2,match_indices,match_distances,2,

                cv::flann::SearchParams()); //SearchParams is ignored as I used auto-create above

        for (int i=0;i<img2_points.size();i++) {

                float dist1=match_distances.at<float>(i,0);

                if (dist1<0.6*match_distances.at<float>(i,1)) {

                        img1_inliers.push_back(img1_points.at(match_indices.at<int>(i,0)).pt);

                        img2_inliers.push_back(img2_points.at(i).pt);

                }

        }        

        if (img1_inliers.size()<4) {

                printf ("The number of inliers was only %d\n",img1_inliers.size());

                return false;

        }

        // All the point matches have now been found.  Run RANSAC homography finder to compute

        // the homography

        going_out=findHomography(Mat(img1_inliers),Mat(img2_inliers),(int)cv::RANSAC,DISTANCE_THRESHOLD);

        for (int i=0;i<3;i++)

                memcpy(homography->data.ptr+homography->step*i,going_out.ptr(i),sizeof(float)*3);

        return true;

}