/* AUTHOR: Ravi Prakash
   University of Southern California, Los Angeles
   RANSAC method for point matching, when sift features are given
*/

#include<iostream>
#include<conio.h>
#include<fstream>
#include<cxcore.h>
#include<cvaux.h>
#include<highgui.h>

#define MAXNUMBER 20
#define TOADD 100

using namespace std;

//Converts a string containing a value to a double
double parseDouble(char *str)
{
       double toret=0.0, toMul=0.01;
       int index=0;
       while(str[index]!='\0')
       {
            toMul*=10;
            if(str[index]=='.') break;
            else index++;
       }
       int i=0;
       while(str[i]!='\0')
       {
               if(str[i]=='.') {i++; continue;}
               toret+=toMul*(str[i]-48);
               toMul/=10;
               i++;
       }
       return toret;
}

//Converts a string containing a value to an integer
int parseInt(char *str)
{
       int toret=0, index=0;
       while(str[index]!='\0')
       {
            toret=toret*10+(str[index]-48);
            index++;
       }
       return toret;
}

//Each feature consists of position (x,y), scale, orientation and a 128 integer vector representing the feature.
class feature
{
      public:
      double x, y, scale, orientation;
      int featVect[128];
};

//Linked list for features. Adds only at the end of the LL. Converts Linked list to array
class ll
{
      //A node in the Linked list
      class node
      {
            public:
            feature *now;
            node *next;
      };

      public:
      int noOfNodes;

      node* Start;
      node* last;

      ll()
      {Start=NULL; last=NULL; noOfNodes=0;}

      //Add a node to the Linked list at the end
      void add(feature *f)
      {
           node *n=new node;
           (*n).now=f;
           (*n).next=NULL;
           if(Start==NULL)
           {
                Start=n;
                last=n;
           }
           else
           {
               (*last).next=n;
               last=n;
           }
           noOfNodes++;
      }

      //Store the pointers to the features in the linked list, in an array
      feature** arrize()
      {
                feature** toret = new feature*[noOfNodes];
                node *temp=Start;
                for(int i=0; i<noOfNodes; i++)
                {
                        toret[i]=(*temp).now;
                        temp=(*temp).next;
                }
                return toret;
      }
};

// Abstracts the image's features.
class image
{
      public:
      int noOfFeatures;
      feature** featArr;

      ~image()
      {
              for(int i=0; i<noOfFeatures; i++) delete featArr[i];
              delete featArr;
      }
      
      //Read all the features of the image from a file
      void readKeyFile(char *fileName)
      {
           ifstream file1(fileName);
           if(!file1)
           {
                cout<<"\nFile open failed";
                exit(0);
           }

           ll tempLL;
           char number[MAXNUMBER];
           while(!file1.eof())
           {
                 feature *temp=new feature;
                 file1>>number; (*temp).x=parseDouble(number);
                 file1>>number; (*temp).y=parseDouble(number);
                 file1>>number; (*temp).scale=parseDouble(number);
                 file1>>number; (*temp).orientation=parseDouble(number);
                 for(int i=0; i<128; i++)
                 {
                       file1>>number;
                      (*temp).featVect[i]=parseInt(number);
                 }
                 tempLL.add(temp);
           }
           file1.close();
           featArr=tempLL.arrize();
           noOfFeatures=tempLL.noOfNodes;
      }

      //Print out the position, scale and orientation of all features. For testing only.
      void test()
      {
           feature *f;
           for(int i=0; i<10; i++)
           {
                   f=featArr[i];
                   printf("\n%f %f %f %f ", (*f).x, (*f).y, (*f).scale, (*f).orientation);
                   getch();
           }
      }
};

//Abstracts the matching of two images
class matcher
{
      public:
      image *i1, *i2;
      double **matches;
      int count;

      matcher(image *ti1, image *ti2)
      {
           if(ti1->noOfFeatures<ti2->noOfFeatures) //i1 will always have equal or more no of features than i2. For programming convenience :D
           {      i1=ti2; i2=ti1;
           }
           else
           {      i1=ti1; i2=ti2;
           }

      }

