#include "StereoVision.h"
#include <cmath>


StereoVision::StereoVision(void)
{
}

StereoVision::~StereoVision(void)
{
}
std::vector<candidate*> StereoVision::findCorrespPointsPMF(BV::ImageRGB &leftScanline, BV::ImageRGB &rightScanline)
{
	std::vector<unsigned> leftEdges;
	std::vector<unsigned> rightEdges;
	StereoVision::findEdges(leftScanline,rightScanline,leftEdges,rightEdges);
	
	std::list<candidatePair> candidatePairs;
	std::vector<candidate*> candidates;
	std::list<candidatePair>::iterator i;
	double dgThreshold = 1.0;

	// generate all candidates of corresponding points (only x coordinates are relevant)
	if(leftEdges.size() == 0 || rightEdges.size() == 0) {
		return candidates;
	}

	for(size_t lc = 0;lc<leftEdges.size();lc++) {
		for(size_t rc = 0;rc<rightEdges.size();rc++) {
			BV::RGBPixel leftImgPx = leftScanline.getPixel(leftEdges[lc],0);
			BV::RGBPixel rightImgPx = rightScanline.getPixel(rightEdges[rc],0);
			BV::RGBPixel leftImgPxBefore = leftScanline.getPixel(leftEdges[lc]-1,0);
			BV::RGBPixel rightImgPxBefore = rightScanline.getPixel(rightEdges[rc]-1,0);
			BV::RGBPixel leftImgPxAfter = leftScanline.getPixel(leftEdges[lc]+1,0);
			BV::RGBPixel rightImgPxAfter = rightScanline.getPixel(rightEdges[rc]+1,0);

			if (arePixelSimilar(leftImgPx,rightImgPx) 
				&& arePixelSimilar(leftImgPxBefore,rightImgPxBefore) && arePixelSimilar(leftImgPxAfter,rightImgPxAfter) )  {
				candidate* c = new candidate();
				c->xLeft=leftEdges[lc];
				c->xRight=rightEdges[rc];
				c->propapilityCounter = 0.0;
				c->isAssigned = false;
				candidates.push_back(c);
			}
		}
	}
	size_t numCandidates = candidates.size();
	if(numCandidates == 1) {
		return candidates;
	}
	//long pairCounter = 0;
	// generate candidate Pairs (each with each)
	for(std::vector<candidate*>::iterator n=candidates.begin();n!=candidates.end();n++) {
		for(std::vector<candidate*>::iterator m=n;m!=candidates.end();m++) {
			double dispDist = std::abs(getDisparityDifference((*n)->xLeft,(*n)->xRight,(*m)->xLeft,(*m)->xRight));
			// check that candidates are not to far away from each other
			if( dispDist < DISP_DIST_THRESHOLD) {
				double dg = getDisparityGradient((*n)->xLeft,(*n)->xRight,(*m)->xLeft,(*m)->xRight);
				if( (dg != ERROR_DISP_GRADIENT) && (dg <= MAX_DISP_GRADIENT) ) {
					candidatePair cp;
					cp.firstCand = *n;
					cp.secondCand = *m;
					cp.disparityGradient = dg;
					cp.isAssigned = false;
					cp.isDeleted = false;
					cp.candidateDistance = dispDist;
					candidatePairs.push_back(cp);
				}	
			}
		}
	}
	size_t numCandidatePairs = candidatePairs.size();
	bool allAssigned = false;
	long loopCounter = 0;
	do {
		// for each candidate set propapilityCounter to 0
		for(std::vector<candidate*>::iterator n=candidates.begin();n!=candidates.end();n++) {
			if(!(*n)->isAssigned) {
				(*n)->propapilityCounter = 0.0;
			}
		}
		// set propapility counters for all unassigned pairs
		for(i = candidatePairs.begin(); i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				if(i->disparityGradient <= dgThreshold) {
					// weight with reciprocal of distance
					i->firstCand->propapilityCounter+=1/(i->candidateDistance);
					i->secondCand->propapilityCounter+=1/(i->candidateDistance);
				}
			}
		}
		std::vector<candidatePair*> winnerPairs;
		double maxPropability = 0;
		// find first pair, which is not assigned and set this as max propability
		for(i = candidatePairs.begin(); i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				winnerPairs.push_back(&(*i));
				maxPropability = i->firstCand->propapilityCounter + i->secondCand->propapilityCounter;
				break;
			}
		}
		// find the pair with greatest sum of both propability counters
		for(i = candidatePairs.begin(); i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				double propapility = i->firstCand->propapilityCounter + i->secondCand->propapilityCounter;
				if (maxPropability < propapility )
				{
					winnerPairs.clear();
					winnerPairs.push_back(&(*i));
					maxPropability = propapility;
				} else if(maxPropability == propapility) {
					winnerPairs.push_back(&(*i));
				}
			}
		}
		// set each winner pair to state assigned
		size_t numWinPairs = winnerPairs.size();
		for(std::vector<candidatePair*>::iterator w = winnerPairs.begin();w!=winnerPairs.end();w++) {
			(*w)->isAssigned = true;
			(*w)->firstCand->isAssigned = true;
			(*w)->secondCand->isAssigned = true;
		}
		// delete all candidate pairs, which contain one the candidates of the "winner" candidate pairs
		for(i=candidatePairs.begin();i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				for(std::vector<candidatePair*>::iterator w = winnerPairs.begin();w!=winnerPairs.end();w++) {
					if(! ( (*i) == *(*w) ) ) {
						if( (*w)->hasCandidate(i->firstCand) || (*w)->hasCandidate(i->secondCand) ) {
							i->isDeleted = true;
						}
					}
				}
			}
		}
		// check if one more iteration is needed
		int numOfUnassignedPairs = 0;
		for(i=candidatePairs.begin();i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				numOfUnassignedPairs++;
			}
		}
		if(numOfUnassignedPairs == 0) {
			allAssigned = true;
		}
	} while(!allAssigned);

	/* for debugging only
	for(i=candidatePairs.begin();i!=candidatePairs.end();i++) {
		if(i->isAssigned && (!i->firstCand->isAssigned || !i->secondCand->isAssigned)) {
			int x = 7;
		}
	}
	*/
	std::vector<candidate*> res;
	/*
	for(std::vector<candidate*>::iterator ci = candidates.begin();ci!=candidates.end();ci++) {
		if((*ci)->isAssigned) {
			res.push_back(*ci);
		}
	}*/
	for(i=candidatePairs.begin();i!=candidatePairs.end();i++) {
		if(i->isAssigned && !i->isDeleted) {
			bool alreadyInResult = false;
			for(std::vector<candidate*>::iterator ri = res.begin();ri!=res.end();ri++) {
				if( (*ri) == i->firstCand ) {
					alreadyInResult = true;
					break;
				}
			}
			if(!alreadyInResult) {
				res.push_back(i->firstCand);
			}
			alreadyInResult = false;
			for(std::vector<candidate*>::iterator ri = res.begin();ri!=res.end();ri++) {
				if( (*ri) == i->secondCand ) {
					alreadyInResult = true;
					break;
				}
			}
			if(!alreadyInResult) {
				res.push_back(i->secondCand);
			}
		}
	}
	return res;
}

