#include "observer.h"
#include "util.h"

// Draws the normals from each of the point
void Observer::drawNormals(IplImage* image, Spline *sp)
{
	if (sp->points.size() < 4) 
		return;

	for(int i=0; i<sp->points.size();i++) {
		vec3d n = sp->getNormalPoint(i);
		cvLine( image,  VEC2POINT(sp->points[i]),
				VEC2POINT(sp->points[i]+n*normalLength), cvScalar(0xff00ff), 1);
		cvLine( image,  VEC2POINT(sp->points[i]),
				VEC2POINT(sp->points[i]-n*normalLength), cvScalar(0x00ff00), 1);
	}
}

// Search along a line to find the nearest hit.  Parameters: image, obvious.
// pt1 is where to start search. pt2 is where to end it. returns the vector of
// the first edge, or 1 in the z coordinate if no edge is found.
vec3d Observer::findNearestEdge(IplImage* image, CvPoint pt1, CvPoint pt2)
{
	// Need to bound pt1 and pt2 inside the image.
	if (pt1.x >= image->width || pt1.x < 0 ||
			pt1.y >= image->height || pt1.y < 0) {
		return vec3d(0,0,1);
	}
	if (pt2.x >= image->width || pt2.x < 0 ||
			pt2.y >= image->height || pt2.y < 0) {
		return vec3d(0,0,1);
	}


	// Image should be 1-channel
	CvLineIterator iterator;
	int count = cvInitLineIterator( image, pt1, pt2, &iterator, 8 );

	for( int i = 0; i < count; i++ ){
		unsigned int in = (unsigned int)iterator.ptr[0];
		if (in == 255) {
			// We found an edge. Note that this is in absolute
			// image coords
			int offset, x, y;
			offset = iterator.ptr - (uchar*)(image->imageData);
			y = offset/image->widthStep;
			x = (offset - y*image->widthStep)/sizeof(uchar);
			return vec3d(x,y,0);
		}
		CV_NEXT_LINE_POINT(iterator);
	}
	return vec3d(0,0,1);
}

// Adjusts each control point to the nearest edge in the image.
void Observer::fullyTweakSpline(Spline& d)
{
	if (d.points.size() < 4) 
		return;

	vec3d v1, v2;

	for(int i=0; i<d.points.size();i++) {
		vec3d n = d.getNormalPoint(i);
		v1 = d.points[i]-n*normalLength;
		v2 = d.points[i]+n*normalLength;

		vec3d e1 = findNearestEdge(lastImage, VEC2POINT(d.points[i]), VEC2POINT(v1));
		vec3d e2 = findNearestEdge(lastImage, VEC2POINT(d.points[i]), VEC2POINT(v2));

		if (e1[2] != 0 && e2[2] != 0) {
			d.points[i] = vec3d(0,0,2*normalLength);
		} else if (e1[2] !=0) {
			d.points[i] = e2;
		} else if (e2[2] !=0) {
			d.points[i] = e1;
//		} else if (3*(e1-d.points[i]).mag() < (e1-d.points[i]).mag()) {
		} else if ((e1-d.points[i]).mag() < (e1-d.points[i]).mag()) {
			d.points[i] = e1;
		} else {
			d.points[i] = e2;
		}
	}

}

void Observer::fullyTweakSpline(IplImage* image)
{
	if (d == NULL) {
		return;
	}

	// Only valid when we have > 4 coords in the spline
	if (d->points.size() < 4) 
		return;

	vec3d v1, v2;

	// We want to take curve points that didn't hit anything, and move them
	// by the average displacement of the other points.
	vector<int> missinds;
	vec3d dispsum;
	int hits;

	for(int i=0; i<d->points.size();i++) {
		vec3d n = d->getNormalPoint(i);
		v1 = d->points[i]-n*normalLength;
		v2 = d->points[i]+n*normalLength;

		// Find nearest edge!
		vec3d e1 = findNearestEdge(image, VEC2POINT(d->points[i]), VEC2POINT(v1));
		vec3d e2 = findNearestEdge(image, VEC2POINT(d->points[i]), VEC2POINT(v2));

		if (e1[2] != 0 && e2[2] != 0) {
			missinds.push_back(i);
			hits += 1;
		} else if (e1[2] !=0) {
			dispsum += (e2-d->points[i]);
			d->points[i] = e2;
		} else if (e2[2] !=0) {
			dispsum += (e1-d->points[i]);
			d->points[i] = e1;
		// Note the factor of 3 for e1. This helps when the curve is
		// wound clockwise; the image features tend to be bunched
		// inside the object which leads to spline collapse. So we
		// penalize the inner normal.
//		} else if (3*(e1-d->points[i]).mag() < (e1-d->points[i]).mag()) {
		} else if ((e1-d->points[i]).mag() < (e1-d->points[i]).mag()) {
			dispsum += (e1-d->points[i]);
			d->points[i] = e1;
		} else {
			dispsum += (e2-d->points[i]);
			d->points[i] = e2;
		}
	}

	// Get real hits and move control points that didn't hit anything over
	// by the average displacement
	hits = d->points.size()-hits;
	dispsum /= (double) hits;
	for (int i=0; i<missinds.size(); i++) {
		d->points[missinds[i]] += dispsum;
	}
}