      //find Euclidean distance between two features. Sum of squares of difference
      inline long getEucDist(feature *f1, feature *f2)
      {
             long toret=0;
             for(int i=0; i<128; i++)
             {
                     toret+=((f1->featVect[i])-(f2->featVect[i]))*((f1->featVect[i])-(f2->featVect[i]));
             }
             return toret;
      }

      //Preliminary matching. A feature A in image 1 is matched with a feature B in image 2 iff:
      // 1) the A is the best match for B AND B is the best match for A
      // 2) The ratio between best match and second best match is less than ratio
      int* makePairs(double ratio)
      {
           int *im1FeatMatch=new int[i1->noOfFeatures];
           int *im2FeatMatch=new int[i2->noOfFeatures];

           long best, secBest, temp;
           int bestInd;

           for(int i=0; i<i1->noOfFeatures; i++)
           {
                   best=LONG_MAX;
                   secBest=LONG_MAX;
                   for(int j=0; j<i2->noOfFeatures; j++)
                   {
                           temp=getEucDist(i1->featArr[i],i2->featArr[j]);
                           if (temp<best)
                           {
                                secBest=best;
                                best=temp; bestInd=j;
                           }
                           else if(temp<secBest)
                           {
                                secBest=temp;
                           }
                   }
                   if(ratio*best<secBest) im1FeatMatch[i]=bestInd;
                   else im1FeatMatch[i]=-1;
           }
           printf("\nHalfway thru\n");
           for(int i=0; i<i2->noOfFeatures; i++)
           {
                   best=LONG_MAX;
                   secBest=LONG_MAX;
                   for(int j=0; j<i1->noOfFeatures; j++)
                   {
                           temp=getEucDist(i2->featArr[i],i1->featArr[j]);
                           if (temp<best)
                           {
                                secBest=best;
                                best=temp; bestInd=j;
                           }
                           else if(temp<secBest)
                           {
                                secBest=temp;
                           }
                   }
                   if(ratio*best<secBest) im2FeatMatch[i]=bestInd;
                   else im2FeatMatch[i]=-1;
           }

           count=0;
           for(int i=0; i<i1->noOfFeatures; i++)
           {
                   if(im1FeatMatch[i] != -1)
                          if( im2FeatMatch[ im1FeatMatch[i] ]!=i) im1FeatMatch[i]=-1;
                          else count++;
           }

           delete im2FeatMatch;
           printf("\nAll the way through step 1. \nFeature Points in image 1: %d\nFeature Points in image 2: %d\nNo. of matches found: %d", i1->noOfFeatures, i2->noOfFeatures, count);
           return im1FeatMatch;
      }
      
      //Calculates the error when this match is selected with <fMat> (fundamental matrix)
      inline double matMul(double *point1, CvMat *fMat, double *point2)
      {
             double temp[3];
             temp[0]=temp[1]=temp[2]=0.0;
             for(int i=0; i<3; i++)
                     for(int j=0; j<3; j++)
                             temp[i]+=point1[j]*fMat->data.fl[i+j*3];
             return (temp[0] * point2[0] + temp[1] * point2[1] + temp[2] * point2[2]);
      }