std::vector<candidate*> StereoVision::findCorrespPointsPMF(BV::Image &leftScanline, BV::Image &rightScanline)
{
	std::vector<unsigned> leftEdges;
	std::vector<unsigned> rightEdges;
	StereoVision::findEdges(leftScanline,rightScanline,leftEdges,rightEdges);
	
	std::list<candidatePair> candidatePairs;
	std::vector<candidate*> candidates;
	std::list<candidatePair>::iterator i;
	double dgThreshold = 1.0;

	// generate all candidate points (only x coordinates are relevant)
	for(size_t lc = 0;lc<leftEdges.size();lc++) {
		for(size_t rc = 0;rc<rightEdges.size();rc++) {
			if(leftEdges[lc] == 0 || rightEdges[rc] == 0 || leftEdges[lc] == 255 || rightEdges[rc] == 255) {
				continue;
			}
			int leftVal = leftScanline.getPixel(leftEdges[lc],0);
			int rightVal = rightScanline.getPixel(rightEdges[rc],0);
			if ( leftVal == rightVal )
			{
				candidate* c = new candidate();
				c->xLeft=leftEdges[lc];
				c->xRight=rightEdges[rc];
				c->propapilityCounter = 0.0;
				c->isAssigned = false;
				candidates.push_back(c);
			}
		}
	}
	size_t numCandidates = candidates.size();
	
	//long pairCounter = 0;
	// generate candidate Pairs (each with each)
	for(std::vector<candidate*>::iterator n=candidates.begin();n!=candidates.end();n++) {
		for(std::vector<candidate*>::iterator m=n;m!=candidates.end();m++) {
			double dispDist = std::abs(getCyclopDistanceEpipolar((*n)->xLeft,(*n)->xRight,(*m)->xLeft,(*m)->xRight));
			// check that candidates are not to far away from each other
			if( dispDist < DISP_DIST_THRESHOLD) {
				double dg = getDisparityGradient((*n)->xLeft,(*n)->xRight,(*m)->xLeft,(*m)->xRight);
				if( (dg != ERROR_DISP_GRADIENT) && (dg < MAX_DISP_GRADIENT) ) {
					candidatePair cp;
					cp.firstCand = *n;
					cp.secondCand = *m;
					cp.disparityGradient = dg;
					cp.isAssigned = false;
					cp.isDeleted = false;
					cp.candidateDistance = dispDist;
					candidatePairs.push_back(cp);
				} else {
					continue;
				}		
			}
		}
	}
	size_t numCandidatePairs = candidatePairs.size();
	bool allAssigned = false;
	long loopCounter = 0;
	do {
		// for each candidate set propapilityCounter to 0
		for(std::vector<candidate*>::iterator n=candidates.begin();n!=candidates.end();n++) {
			if(!(*n)->isAssigned) {
				(*n)->propapilityCounter = 0.0;
			}
		}
		// set propapility counters for all unassigned pairs
		for(i = candidatePairs.begin(); i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				if(i->disparityGradient < dgThreshold) {
					// weight with reciprocal of distance
					i->firstCand->propapilityCounter+=1.0;
					i->secondCand->propapilityCounter+=1.0;
				}
				break;
			}
		}
		std::vector<candidatePair*> winnerPairs;
		double maxPropability = 0;
		// find first pair, which is not assigned and set this as max propability
		for(i = candidatePairs.begin(); i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				winnerPairs.push_back(&(*i));
				// weight with reciprocal of distance
				maxPropability = (i->firstCand->propapilityCounter + i->secondCand->propapilityCounter) / (100*i->candidateDistance);
				break;
			}
		}
		// find the pair with greatest sum of both propability counters
		for(i = candidatePairs.begin(); i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				// weight with reciprocal of distance
				double propapility = (i->firstCand->propapilityCounter + i->secondCand->propapilityCounter) / (100*i->candidateDistance);
				if (maxPropability < propapility )
				{
					winnerPairs.clear();
					winnerPairs.push_back(&(*i));
					maxPropability = propapility;
				} else if(maxPropability == propapility) {
					winnerPairs.push_back(&(*i));
				}
			}
		}
		// set each winner pair to state assigned
		size_t numWinPairs = winnerPairs.size();
		for(std::vector<candidatePair*>::iterator w = winnerPairs.begin();w!=winnerPairs.end();w++) {
			(*w)->isAssigned = true;
			(*w)->firstCand->isAssigned = true;
			(*w)->secondCand->isAssigned = true;
		}
		// delete all candidate pairs, which contain one the candidates of the "winner" candidate pairs
		for(i=candidatePairs.begin();i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				for(std::vector<candidatePair*>::iterator w = winnerPairs.begin();w!=winnerPairs.end();w++) {
					if(! ( (*i) == *(*w) ) ) {
						if( (*w)->hasCandidate(i->firstCand) || (*w)->hasCandidate(i->secondCand) ) {
							i->isDeleted = true;
						}
					}
				}
			}
		}
		// check if one more iteration is needed
		int numOfUnassignedPairs = 0;
		for(i=candidatePairs.begin();i!=candidatePairs.end();i++) {
			if(!i->isAssigned && !i->isDeleted) {
				numOfUnassignedPairs++;
			}
		}
		if(numOfUnassignedPairs == 0) {
			allAssigned = true;
		}
	} while(!allAssigned);

	/* for debugging only
	for(i=candidatePairs.begin();i!=candidatePairs.end();i++) {
		if(i->isAssigned && (!i->firstCand->isAssigned || !i->secondCand->isAssigned)) {
			int x = 7;
		}
	}
	*/
	std::vector<candidate*> res;

	for(std::vector<candidate*>::iterator ci = candidates.begin();ci!=candidates.end();ci++) {
		if((*ci)->isAssigned) {
			res.push_back(*ci);
		}
	}
	return res;
}

