#include <iostream>
#include "CC_Alloc.h"


#include "PointMatcher.h"
#include "math.h"

using namespace std;

#define  SQR(a)             ((a)*(a))
#define  ABS(a)             ((a) < 0.0 ? -(a) : (a))

#define MEAN_THRESHOLD_FOR_CLOSE_FRAMES 5


#define MEMALLOC(ptr, type, number) \
	if ((ptr = (type *)CC_malloc(sizeof(type) * (number))) == NULL)\
{ \
	fprintf(stderr, "> Can't alloc memory for variable " #ptr "\n"); \
	fprintf(stderr, "> (Asked for %d instances of `" #type "').\n", (number));\
}

FptPointMatcher::~FptPointMatcher()
{
	freeBuckets(&buckets1, &bucketsSizes1);
	freeBuckets(&buckets2, &bucketsSizes2);
}

void FptPointMatcher::allocBuckets(FptFeaturePoint_ ***** buckets, int *** bucketsSizes, 
								   int maxPointNumberPerBucket)
{
	int i, j;

	freeBuckets(buckets, bucketsSizes);

	MEMALLOC(*buckets, FptFeaturePoint_***, bucketsWidth);
	MEMALLOC(*bucketsSizes, int*, bucketsWidth);
	for(i = 0; i < bucketsWidth; i++)
	{
		MEMALLOC((*buckets)[i], FptFeaturePoint_**, bucketsHeight);
		MEMALLOC((*bucketsSizes)[i], int, bucketsHeight);
		for(j = 0; j < bucketsHeight; j++)
		{
			MEMALLOC((*buckets)[i][j], FptFeaturePoint_*, maxPointNumberPerBucket);
			(*bucketsSizes)[i][j] = 0;
		}
	}
}

void FptPointMatcher::freeBuckets(FptFeaturePoint_ ***** buckets, int *** bucketsSizes)
{
	int i, j;

	if (*buckets != NULL)
	{
		for(j = 0; j < bucketsWidth; j++)
		{
			for(i = 0; i < bucketsHeight; i++)
				CC_free((*buckets)[j][i]);
			CC_free((*buckets)[j]);
			CC_free((*bucketsSizes)[j]);
		}
		CC_free(*buckets);
		CC_free(*bucketsSizes);
		*buckets = NULL;
	}
}

void FptPointMatcher::init(int p_width, int p_height, 
						   int p_maxDisparityU, int p_maxDisparityV, int p_correlationWindowSize,
						   int p_maxPointNumberPerBucket)
{
	width = p_width;
	height = p_height;

	maxDisparityU = p_maxDisparityU;
	maxDisparityV = p_maxDisparityV;

	correlationWindowSize = p_correlationWindowSize;
	maxPointNumberPerBucket = p_maxPointNumberPerBucket;

	bucketsWidth = width   / maxDisparityU + 1;
	bucketsHeight = height / maxDisparityV + 1;

	allocBuckets(&buckets1, &bucketsSizes1, maxPointNumberPerBucket);
	allocBuckets(&buckets2, &bucketsSizes2, maxPointNumberPerBucket);
}

void FptPointMatcher::fillBuckets(FptFeaturePoint_ **** buckets, int ** bucketsSizes, 
								  FptFeaturePointsArray_ * points)
{
	int i, j;
	int lostPointsNumber = 0;

	for(i = 0; i < bucketsWidth; i++)
		for(j = 0; j < bucketsHeight; j++)
			bucketsSizes[i][j] = 0;

	for(i = 0; i < points->size; i++)
	{
		FptFeaturePoint_ * point = &(points->array[i]);

		if (point->step2U > correlationWindowSize + 1 &&
			point->step2U < width - correlationWindowSize - 1 &&
			point->step2V > correlationWindowSize + 1 &&
			point->step2V < height - correlationWindowSize - 1)
		{
			point->iBucket = point->step1U / maxDisparityU;
			point->jBucket = point->step1V / maxDisparityV;

			if (bucketsSizes[point->iBucket][point->jBucket] < maxPointNumberPerBucket)
			{
				buckets[point->iBucket][point->jBucket][bucketsSizes[point->iBucket][point->jBucket]] = point;
				bucketsSizes[point->iBucket][point->jBucket]++;
			}
			else
				lostPointsNumber++;
		}
	}
	if (lostPointsNumber > 0)
		printf("%d points lost when filling buckets (FptPointMatcher::fillBuckets()).\n", lostPointsNumber); 
}


void FptPointMatcher::computeIntensityMeanAndSigma(FptFeaturePointsArray_ *points,
												   unsigned char * image, int imageWidth, int imageHeight, int imagePitch,
												   int correlationWindowSize) 
{
	int i, j, cpt, halfSize = correlationWindowSize / 2;
	FptFeaturePoint_ *pt;
	unsigned char *ptr;

	for(cpt = 0, pt = points->array; cpt < points->size; cpt++, pt++) 
	{
		double sum = 0., sum2 = 0.;

		if (pt->step2U > correlationWindowSize + 1 &&
			pt->step2U < imageWidth - correlationWindowSize - 1 &&
			pt->step2V > correlationWindowSize + 1 &&
			pt->step2V < imageHeight - correlationWindowSize - 1)
		{
			for (i = -halfSize; i <= halfSize; i++) 
			{
				ptr = image +
					(pt->step2V + i) * imagePitch + 
					pt->step2U - halfSize;

				for (j = -halfSize; j <= halfSize; j++) 
				{
					sum  += (int)(unsigned char) *ptr;
					sum2 += SQR((int) (unsigned char)*ptr);  
					ptr++;
				}
			}
		}
		float mean = sum / SQR(correlationWindowSize); 
		pt->meanI = mean;
		pt->sigmaI = sqrt(sum2 / SQR(correlationWindowSize) - SQR(mean));
	}
}

double FptPointMatcher::crossCorrelationScore(const FptFeaturePoint_ * pointPrev, const FptFeaturePoint_ * point,
											  const unsigned char * imagePrev, const unsigned char * image, int imageWidth, int imageHeight, int imagePitch,
											  const int correlationWindowSize)
{
	float sum = 0.0;
	int i, j;
	int halfSize = correlationWindowSize / 2;
	unsigned char * ptr1;
	unsigned char * ptr2;

	for(i = -halfSize; i <= +halfSize; i++)
	{

		ptr1 = (unsigned char *)(imagePrev+
			(pointPrev->step2V + i) * imagePitch +
			pointPrev->step2U - halfSize);
		ptr2 = (unsigned char *)(image+ 
			(point->step2V + i) * imagePitch+
			point->step2U - halfSize);

		for(j = 0; j < correlationWindowSize; j++)
		{
			float d1, d2;

			d1 = ptr1[j] - pointPrev->meanI;
			d2 = ptr2[j] - point->meanI;
			sum += d1 * d2;
		}

		// TEST !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
		if (sum < 0) return -1.;
	}

	return sum / (SQR(correlationWindowSize) * pointPrev->sigmaI * point->sigmaI);
}




void FptPointMatcher::lookForCorrespondents(FptFeaturePointsArray_ *points1, FptFeaturePointsArray_ *points2,
											unsigned char * image1, unsigned char * image2, int imageWidth, int imageHeight, int imagePitch,
											FptFeaturePoint_ **** buckets1, int ** bucketsSizes1,
											FptFeaturePoint_ **** buckets2, int ** bucketsSizes2,
											int correlationWindowSize, float correlationThreshold,
											int framesPairType)
{
	int i, j;
	float meanThreshold;

	meanThreshold = MEAN_THRESHOLD_FOR_CLOSE_FRAMES;

	for (i = 0; i < points1->size; i++)
	{
		points1->array[i].tempCorrespondent = NULL;
		points1->array[i].score = -2;
	}
	for (i = 0; i < points2->size; i++)
	{
		points2->array[i].tempCorrespondent = NULL;
		points2->array[i].score = -2;
	}

	for (i = 0; i < points1->size; i++)
	{
		FptFeaturePoint_ * pointA = &(points1->array[i]);
		short uA, vA;

		pointA->tempCorrespondent = NULL;

		uA = pointA->step1U;
		vA = pointA->step1V;

		if (uA > correlationWindowSize + 1 &&
			uA < imageWidth - correlationWindowSize - 1 &&
			vA > correlationWindowSize + 1 &&
			vA < imageHeight - correlationWindowSize - 1)
		{
			double bestScore = -2;
			FptFeaturePoint_ * bestCorrespondent = NULL;
			bool thereIsAMatch = false;
			int diBucket, djBucket;
			int iBucket, jBucket;

			for(diBucket = -1; diBucket <= +1; diBucket++)
			{
				iBucket = pointA->iBucket + diBucket;
				if (iBucket >= 0 && iBucket < bucketsWidth)
				{
					for(djBucket = -1; djBucket <= +1; djBucket++)
					{
						jBucket = pointA->jBucket + djBucket;

						if (jBucket >= 0 && jBucket < bucketsHeight)
						{
							FptFeaturePoint_ ** pointsB = buckets2[iBucket][jBucket];

							for(j = 0; j < bucketsSizes2[iBucket][jBucket]; j++)
							{
								FptFeaturePoint_ * pointB = pointsB[j];

								// LE TEST QUI SUIT EST UN ESSAI POUR REDUIRE LE NOMBRE D'APPELS A CROSS-CORRELATION-SCORE:
								if (ABS(pointA->meanI - pointsB[j]->meanI) < meanThreshold)
								{
									short uB, vB;

									uB = pointB->step2U;
									vB = pointB->step2V;

									// Eventuellement rajouter ici un test sur la disparite:
									// mais la recherche est deja restreinte aux buckets
									//if (inWindow(uB, vB, &window)) 
									{
										double score = crossCorrelationScore(pointA, pointB,
											image1, image2,imageWidth, imageHeight, imagePitch,
											correlationWindowSize);
										if (score > correlationThreshold) 
										{
											if (score > pointA->score) 
											{
												pointA->score = score;
												pointA->tempCorrespondent = pointB;
											}
											if (score > pointB->score) 
											{
												pointB->score = score;
												pointB->tempCorrespondent = pointA;
											}
										}
									}
								}
							}
						}
					}
				}
			}
		}
	}
}






void FptPointMatcher::matchPoints(unsigned char *image1,unsigned char *image2,int imageWidth, int imageHeight, int imagePitch,
								  FptFeaturePointsArray_ *pointList1,FptFeaturePointsArray_ *pointList2,
								  float correlationThreshold, AppariementPDI * matches, int * nbMatches,
								  int method)
{
	int i;
#ifdef _DEBUG		
	printf(" - Matching..."); fflush(stdout);
#endif

	for(i = 0; i < pointList1->size; i++) 
	{
		pointList1->array[i].step1U = pointList1->array[i].originalU + 0.5;
		pointList1->array[i].step1V = (pointList1->array[i].originalV + 0.5);
		pointList1->array[i].step2U = pointList1->array[i].originalU + 0.5;
		pointList1->array[i].step2V = (pointList1->array[i].originalV + 0.5);
	}
	for(i = 0; i < pointList2->size; i++) 
	{
		pointList2->array[i].step1U =  pointList2->array[i].originalU + 0.5;
		pointList2->array[i].step1V = (pointList2->array[i].originalV + 0.5);
		pointList2->array[i].step2U =  pointList2->array[i].originalU + 0.5;
		pointList2->array[i].step2V = (pointList2->array[i].originalV + 0.5);
	}


	computeIntensityMeanAndSigma(pointList1, image1, imageWidth, imageHeight, imagePitch, correlationWindowSize);
	computeIntensityMeanAndSigma(pointList2, image2, imageWidth, imageHeight, imagePitch,correlationWindowSize);

	fillBuckets(buckets1, bucketsSizes1, pointList1);
	fillBuckets(buckets2, bucketsSizes2, pointList2);


	lookForCorrespondents(pointList1, pointList2,
		image1, image2, imageWidth, imageHeight, imagePitch,
		buckets1, bucketsSizes1, 
		buckets2, bucketsSizes2, 
		correlationWindowSize, correlationThreshold,
		0/*CLOSE_FRAMES*/);

	*nbMatches = 0;

	for (i = 0; i < pointList1->size; i++) 
		if (pointList1->array[i].tempCorrespondent != NULL &&
			(pointList1->array[i].tempCorrespondent)->tempCorrespondent == &(pointList1->array[i])) 
		{
			matches[(*nbMatches)].p1 = pointList1->array[i].tempCorrespondent;
			matches[(*nbMatches)].p2 = &(pointList1->array[i]);
			matches[(*nbMatches)].score = pointList1->array[i].score;
			(*nbMatches)++;
		}
		else
			pointList1->array[i].tempCorrespondent = NULL;

	for (i = 0; i < pointList2->size; i++) 
		if (pointList2->array[i].tempCorrespondent != NULL &&
			(pointList2->array[i].tempCorrespondent)->tempCorrespondent != &(pointList2->array[i])) 
			pointList2->array[i].tempCorrespondent = NULL; 
#ifdef _DEBUG		
	printf(" %d matched points.\n", *nbMatches);
#endif

}