#include "main.h"
#include "Alignment.h"
#include "Statistics.h"
#include "Filter.h"
#include "Arithmetic.h"
#include "File.h"
#include "Color.h"
#include "LinearAlgebra.h"
#include <algorithm>
#include "Geometry.h"
#include "Convolve.h"

// First we define the various types of transformations we may wish to
// solve for. All of these classes lean heavily on the least squares
// solving in LinearAlgebra.h
class Transform {
public:
    virtual ~Transform() {}

    // Add a new constraint (ie (x1,y1) must warp to (x2, y2))
    virtual void addCorrespondence(float x1, float y1, float x2, float y2) = 0;

    // Once all the constraints are added, solve for the optimal warp
    virtual void solve() = 0;

    // After solving, we can apply the warp to a given (x1, y1) to produce (x2, y2)
    virtual void apply(float x1, float y1, float *x2, float *y2) = 0;

    // Different types of transformations required different numbers
    // of constraints in order to produce a model
    virtual int constraintsRequired() = 0;

    // Forget all the correspondences so far, and start again. This is
    // useful in RANSAC if we picked bad initial correspondences.
    virtual void reset() = 0;
};

// Solving for a least squares translation is easy, we just average
// the translations of all the correspondences given.
class Translation : public Transform {
public:
    Translation() {
	reset();
    }

    virtual ~Translation() {}

    void reset() {
	dx = dy = dxSum = dySum = 0;
	count = 0;
    }

    void addCorrespondence(float x1, float y1, float x2, float y2) {
	dxSum += x2 - x1;
	dySum += y2 - y1;
	count++;
    }

    void solve() {
	dx = dxSum / count;
	dy = dySum / count;
    }

    int constraintsRequired() {
	return 1;
    }

    void apply(float x1, float y1, float *x2, float *y2) {
	*x2 = x1 + dx;
	*y2 = y1 + dy;
    }

private:
    // the transform parameters
    float dx, dy;
    // the internal state to keep track of
    float dxSum, dySum;
    int count;
};

// Solving for a least squares similarity (ie rotation, translation,
// and scale) matrix is a little tricker. There are four free
// parameters. We use the LeastSquares solver class to solve it.
class Similarity : public Transform {
public:
    Similarity() {
	reset();
    }

    virtual ~Similarity() {}

    void reset() {
	params[0] = params[1] = params[2] = params[3] = 0;
	solver.reset();
    }

    void addCorrespondence(float x1, float y1, float x2, float y2) {
	{
	    float in[4] = {x1, y1, 1, 0};
	    float out[1] = {x2};
	    solver.addCorrespondence(in, out);
	}
	{
	    float in[4] = {y1, -x1, 0, 1};
	    float out[1] = {y2};
	    solver.addCorrespondence(in, out);
	}
    }

    void solve() {
	solver.solve(params);
    }

    int constraintsRequired() {
	return 2;
    }

    void apply(float x1, float y1, float *x2, float *y2) {
	*x2 = params[0]*x1 + params[1]*y1 + params[2];
	*y2 = params[0]*y1 - params[1]*x1 + params[3];
    }

protected:
    // the transform parameters
    double params[4];

    // the internal state
    LeastSquaresSolver<4, 1> solver;	
};

// Solving for a 2D rigid transformation (rotation and translation but
// no scale) is actually surprisingly hard. We dodge the problem by
// just solving for a similarity transform and altering the parameters
// to remove the scale factor.
class Rigid : public Similarity {
public:
    virtual ~Rigid() {}

    void solve() {
	solver.solve(params);
	double l = sqrt(params[0]*params[0] + params[1]*params[1]);
	params[0] /= l;
	params[1] /= l;
    }
};

// Solving for an affine transform is a fairly standard least squares
// solve, because it's easy to represent an affine transform as a
// 2x3 matrix.
class Affine : public Transform {
public:
    Affine() {
	reset();
    }

    virtual ~Affine() {}

    void reset() {
	params[0] = params[1] = params[2] = params[3] = params[4] = params[5] = 0;
	solver.reset();
    }

    void addCorrespondence(float x1, float y1, float x2, float y2) {
	float in[3] = {x1, y1, 1};
	float out[2] = {x2, y2};
	solver.addCorrespondence(in, out);
    }

    void solve() {
	solver.solve(params);
    }

    int constraintsRequired() {
	return 3;
    }