void StereoVision::findEdges(BV::ImageRGB &scanlineLeft, BV::ImageRGB &scanlineRight, std::vector<unsigned> &leftEdges, std::vector<unsigned>&rightEdges)
{
	unsigned width = scanlineLeft.getWidth();;
	unsigned height = 0;
	BV::RGBPixel prevLeftPixel = scanlineLeft.getPixel(0,height);
	BV::RGBPixel prevRightPixel = scanlineRight.getPixel(0,height);
	BV::RGBPixel actPixel;
	for(unsigned i=1;i<width-1;i++) {
		actPixel = scanlineLeft.getPixel(i,height);
		if( !arePixelSimilar(prevLeftPixel,actPixel) ) 
		{
			leftEdges.push_back(i);
		}
		prevLeftPixel = actPixel;

		actPixel = scanlineRight.getPixel(i,height);
		if( !arePixelSimilar(prevRightPixel,actPixel) ) 
		{
			rightEdges.push_back(i);
		}
		prevRightPixel = actPixel;
	}
}

void StereoVision::findEdges(BV::Image &scanlineLeft, BV::Image &scanlineRight, std::vector<unsigned> &leftEdges, std::vector<unsigned>&rightEdges)
{
	unsigned width = scanlineLeft.getWidth();;
	unsigned height = 0;
	int prevLeftPixel = scanlineLeft.getPixel(0,height);
	int prevRightPixel = scanlineRight.getPixel(0,height);
	int actPixel;
	int sky = 0;
	for(unsigned i=1;i<width;i++) {
		actPixel = scanlineLeft.getPixel(i,height);
		if( prevLeftPixel != actPixel ) 
		{
			if(prevLeftPixel == sky) {
				leftEdges.push_back(i);
			} else {
				leftEdges.push_back(i-1);
			}
		}
		prevLeftPixel = actPixel;

		actPixel = scanlineRight.getPixel(i,height);
		if( prevRightPixel != actPixel) 
		{
			if(prevRightPixel == sky) {
				rightEdges.push_back(i);
			} else {
				rightEdges.push_back(i-1);
			}
		}
		prevRightPixel = actPixel;
	}
}