      //Matches the points using RANSAC
      double** matchPoints(double ratio, int iterations, double deviation)
      {
           int* pairArray, size;
           int index, ranind;
           double sum, bestSum;
           CvMat *points1, *points2, *fMat, *status, *bestFMat;
           
           points1 = cvCreateMat(1,8,CV_32FC2);
           points2 = cvCreateMat(1,8,CV_32FC2);
           status =  cvCreateMat(1,8,CV_8UC1);
           bestFMat = cvCreateMat(1,1,CV_8UC1);
           fMat=cvCreateMat(3,3, CV_32FC1);
           size=i1->noOfFeatures;

           pairArray=makePairs(ratio);

//           printf("Picking random points");getch();

           bestSum=DBL_MAX;
           for(int i=0; i<iterations; i++)
           {
                   index=0;
                   while(index<8)
                   {
                            ranind=rand()%size;
                            if(pairArray[ranind]!=-1)
                            {
                                     points1->data.fl[index*2] = i1->featArr[ranind]->x;
                                     points1->data.fl[index*2+1] = i1->featArr[ranind]->y;
                                     points2->data.fl[index*2] = i2->featArr[pairArray[ranind]]->x;
                                     points2->data.fl[index*2+1] = i2->featArr[pairArray[ranind]]->y;
                                     index++;
                            }
                   }

                   //Randomly, 8 points picked.

                   cvFindFundamentalMat(points1, points2, fMat, CV_FM_8POINT, 1, 0.99, status);

                   //Now calculate error (p1 F p2) for EVERY POINT and sum this error for this choice of 8 points.
                   sum=0.0; double p1[3], p2[3];
                   p1[2]=p2[2]=1;
                   for(int j=0; j<size; j++)
                   {
                           if(pairArray[j]==-1) continue;
                           p1[0]=i1->featArr[j]->x; p1[1]=i1->featArr[j]->y;
                           p2[0]=i2->featArr[pairArray[j]]->x; p2[1]=i2->featArr[pairArray[j]]->y;

                           //Matrix multiplication done linearly... Phew.... ABSOLUTE VALUE OF ERROR NEEDS TO BE SUMMED
//                             sum+=fabs(((i2->featArr[pairArray[j]]->x)* ((i1->featArr[j]->x)*(fMat->data.fl[0])+(i1->featArr[j]->y)*(fMat->data.fl[3])+1*(fMat->data.fl[6]) ) + (i2->featArr[pairArray[j]]->y)* ((i1->featArr[j]->x)*(fMat->data.fl[1])+(i1->featArr[j]->y)*(fMat->data.fl[4])+1*(fMat->data.fl[7])) + 1 * ((i1->featArr[j]->x)*(fMat->data.fl[2])+(i1->featArr[j]->y)*(fMat->data.fl[5])+1*(fMat->data.fl[8]))));
                               sum+=fabs(matMul(p1, fMat, p2));
                   }
                   //repeat this process. Remember the matrix that gave the lowest sum. This is the true F matrix.
                   if(sum<bestSum)
                   {
                                  CvMat *temp=bestFMat;
                                  cvReleaseMat(&temp);
                                  bestFMat=cvCloneMat(fMat);
                                  bestSum=sum;
                   }

           }
           // Now that we have F, choose the pairs of points for which (p1 F p2) is below a certain threshhold. These are the correct matches.

           bestSum/=count; //Should be /(count-8) , but its the same
           double temp;
           count=0;
           double p1[3], p2[3];
           p1[2]=p2[2]=1;
           for(int j=0; j<size; j++)
           {
                   if(pairArray[j]==-1) continue;
                   //Matrix multiplication done linearly... Phew.... ABSOLUTE VALUE OF ERROR NEEDS TO BE CHECKED
//                   temp=fabs(((i2->featArr[pairArray[j]]->x)* ((i1->featArr[j]->x)*(bestFMat->data.fl[0])+(i1->featArr[j]->y)*(bestFMat->data.fl[3])+1*(bestFMat->data.fl[6]) ) + (i2->featArr[pairArray[j]]->y)* ((i1->featArr[j]->x)*(bestFMat->data.fl[1])+(i1->featArr[j]->y)*(bestFMat->data.fl[4])+1*(bestFMat->data.fl[7])) + 1 * ((i1->featArr[j]->x)*(bestFMat->data.fl[2])+(i1->featArr[j]->y)*(bestFMat->data.fl[5])+1*(bestFMat->data.fl[8]))));
                   p1[0]=i1->featArr[j]->x; p1[1]=i1->featArr[j]->y;
                   p2[0]=i2->featArr[pairArray[j]]->x; p2[1]=i2->featArr[pairArray[j]]->y;
                   temp=fabs(matMul(p1, bestFMat, p2));
                   if(temp<bestSum*deviation)
                         count++;
           }
//           printf("\n\nBest:%.2f", bestSum);

           matches=new double*[count];
           for(int j=0; j<count; j++) matches[j]=new double[4];
           int ind=0;
           for(int j=0; j<size; j++)
           {
                   if(pairArray[j]==-1) continue;
                   //Matrix multiplication done linearly... Phew.... ABSOLUTE VALUE OF ERROR NEEDS TO BE CHECKED
//                   temp=fabs(((i2->featArr[pairArray[j]]->x)* ((i1->featArr[j]->x)*(bestFMat->data.fl[0])+(i1->featArr[j]->y)*(bestFMat->data.fl[3])+1*(bestFMat->data.fl[6]) ) + (i2->featArr[pairArray[j]]->y)* ((i1->featArr[j]->x)*(bestFMat->data.fl[1])+(i1->featArr[j]->y)*(bestFMat->data.fl[4])+1*(bestFMat->data.fl[7])) + 1 * ((i1->featArr[j]->x)*(bestFMat->data.fl[2])+(i1->featArr[j]->y)*(bestFMat->data.fl[5])+1*(bestFMat->data.fl[8]))));
                   p1[0]=i1->featArr[j]->x; p1[1]=i1->featArr[j]->y;
                   p2[0]=i2->featArr[pairArray[j]]->x; p2[1]=i2->featArr[pairArray[j]]->y;
                   temp=fabs(matMul(p1, bestFMat, p2));
                   if(temp<bestSum*deviation)
                   {
                         matches[ind][0]=i1->featArr[j]->x; matches[ind][1]=i1->featArr[j]->y;
                         matches[ind][2]=i2->featArr[pairArray[j]]->x; matches[ind][3]=i2->featArr[pairArray[j]]->y;
                         ind++;
                   }
           }

           delete pairArray;
           cvReleaseMat(&points1); cvReleaseMat(&points2); cvReleaseMat(&status); cvReleaseMat(&fMat); cvReleaseMat(&bestFMat);
           return matches;
      }
    
};