    void apply(float x1, float y1, float *x2, float *y2) {
	*x2 = params[0]*x1 + params[2]*y1 + params[4];
	*y2 = params[1]*x1 + params[3]*y1 + params[5];
    }

private:
    // the transform parameters
    double params[6];

    // the internal state
    LeastSquaresSolver<3, 2> solver;	
};

// Solving for 2D perspective warps involves some algebraic
// manipulation.  In general, a 2D perspective warp can be expressed
// by a 3x3 matrix, which maps (x1, y1, 1) to some homogeneous
// representation of (x2, y2) - ie (w.x2, w.y2, w).
//
// If you write this out and substitute out w, then shuffle terms
// around to make it linear, you can solve for the eight parameters
// of the transform. Why only 8 parameters in a 3x3 matrix? Because
// its outputs are homogeneous vectors, the matrix is invariant to
// scale, so we can assume WLOG that the bottom right entry is 1.

class Perspective : public Transform {
public:
    Perspective() {
	reset();
    }

    virtual ~Perspective() {}
 
    void reset() {
	for (int i = 0; i < 8; i++) params[i] = 0;
	solver.reset();
    }

    void addCorrespondence(float x1, float y1, float x2, float y2) {
	{
	    float in[8] = {-x1*x2, -y1*x2, x1, y1, 1, 0, 0, 0};
	    float out[1] = {x2};
	    solver.addCorrespondence(in, out);
	}
	{
	    float in[8] = {-x1*y2, -y1*y2, 0, 0, 0, x1, y1, 1};
	    float out[1] = {y2};
	    solver.addCorrespondence(in, out);
	}
    }

    void solve() {
	solver.solve(params);
    }

    int constraintsRequired() {
	return 4;
    }

    void apply(float x1, float y1, float *x2, float *y2) {
	float w = 1.0f/(params[0]*x1 + params[1]*y1 + 1);	
	*x2 = (params[2]*x1 + params[3]*y1 + params[4])*w;
	*y2 = (params[5]*x1 + params[6]*y1 + params[7])*w;
    }

private:
    // the transform parameters
    double params[8];

    // the internal state
    LeastSquaresSolver<8, 1> solver;	    
};

// A Digest is a data structure that gathers together all the features
// extracted from a single image.
class Digest {
public:

    // An image feature is a local maximum in the image (after
    // applying some corner detector filter), with a descriptor
    // vector. In this case we use a 7x7 patch.
    struct Feature : public LocalMaxima::Maximum {
    public:
	Feature(LocalMaxima::Maximum m, Window im) {
	    x = m.x;
	    y = m.y;
	    t = floor(m.t + 0.5);
	    int ix = (int)(x + 0.5);
	    int iy = (int)(y + 0.5);
	    value = m.value;
	    assert( ( ix > 2 && ix < im.width-3 &&
		      iy > 2 && iy < im.height-3 &&
		    t >= 0 && t < im.frames ),
		   "Requested feature out of bounds\n");
	    patch = Window(im, 
			   (int)t, ix-3, iy-3,
			   1, 7, 7);
	    usage = 0;
	}

	// Distance between two features is just the sum of squared differences between the two patches.
	float distance(Feature *other) {
	    float dist = 0;
	    for (int t = 0; t < patch.frames; t++) {
		for (int y = 0; y < patch.height; y++) {
		    float *thisPtr = patch(t, 0, y);
		    float *otherPtr = other->patch(t, 0, y);
		    for (int x = 0; x < patch.channels * patch.width; x++) {
			float d = *thisPtr++ - *otherPtr++;
			dist += d*d;
		    }
		}
	    }
	    return dist;
	}

	// It's useful to keep track of how many times any one given
	// features is used, so we don't depend too heavily on a
	// single feature.
	int usage;	

	Window patch;
    };
    
    // A correspondences is a pair of features that hopefully match.
    struct Correspondence {
	Correspondence(Feature *a_, Feature *b_) {
	    a = a_;
	    b = b_;
	    distance = a->distance(b);
	}
	float distance;
	Feature *a, *b;

	// Correspondences with lower distances between their features
	// are better, so there is an ordering on correspondences.
	bool operator<(const Correspondence &other) const {
	    return distance < other.distance;
	}
    };
    