double StereoVision::getCyclopDistanceEpipolar(const unsigned int &xAl, const unsigned int &xAr,const unsigned&xBl,const unsigned&xBr)
{
	long leftDisparity = xAl - xBl;
	long rightDisparity = xAr - xBr;
	return 0.5*(leftDisparity+rightDisparity);
}
long StereoVision::getDisparityDifference(const unsigned int &xAl, const unsigned int &xAr,const unsigned&xBl,const unsigned&xBr)
{
	long leftDisparity = xAl - xAr;
	long rightDisparity = xBl - xBr;
	return std::abs(rightDisparity-leftDisparity);
}
double StereoVision::getDisparityGradient(const unsigned int &xAl, const unsigned int &xAr,const unsigned&xBl,const unsigned&xBr)
{
	double cyclopDist = getCyclopDistanceEpipolar(xAl,xAr,xBl,xBr);
	if(cyclopDist != 0.0) {
		return (getDisparityDifference(xAl,xAr,xBl,xBr)/ std::abs(cyclopDist) );
	} else {
		return ERROR_DISP_GRADIENT;
	}
	
}
bool StereoVision::arePixelSimilar(const BV::RGBPixel& left,const BV::RGBPixel& right) {
	bool res;
	res = (MAX_COLOR_DIFF >= std::abs(left.getRed() - right.getRed())) && (MAX_COLOR_DIFF >= std::abs(left.getGreen() - right.getGreen())) && (MAX_COLOR_DIFF >= std::abs(left.getBlue() - right.getBlue()));
	return res;
}