/*
 * Tracker.cpp
 *
 *  Created on: Jun 29, 2014
 *      Author: Daniel Wesierski
 */


#include "tracker.h"
#include "helpers.h"
#include <algorithm>
#include <iterator>

#define INF							1.0e10
#define ZERO						1.0e-10

void Tracker::run() {
	cv::Mat image;
	int frameNum = 1; // frame idx starts from 0
	while (cap.read(image)) {
		double time = start_clock();
		detect(image);
		update(image.cols, image.rows);
		time = stop_clock(time);

		std::cout << "Frame " << frameNum++ << " processed in: " << time << " [ms] \n" << std::endl;
		std::cout << "-----" << std::endl;
		writeframeresult();

		eo.draw(image, 1, true);
		eo.draw(image, 0, false);
		show_img(image, true);		
	}	

	// save trajectories of object parts to file
	cv::Mat res(result.size(), result[0].size(), CV_64F);
	for (int i = 0; i != res.rows; ++i) {
		for (int j = 0; j != res.cols; ++j) {
			res.at<double>(i, j) = result[i][j];
		}
	}
	cvmat2file<double>(res, "result.txt");
}

void Tracker::detect(const cv::Mat& image) {
	gen_candidates();	
	build_graph(image);	
	infer();
}

void Tracker::gen_candidates() {

	for (unsigned int i=0; i != eo.parts.size(); ++i) {

		canEOs[i].resize(eo.grids[i].stepYNum * eo.grids[i].stepXNum * eo.grids[i].bboxes.size());

		unsigned int hyp = 0;

		for (unsigned int j=0; j!=eo.grids[i].stepYNum; ++j) { // generate part locations at different scales on regular 2D grid
			for (unsigned int k=0; k!=eo.grids[i].stepXNum; ++k) {

				unsigned int xCurr = eo.grids[i].x + k*eo.grids[i].stepPlaneX;
				unsigned int yCurr = eo.grids[i].y + j*eo.grids[i].stepPlaneY;

				for (unsigned int l=0; l!=eo.grids[i].bboxes.size(); ++l) {

					cv::Rect bbox(xCurr + eo.grids[i].bboxes[l].x, yCurr + eo.grids[i].bboxes[l].y,
							      eo.grids[i].bboxes[l].width, eo.grids[i].bboxes[l].height);

					canEOs[i][hyp++] = new BBox(bbox);
				}
			}
		}
	}
}

void Tracker::build_graph(const cv::Mat& image) {

	//draw_grid(cv::Mat(image));

	// re-init containers
	unsigned int sz = canEOs.size();
	for (unsigned int i = 0; i != sz; ++i)
		app[i].resize(canEOs[i].size());

	for (unsigned int i = 0; i != sz - 1; ++i)
		def[i].resize(canEOs[i].size()*canEOs[i + 1].size());

	//-- 1. match part appearance
	for (unsigned int i = 0; i != sz; ++i)
		for (unsigned int j = 0; j != canEOs[i].size(); ++j)
			app[i][j] = eo.appVar*chisquare_score(eo.parts[i].hist,
			eo.parts[i].compute_hist(image, canEOs[i][j]->bbox, eo.nBins));

	//-- 2. match neighbor parts shape
	double angStdDevInv = deg() / eo.angVar;

	for (unsigned int i = 0; i != sz - 1; ++i) {

		// precompute rescaling factors
		std::vector<double> rescaleInv(eo.grids[i].scaleF.size()*eo.grids[i + 1].scaleF.size());
		for (unsigned int j = 0; j != eo.grids[i].scaleF.size(); ++j)
			for (unsigned int k = 0; k != eo.grids[i + 1].scaleF.size(); ++k)
				rescaleInv[k + j*eo.grids[i + 1].scaleF.size()] = (1.0 / (0.5*(eo.grids[i].scaleF[j] + eo.grids[i + 1].scaleF[k])));

		double distMean = eo.kinMean[i];
		double angMean = eo.kinAngOff[i];
		double distStdDevInv = 1.0 / eo.kinStdDev[i];

		// single global orientation rotated by predefined angular offset
		double cAng = std::cos(angMean);
		double sAng = std::sin(angMean);
		double localOrientationX = cAng*eo.segCS[i].x - sAng*eo.segCS[i].y;
		double localOrientationY = sAng*eo.segCS[i].x + cAng*eo.segCS[i].y;

		// get def weights
		int canEOsNextNum = canEOs[i + 1].size();

		// for each candidate of i-th part...
		for (unsigned int j = 0; j != canEOs[i].size(); j += eo.grids[i].scaleF.size())
		{
			// ...scan through all candidates of i+1-th part
			for (unsigned int k = 0; k != canEOs[i + 1].size(); k += eo.grids[i + 1].scaleF.size())
			{
				cv::Point2d orient = canEOs[i + 1][k]->center - canEOs[i][j]->center;

				double distance12 = std::sqrt(orient.x*orient.x + orient.y*orient.y);
				double angle = 0;

				if (distance12 < ZERO)  // two parts coincide
				{
					angle = angMean;
				}
				else
				{
					orient.x /= distance12;   // local orientation between candidates
					orient.y /= distance12;

					double nr = std::min(1.0, orient.x*localOrientationX + orient.y*localOrientationY);
					angle = std::acos(nr);
				}

				// add angular mean between pairs of candidates
				double angularScore = angle*angStdDevInv;
				angularScore = angularScore*angularScore;

				int s = 0;
				for (unsigned int j1 = j; j1 != j + eo.grids[i].scaleF.size(); ++j1)
				{
					for (unsigned int k1 = k; k1 != k + eo.grids[i + 1].scaleF.size(); ++k1)
					{
						// rescaled distances
						double distScore = (rescaleInv[s++] * distance12 - distMean)*distStdDevInv;
						distScore = distScore*distScore;
						def[i][j1*canEOsNextNum + k1] = distScore + angularScore;
					}
				}
			}
		}
	}
}