    Digest(Window im) {
	// To make a digest we first find some corner features by
	// converting to grayscale ...
	vector<float> grayMatrix;
	for (int i = 0; i < im.channels; i++) {
	    grayMatrix.push_back(1.0f/im.channels);
	}
 	Image gray = ColorMatrix::apply(im, grayMatrix);

	// ... then taking the difference of two Gaussians ...
	Image blurry = GaussianBlur::apply(gray, 0, 3, 3);
	Image blurrier = GaussianBlur::apply(blurry, 0, 5, 5);
	Subtract::apply(blurrier, blurry);

	// ... then finding local maxima.
	vector<LocalMaxima::Maximum> maxima = LocalMaxima::apply(blurrier, false, true, true, 0.00001);
	
	// We sort the maxima to put the strongest ones at the front
	::std::sort(maxima.begin(), maxima.end());

	// We reject maxima too close to the image edges - it makes it hard
	// to extract descriptors
	for (int i = (int)maxima.size()-1, j=0; i >= 0 && j < 256; i--, j++) {
	    if (maxima[i].x < 3 || maxima[i].x > im.width-4 ||
		maxima[i].y < 3 || maxima[i].y > im.height-4) {
		j--;
		continue;
	    }
	    corners.push_back(Feature(maxima[i], im));
	}

	printf("%u corner features found\n", (unsigned int)maxima.size());
    }

    // Once we have computed a digest for each of the images to align,
    // we can attempt to solve for the best alignment using RANSAC and
    // least squares
    Transform *align(Digest &other, Align::Mode m) {
	Transform *transform = NULL, *refined = NULL;
	if (m == Align::TRANSLATE) {
	    transform = new Translation();
	    refined   = new Translation();
	} else if (m == Align::SIMILARITY) {
	    transform = new Similarity();
	    refined   = new Similarity();
	} else if (m == Align::RIGID) {
	    transform = new Rigid();
	    refined   = new Rigid();
	} else if (m == Align::AFFINE) {
	    transform = new Affine();
	    refined   = new Affine();
	} else if (m == Align::PERSPECTIVE) {
	    transform = new Perspective();
	    refined   = new Perspective();
	} else {
	    panic("Unknown transform type: %i\n", m);
	}

	// Associate the features with other features to produce
	// correspondences.
	vector<Correspondence> allCorrespondences, correspondences;

	for (unsigned i = 0; i < corners.size(); i++) {
	    for (unsigned j = 0; j < other.corners.size(); j++) {
		allCorrespondences.push_back(Correspondence(&corners[i], &other.corners[j]));
	    }
	}

	// Sort the correspondences by how good they are. Ones with a
	// low distance between their features will be at the start of
	// this list. This is a little inefficient, given that we're
	// going to throw out most of these, but compared to the image
	// processing steps, everything is cheap.
	::std::sort(allCorrespondences.begin(), allCorrespondences.end());	

	// Select up to 256 of the best correspondences.
	for (unsigned i = 0; i < allCorrespondences.size() && correspondences.size() < 256; i++) {
	    // No feature may be selected more than three times. If
	    // you get a single image patch that matches everything,
	    // it can make a big mess.
	    if (allCorrespondences[i].a->usage < 3 &&
		allCorrespondences[i].b->usage < 3) {
		correspondences.push_back(allCorrespondences[i]);
		allCorrespondences[i].a->usage++;
		allCorrespondences[i].b->usage++;
	    }
	}

	// Print out the correspondences found for debugging.
	for (unsigned i = 0; i < correspondences.size(); i++) {
	    printf("%f %f -> %f %f (%f)\n", 
		   correspondences[i].a->x,
		   correspondences[i].a->y,
		   correspondences[i].b->x,
		   correspondences[i].b->y,
		   correspondences[i].distance);
	}

	// Run RANSAC
	int bestSeed = 0;
	float bestScore = 0;

	for (int iter = 0; iter < 100000; iter++) {
	    // Reset the transform
	    transform->reset();

	    // Choose a random seed
	    int seed = rand();
	    srand(seed);

	    // Pick the minimum number of correspondences required to generate a model
	    for (int i = 0; i < transform->constraintsRequired(); i++) {
		int j = rand() % correspondences.size();
		transform->addCorrespondence(correspondences[j].a->x,
					     correspondences[j].a->y,
					     correspondences[j].b->x,
					     correspondences[j].b->y);
	    }

	    // Do a least squares solve using the minimal number of constraints
	    transform->solve();

	    // Test the remaining correspondences against the model, counting the inliers
	    float score = 0;
	    for (unsigned i = 0; i < correspondences.size(); i++) {
		float x, y;
		transform->apply(correspondences[i].a->x,
				 correspondences[i].a->y,
				 &x, &y);
		x -= correspondences[i].b->x;
		y -= correspondences[i].b->y;

		// When does something count as an inlier? Using this
		// formula, a perfect match is 1, 1 pixel off is 0.5,
		// and it tails off with distance squared.
		score += 1.0/(x*x + y*y + 1);
	    }

	    // See if this is the best model we've found so far (highest number of inliers)
	    if (score > bestScore) {
		bestScore = score;
		bestSeed = seed;
		printf("%i %f\n", bestSeed, bestScore);
	    }
	    
	}

	// Use the best seed we found again to compute its model
	transform->reset();
	srand(bestSeed);
	for (int i = 0; i < transform->constraintsRequired(); i++) {
	    int j = rand() % correspondences.size();
	    printf("Using constraint: %f %f -> %f %f\n",
		   correspondences[j].a->x,
		   correspondences[j].a->y,
		   correspondences[j].b->x,
		   correspondences[j].b->y);
	    transform->addCorrespondence(correspondences[j].a->x,
					 correspondences[j].a->y,
					 correspondences[j].b->x,
					 correspondences[j].b->y);
	}
	transform->solve();	

	// Now we're going to throw in all the inliers under that
	// model into a single big least squares solve to refine the
	// solution.
	refined->reset();
	for (unsigned i = 0; i < correspondences.size(); i++) {
	    float x, y;
	    transform->apply(correspondences[i].a->x,
			     correspondences[i].a->y,
			     &x, &y);
	    x -= correspondences[i].b->x;
	    y -= correspondences[i].b->y;
	    if (x*x + y*y < 1) {
		printf("Inlier: %f %f -> %f %f\n",
		       correspondences[i].a->x,
		       correspondences[i].a->y,
		       correspondences[i].b->x,
		       correspondences[i].b->y);
		refined->addCorrespondence(correspondences[i].a->x,
					   correspondences[i].a->y,
					   correspondences[i].b->x,
					   correspondences[i].b->y);
	    }
	}
	refined->solve();	

	// Done! Return the refined solution.

	delete transform;
	return refined;
    }