void Observer::draw(IplImage *image)
{
	if (d != NULL) {
		d->drawSpline(image);
		drawNormals(image, d);
	}
}

double Observer::fitTransform( ShapeState& x )
{
	st.setDerived(x);
	Spline t = st.d;
	fullyTweakSpline(t);
	return st.mse(t);
}

void Observer::startObserving()
{
	cout << "Starting to observe\n";
	if (observing) {
		cout << "Already observing! Error!\n";
		return;
	}
	observing = true;
	centroid = st.s.findCentroid();
	last[0] = centroid[0]; // x
	last[1] = centroid[1]; // y
	last[2] = 0; // Scale / rotation
	last[3] = 0;   
	st.setTemplate();
	d = &st.d;
	st.setDerived(last);

	Nb =  st.s.points.size();

	// Set up the W matrix for translations
	W = cvCreateMat(Nb*2, 4, CV_32FC1);
	for (int i = 0; i< Nb; i++) {
		// 1 0 Qx -Qy
		cvmSet(W, i, 0, 1);
		cvmSet(W, i, 1, 0);
		cvmSet(W, i, 2,  st.s.points[i].x());
		cvmSet(W, i, 3, -st.s.points[i].y());
		// 0 1 Qy Qx
		cvmSet(W, 2*i, 0, 0);
		cvmSet(W, 2*i, 1, 1);
		cvmSet(W, 2*i, 2, st.s.points[i].y());
		cvmSet(W, 2*i, 3, st.s.points[i].x());
	}

	// Set up the W matrix for translations
	Q0 = cvCreateMat(Nb*2, 1, CV_32FC1);
	Qf = cvCreateMat(Nb*2, 1, CV_32FC1);
	Qt = cvCreateMat(Nb*2, 1, CV_32FC1);
	for (int i = 0; i< Nb; i++) {
		cvmSet(Q0,   i, 0, st.s.points[i].x());
		cvmSet(Q0, 2*i, 0, st.s.points[i].y());
	}

	// If we are nasty and set U = I, then projection becomes vastly
	// simpler. Here we solve for the pseudo inverse, assuming U = I (i.e.
	// no spline at all, just control points)
	Wp = cvCreateMat(4, Nb*2, CV_32FC1);
	cvInvert(W, Wp, CV_SVD);
}

ShapeState Observer::observe(IplImage *image, ShapeState estimate)
{
	// Here we must make our estimate, then try to fit a curve and get the
	// ShapeState for the deformed curve.
	
	lastImage = image;

	if (d == NULL) 
		return ShapeState();
	
	// Now make the estimate from the particle filter the spline we're
	// going to deform (rather than our last observation)
	st.setDerived(estimate);

	// Deform the spline
	fullyTweakSpline(image);

	// Now project it into the shape space of Euclidean similarities
	for (int i = 0; i< Nb; i++) {
		cvmSet(Qf,   i, 0, d->points[i].x());
		cvmSet(Qf, 2*i, 0, d->points[i].y());
	}
	
	// X = Wp ( Qf - Q0 )
	cvSub(Qf,Q0,Qt);
	CvMat *X = cvCreateMat(4, 1, CV_32FC1);
	cvMatMulAdd(Wp, Qt, NULL, X);
	for (int i=0; i < 4; i++) {
		last[i] = cvmGet(X,i,0);
	}

	// Now update the derived spline
	st.setDerived(last);
	return last;
}

void Observer::addPoint(int x, int y)
{
	if (!observing) {
		st.addPoint(x,y);
	}
}