void Tracker::infer() {
	pass_msg();
	backtrack();
}

void Tracker::pass_msg() {
	for (unsigned int i=app.size()-1; i!=0; --i) {

		ptr[i-1].resize(app[i-1].size());

		for (unsigned int j=0; j!=app[i-1].size(); ++j) {
			std::vector<double> score0(app[i].size());
			std::transform(app[i].begin(), app[i].end(), def[i-1].begin()+j*app[i].size(), score0.begin(), std::plus<double>());
			std::vector<double>::iterator bestCan = std::min_element(score0.begin(), score0.end());
			app[i-1][j] += *bestCan;
			ptr[i-1][j] = std::ptrdiff_t(bestCan-score0.begin());
		}
	}
}

void Tracker::backtrack() {
	bestCansIdx[0] = std::ptrdiff_t(std::min_element(app[0].begin(), app[0].end())-app[0].begin());
	bestEO[0] = BBox(canEOs[0][bestCansIdx[0]]->bbox);

	for(unsigned int i=0; i!=ptr.size(); ++i) {
		bestCansIdx[i+1] = ptr[i][bestCansIdx[i]];
		bestEO[i+1] = BBox(canEOs[i+1][bestCansIdx[i+1]]->bbox);
	}
}

void Tracker::update(int imW, int imH) {

	update_scale();
	update_orientation();
	update_location(imW, imH);
	clear();
}

void Tracker::update_scale() {

	double scaleNew=0;
	for (unsigned int i=0; i!=bestEO.size(); ++i)
		scaleNew += double(bestEO[i].bbox.width)/eo.parts[i].bbox.width;

	scaleNew /= bestEO.size();
	eo.scale = (1.0-coeffIIR)*eo.scale + coeffIIR*scaleNew;
}

void Tracker::update_orientation() {

	for (unsigned int i=0; i!=eo.partSegments.size(); ++i) {

		cv::Mat prevLoc(2, eo.partSegments[i].size(), CV_64F);
		cv::Mat currLoc(2, eo.partSegments[i].size(), CV_64F);

		for (unsigned int j=0; j!=eo.partSegments[i].size(); ++j) {		
			prevLoc.at<double>(0,j) = eo.scale*eo.parts[eo.partSegments[i][j]].center.x;
			prevLoc.at<double>(1,j) = eo.scale*eo.parts[eo.partSegments[i][j]].center.y;

			currLoc.at<double>(0,j) = bestEO[eo.partSegments[i][j]].center.x;
			currLoc.at<double>(1,j) = bestEO[eo.partSegments[i][j]].center.y;
		}

		cv::Mat rot = Kabsch(prevLoc,currLoc);		
		cv::Mat currOrient(2,1,CV_64F);
		currOrient.at<double>(0,0) = rot.at<double>(0,0)*eo.segCS[eo.partSegments[i][0]].x + rot.at<double>(0,1)*eo.segCS[eo.partSegments[i][0]].y;
		currOrient.at<double>(1,0) = rot.at<double>(1,0)*eo.segCS[eo.partSegments[i][0]].x + rot.at<double>(1,1)*eo.segCS[eo.partSegments[i][0]].y;

		if (std::abs(1-std::sqrt(currOrient.at<double>(0,0)*currOrient.at<double>(0,0)+currOrient.at<double>(1,0)*currOrient.at<double>(1,0))) > 0.0001)
			std::cout << "Error" << std::endl;

		for (unsigned int j=0; j!=eo.partSegments[i].size()-1; ++j)
			eo.segCS[eo.partSegments[i][j]] = cv::Point2d(currOrient);
	}
}

