// Imagine++ project
// Project:  Fundamental
// Author:   Pascal Monasse
// Date:     2013/10/08

#include "./Imagine/Features.h"
#include <Imagine/Graphics.h>
#include <Imagine/LinAlg.h>
#include <vector>
#include <cstdlib>
#include <ctime>
using namespace Imagine;
using namespace std;

static const float BETA = 0.01f; // Probability of failure

struct Match {
    float x1, y1, x2, y2;
};

// Display SIFT points and fill vector of point correspondences
void algoSIFT(Image<Color,2> I1, Image<Color,2> I2,
              vector<Match>& matches) {
    // Find interest points
    SIFTDetector D;
    D.setFirstOctave(-1);
    Array<SIFTDetector::Feature> feats1 = D.run(I1);
    drawFeatures(feats1, Coords<2>(0,0));
    cout << "Im1: " << feats1.size() << flush;
    Array<SIFTDetector::Feature> feats2 = D.run(I2);
    drawFeatures(feats2, Coords<2>(I1.width(),0));
    cout << " Im2: " << feats2.size() << flush;

    const double MAX_DISTANCE = 100.0*100.0;
    for(size_t i=0; i < feats1.size(); i++) {
        SIFTDetector::Feature f1=feats1[i];
        for(size_t j=0; j < feats2.size(); j++) {
            double d = squaredDist(f1.desc, feats2[j].desc);
            if(d < MAX_DISTANCE) {
                Match m;
                m.x1 = f1.pos.x();
                m.y1 = f1.pos.y();
                m.x2 = feats2[j].pos.x();
                m.y2 = feats2[j].pos.y();
                matches.push_back(m);
            }
        }
    }
}

// 8 points method if inliers is of size 8
// Otherwise, solve the minimization problem "min ||Af||2 subject to ||f|| = 1"
FMatrix<float,3,3> computeSingleF(const vector<Match>& matches, const vector<int>& inliers){
    assert(inliers.size() >= 8);
    // create normalization matrix
    FMatrix<float,3,3> N(0.0);
    N(0,0)=N(1,1)=0.001;
    N(2,2)=1;
    // create A
    Matrix<float> A(max(9,(int)inliers.size()), 9);
    Match m;
    for(size_t ii=0; ii<inliers.size(); ii++){
        m = matches[inliers[ii]];
        m.x1 *= N(0,0); m.x2 *= N(0,0); m.y1 *= N(0,0); m.y2 *= N(0,0);
        A(ii,0) = m.x1*m.x2; A(ii,1) = m.x1*m.y2; A(ii,2) = m.x1;
        A(ii,3) = m.y1*m.x2; A(ii,4) = m.y1*m.y2; A(ii,5) = m.y1;
        A(ii,6) = m.x2; A(ii,7) = m.y2; A(ii,8) = 1.0;
    }
    // it is easier to solve svd with square matrices, add a 9th equation T(0)f = 0
    if(inliers.size()==8)
        for(size_t jj=0; jj<9; jj++)
            A(8,jj) = 0;
    // Solve svd
    Vector<float> S;
    Matrix<float> U, Vt;
    svd(A,U,S,Vt);
    // Copy F from Vt;
    FMatrix<float,3,3> F;
    F(0,0) = Vt(8,0); F(0,1) = Vt(8,1); F(0,2) = Vt(8,2);
    F(1,0) = Vt(8,3); F(1,1) = Vt(8,4); F(1,2) = Vt(8,5);
    F(2,0) = Vt(8,6); F(2,1) = Vt(8,7); F(2,2) = Vt(8,8);
    return N*F*N;
}

// Finds inliers from a list of matches
vector<int> findInliers (vector<Match>& matches, FMatrix<float,3,3>& F, float d){
    vector<int> inliers;
    for(size_t ii=0; ii<matches.size(); ii++){
        Match m = matches[ii];
        FVector<float, 3> x(m.x1,m.y1,1);
        FVector<float, 3> y(m.x2,m.y2,1);
        FVector<float, 3> Fy = F*y;
        if( abs(x*Fy) / sqrt(Fy[0]*Fy[0] + Fy[1]*Fy[1]) < d)
            inliers.push_back(ii);
    }
    return inliers;
}