//Puts a colored rectangle of size <size> and color <rgb> at position (<x>, <y>) in image <img>
void putRect(IplImage *img, int x, int y, int size, int* rgb)
{
     for(int i=-size/2; i<size/2; i++)
             for(int j=-size/2; j<size/2; j++)
             {
                     if(y+j>img->height||y+j<0||x+i>img->width||x+y<0) continue;
                     ((uchar*)(img->imageData + img->widthStep*(y+j) ))[(x+i)*3]=rgb[0];
                     ((uchar*)(img->imageData + img->widthStep*(y+j) ))[(x+i)*3+1]=rgb[1];
                     ((uchar*)(img->imageData + img->widthStep*(y+j) ))[(x+i)*3+2]=rgb[2];
             }
}

int main()
{
    image i1, i2;
    double **temp;
    IplImage *im1=cvLoadImage("img3.pgm", 1);
    IplImage *im2=cvLoadImage("img5.pgm", 1);
    i1.readKeyFile("imgs3.key");
    i2.readKeyFile("imgs5.key");
    printf("\nRead both files. Features loaded in memory.\n\n");
    matcher M1(&i1, &i2);
    //More ratio -> fewer points matched. // More deviation allowed -> more points matched
    temp=M1.matchPoints(1.1, 500, 1.9);

    //Mark matched points in images now
    for(int i=0; i<M1.count; i++)
    {
            int rgb[3];
            rgb[0]=rand()%255; rgb[1]=rand()%255; rgb[2]=rand()%255;
            putRect(im1, (int) temp[i][0], (int) temp[i][1], 11, rgb);
            putRect(im2, (int) temp[i][2], (int) temp[i][3], 11, rgb);
    }
//            printf("\n%.2f, %.2f matched with %.2f, %.2f", temp[i][0], temp[i][1], temp[i][2], temp[i][3] );

    printf("\n\nTotal matches: %d", M1.count);
    cvNamedWindow("Image 1 Window", 1);
    cvNamedWindow("Image 2 Window", 1);

    cvShowImage("Image 1 Window", im1);
    cvShowImage("Image 2 Window", im2);
    
    cvSaveImage("Output1.jpg", im1);
    cvSaveImage("Output2.jpg", im2);

    cvWaitKey(0);

    cvDestroyWindow("Image 1 Window");
    cvDestroyWindow("Image 2 Window");

    cvReleaseImage(&im1);
    cvReleaseImage(&im2);
    return 0;
}