void Tracker::update_location(int imW, int imH) {

	// shape params
	for (unsigned int i=0; i!=eo.kinMean.size() ; ++i) {
		eo.kinMean[i] *= eo.scale;
		eo.kinStdDev[i] *= eo.scale;
	}

	// bboxes
	for (unsigned int i=0; i!=eo.parts.size() ; ++i) {		

		eo.parts[i].bbox.width *= eo.scale;
		eo.parts[i].bbox.height *= eo.scale;

		eo.parts[i].bbox.x = std::min(double(imW) - eo.parts[i].bbox.width, 
			                 std::max(0.0, bestEO[i].center.x - eo.parts[i].bbox.width / 2));
		eo.parts[i].bbox.y = std::min(double(imH) - eo.parts[i].bbox.height, 
			                 std::max(0.0, bestEO[i].center.y - eo.parts[i].bbox.height / 2));
				
		eo.parts[i].center = eo.parts[i].get_center();
	}

	// grid
	for(unsigned int i=0; i!=eo.grids.size(); ++i)
		eo.grids[i].update_grid(eo.parts[i].bbox);
}

void Tracker::writeframeresult() {
	std::vector<double> frameresult;
	for (unsigned int i = 0; i != eo.parts.size(); ++i) {
		frameresult.push_back(eo.parts[i].bbox.x);
		frameresult.push_back(eo.parts[i].bbox.y);
		frameresult.push_back(eo.parts[i].bbox.width);
		frameresult.push_back(eo.parts[i].bbox.height);
	}
	result.push_back(frameresult);
}

void Tracker::clear() {
	for (unsigned int i=0; i!=canEOs.size(); ++i)
		for (unsigned int j=0; j!=canEOs[i].size(); ++j)
			delete canEOs[i][j];
}

double Tracker::bhattacharyya_score(const cv::Mat& hist1, const cv::Mat& hist2) {
	// it is assumed that both column vectors have equal dim
	double bhattacharyyaScore = 0;
	for (int i=0; i != hist1.rows; ++i)
		bhattacharyyaScore += std::sqrt(hist1.at<float>(i,0)*hist2.at<float>(i,0));

	return bhattacharyyaScore;
}

double Tracker::chisquare_score(const cv::Mat& hist1, const cv::Mat& hist2) {
	float chiSquare = 0;
	float* hist1Ptr = (float*) hist1.data;
	float* hist2Ptr = (float*) hist2.data;

	for (int i=0; i != hist1.rows; ++i) {
		if (hist1Ptr[i] > ZERO || hist2Ptr[i] > ZERO) {
			float val = hist1Ptr[i]-hist2Ptr[i];
			chiSquare += (val*val)/(hist1Ptr[i]+hist2Ptr[i]);
		}
	}

	return double(0.5*chiSquare);
}

cv::Mat Tracker::Kabsch(cv::Mat& p, cv::Mat& q) {
    //-- number of points
    int N = p.cols;

    //-- compute centroids
    cv::Mat cp(2,1,CV_64F,cv::Scalar(0)), cq(2,1,CV_64F,cv::Scalar(0));
    for (int i=0; i!=p.rows; ++i)
        for (int j=0; j!=p.cols; ++j)
            cp.at<double>(i,0) += p.at<double>(i,j);

    cp /= N;

    for (int i=0; i!=q.rows; ++i)
        for (int j=0; j!=q.cols; ++j)
            cq.at<double>(i,0) += q.at<double>(i,j);

    cq /= N;

    //-- translate to center of origin
    for (int i=0; i!=p.rows; ++i)
        for (int j=0; j!=p.cols; ++j)
            p.at<double>(i,j) -= cp.at<double>(i,0);

    for (int i=0; i!=q.rows; ++i)
        for (int j=0; j!=q.cols; ++j)
            q.at<double>(i,j) -= cq.at<double>(i,0);

    //-- covariance matrix of the coordinates
    cv::Mat c = p*q.t()/N;
    cv::Mat w, u, vt;
    cv::SVD::compute(c, w, u, vt);

    cv::Mat id = cv::Mat::eye(2, 2, CV_64F);

    if (cv::determinant(u*vt)<0)
    	id.at<double>(1,1) = -1;

	return vt.t()*id*u.t();
}

void Tracker::draw_grid(cv::Mat image, unsigned int d)
{
	for (unsigned int i=0; i != canEOs.size(); ++i) {

		for (unsigned int j=0; j != eo.grids[i].scaleNum; ++j)
			canEOs[i][j]->draw(image,color(j+3),1);

		for (unsigned int j=0; j != canEOs[i].size(); ++j) {
			//cv::Point c = canEOs[i][j]->get_center();
			cv::Point c = canEOs[i][j]->bbox.tl();
			cv::rectangle(image, cv::Point(c-cv::Point(d,d)), cv::Point(c+cv::Point(d,d)), color(), 1, 8);
		}
		eo.parts[i].draw(image,color(6),1);
	}
}