    vector<Feature> corners;
};



void Align::help() {
    printf("-align warps the top image on the stack to match the second image on the\n"
	   "stack. align takes one argument, which must be \"translate\", \"similarity\", \n"
	   "\"affine\", \"perspective\", or \"rigid\" and constrains the warp to be of that\n"
	   "type.\n"
	   "\n"
	   "Usage: ImageStack -load a.jpg -load b.jpg -align similarity \\\n"
	   "                  -add -save ab.jpg\n\n");
}

void Align::parse(vector<string> args) {
    assert(args.size() == 1, "-align takes one argument\n");

    Image result;

    if (args[0] == "translate") {
	result = apply(stack(1), stack(0), TRANSLATE);
    } else if (args[0] == "similarity") {
	result = apply(stack(1), stack(0), SIMILARITY);
    } else if (args[0] == "affine") {
	result = apply(stack(1), stack(0), AFFINE);
    } else if (args[0] == "rigid") {
	result = apply(stack(1), stack(0), RIGID);
    } else if (args[0] == "perspective") {
	result = apply(stack(1), stack(0), PERSPECTIVE);
    } else {
	panic("Unknown alignment type: %s. Must be translate, rigid, similarity, affine, or perspective.\n", args[0].c_str());
    }
    pop();
    push(result);
}

// Warp window b to match window a
Image Align::apply(Window a, Window b, Mode m) {

    // First extract features from each frame
    Digest digestA(a);
    Digest digestB(b);

    // now align them
    Transform *transform = digestA.align(digestB, m);

    // now do the warp using lanczos3 sampling
    Image out(a);
    for (int t = 0; t < out.frames; t++) {
	for (int y = 0; y < out.height; y++) {
	    for (int x = 0; x < out.width; x++) {
		float fx, fy;
		transform->apply(x, y, &fx, &fy);
		b.sample2D(t, fx, fy, out(t, x, y));
	    }
	}
    }
    delete transform;
    
    return out;
}