// RANSAC algorithm to compute F from point matches (8-point algorithm)
// Parameter matches is filtered to keep only inliers as output.
FMatrix<float,3,3> computeF(vector<Match>& matches) {
    const float distMax = 1.5f; // Pixel error for inlier/outlier discrimination
    int Niter=100000; // Adjusted dynamically
    FMatrix<float,3,3> bestF;
    vector<int> bestInliers;

    FMatrix<float,3,3> F;
    for(int ii=0; ii<Niter; ii++){
        vector<int> inliers;
        while(inliers.size()<8){
            int kk = rand() % matches.size();
            if(find(inliers.begin(), inliers.end(), kk) == inliers.end())
                inliers.push_back(kk);
        }
        F = computeSingleF(matches, inliers);
        inliers = findInliers(matches, F, distMax);
        if(inliers.size()> bestInliers.size()){
            bestInliers = inliers;
            Niter = (int)ceil(log(BETA)/log(1-pow((float)inliers.size()/(float)matches.size(), 8)));
        }
    }
    cout << bestInliers.size() << " inliers" << endl;
    cout << Niter << " iterations" << endl;
    // refine F with all the inliers
    bestF = computeSingleF(matches, bestInliers);

    // Updating matches with inliers only
    vector<Match> all=matches;
    matches.clear();
    for(size_t i=0; i<bestInliers.size(); i++)
        matches.push_back(all[bestInliers[i]]);
    return bestF;
}

// Expects clicks in one image and show corresponding line in other image.
// Stop at right-click.
void displayEpipolar(Image<Color> I1, Image<Color> I2,
                     const FMatrix<float,3,3>& F) {
    while(true) {
        int x,y;
        if(getMouse(x,y) == 3)
            break;
        drawCircle(x, y, 5, RED);
        drawCircle(x, y, 4, GREEN);
        if(x>I1.width()){
            FVector<float, 3> Y((float)x - I1.width(),(float)y,1.0);
            FVector<float, 3> FY = F*Y;
            drawLine(0,-FY[2]/FY[1],
                    I1.width(),-(FY[2]+FY[0]*I1.width())/FY[1],
                    RED);
        }
        else{
            FVector<float, 3> X((float)x,(float)y,1.0);
            FVector<float, 3> FtX = transpose(F)*X;
            drawLine(I1.width(), (-1)*(FtX[2])/FtX[1],
                    2*I1.width(), (-1)*(FtX[2]+FtX[0]*I1.width())/FtX[1],
                    RED);
        }
    }
}

int main(int argc, char* argv[])
{
    srand((unsigned int)time(0));

    const char* s1 = argc>1? argv[1]: srcPath("im1.jpg");
    const char* s2 = argc>2? argv[2]: srcPath("im2.jpg");

    // Load and display images
    Image<Color,2> I1, I2;
    if( ! load(I1, s1) ||
        ! load(I2, s2) ) {
        cerr<< "Unable to load images" << endl;
        return 1;
    }
    int w = I1.width();
    openWindow(2*w, I1.height());
    display(I1,0,0);
    display(I2,w,0);

    vector<Match> matches;
    algoSIFT(I1, I2, matches);
    cout << " matches: " << matches.size() << endl;
    click();
    
    FMatrix<float,3,3> F = computeF(matches);
    cout << "F="<< endl << F;

    // Redisplay with matches
    display(I1,0,0);
    display(I2,w,0);
    for(size_t i=0; i<matches.size(); i++) {
        Color c(rand()%256,rand()%256,rand()%256);
        fillCircle(matches[i].x1+0, matches[i].y1, 2, c);
        fillCircle(matches[i].x2+w, matches[i].y2, 2, c);        
    }
    click();

    // Redisplay without SIFT points
    display(I1,0,0);
    display(I2,w,0);
    displayEpipolar(I1, I2, F);

    endGraphics();
    return 0;
}