#include <iostream>
#include <vector>
#include <list>
#include <cstdlib>
#include <cmath>
#include <ctime>
using namespace std;

#include "Heartbeat.h"
#include "txs_multires_algo.h"
#include "SlideWriter.h"

void fillPixelsWithRandomImageRGBs(vector<pixelLocn> &pixLocs, TiffImage &img)
{

	__LOG__("filling " << pixLocs.size() << " pixels with random colors from the img");
	vector<pixelLocn>::iterator pit;
	for(pit = pixLocs.begin(); pit != pixLocs.end(); ++pit) {
		pixelLocn loc = *pit;
		
		loc.setRGB(randomRgbFrom(img));
	}
}

int getLoResNhoodSize(int hiSize) {
	return (int)ceil((float)hiSize/2.0);
}

// TEMP TODO
// this function is HACKED AND HARD-CODED!!
int hoodSizeForLevel(int level)
{
	switch(level) {
		case 0: return 3;
		case 1: return 5;
		case 2: return 7;
		case 3: return 9;
		default:
			__ERROR__("we shouldn't be getting this level: " << level);
			return -1;
	}
}

TXS_REAL unitsPerPixelForLevel(TxsPyramid &inPy, int level)
{
	// how big is this level?
	int levSize = inPy.getLevel(level)->getWidth();
	
	// get the finest res
	int finestSize = inPy.getFinestLevel()->getWidth();

	// compute the ratio..
	TXS_REAL unitsPerLevelPixel = (TXS_REAL)finestSize / (TXS_REAL)levSize * UNITS_PER_PIXEL;

	return unitsPerLevelPixel;
}

/*
 * This implements pass 1 of the sythesis algo.
 * See Wei&Levoy 2001
 * this assumes the pixel's level > 0
 */
void extrapPixel(pixelLocn &p, TxsPyramid &inPy)
{
	// form an nhood of the lower level
	NhoodPatch_s lowHood;
	const int lowLev = p.level - 1;
	const int lowSize = hoodSizeForLevel(lowLev);
	lowHood.setTo(p, unitsPerPixelForLevel(inPy, lowLev), lowSize);

	//__DBa(lowSize);

	// fill the nhood image
	list<TxsFace*> &faces = p.face->getImmediateNbors();
	TiffImage lowImg(lowSize, lowSize);
	fillNhoodImage(lowImg, lowHood, lowLev, faces);

	// find the best match for this hood image
	// in the lower level in the input pyramid
	// but don't ignore the middle pixel (so pass in -1)
	TiffImage *lowInImg = inPy.getLevel(lowLev);
	int mx, my;
	lowInImg->getBestMatchingRegion(lowImg, mx, my, -1);

	/////////////
	//__DBa("mx " << mx << " my " << my << " level: " << p.level);

	// now, get the middle pixel of the best match
	// and use it.
	// calculate the middle pixel offset
	const int midOfs = (lowSize - 1) / 2;
	NormRGB bestMatch = lowInImg->getRGB(mx + midOfs, my + midOfs);
	p.setRGB(bestMatch);
	//__DBa("SPP1 done");
	//
	//TEMP - the 2 lines below should yield same thing..about
	//p.setRGB(lowImg.getRGB(midOfs, midOfs));
	//p.setRGB(p.face->getRGB(lowLev, p.pos));
}


/*
 * Pass 2
 * This performs the closest-pixel search using
 * a 2-level search pyramid
 */
void setToClosest2LevPyramid(pixelLocn &p, TxsPyramid &inPy)
{
	// get all the faces necessary to build our nhood
	list<TxsFace*> &faces = p.face->getImmediateNbors();

	// setup the two nhood patches
	const int hiLev = p.level;
	const int loLev = hiLev - 1;
	const int hiHoodSize = hoodSizeForLevel(hiLev);
	const int loHoodSize = getLoResNhoodSize(hiHoodSize);

	NhoodPatch_s hiHood;
	hiHood.setTo(p, unitsPerPixelForLevel(inPy, hiLev), hiHoodSize);
	NhoodPatch_s loHood;
	loHood.setTo(p, unitsPerPixelForLevel(inPy, loLev), loHoodSize);

	// create and fill the nhood pyramid
	TiffImage hiHoodImg(hiHoodSize, hiHoodSize);
	TiffImage loHoodImg(loHoodSize, loHoodSize);
	fillNhoodImage(hiHoodImg, hiHood, hiLev, faces);
	fillNhoodImage(loHoodImg, loHood, loLev, faces);

	// we have our nhood images
	// now find the best match in the input pyramid
	
	// calculate the middle pixel offset
	const int midOfs = (hiHoodSize - 1) / 2;
	int mx = -1;
	int my = -1;
	inPy.getClosestPyramid(
			loHoodImg, hiHoodImg,
			hiLev,
			mx, my, midOfs);
	
	// do the search and
	// set p's color using the middle pixel
	//__LOG__("midOfs is " << midOfs << " BMR is " << mx << " " << my);
	NormRGB bestMatch = inPy.getLevel(hiLev)->getRGB(mx + midOfs, my + midOfs);
	p.setRGB(bestMatch);
}

void setToClosestUsing2LevPyramid(pixelLocn &p, TxsPyramid &inPy)
{
	setToClosest2LevPyramid(p, inPy);
}

/*
 * Accumulates pixel locations from the mesh at the given level
 */
int accumPixelLocs(vector<pixelLocn> &locsOut, TxsMesh &mesh, int level)
{
	__LOG__("Begin accumPixelLocs for level " << level);

	Heartbeat hb;
	list<TxsFace*> faces = mesh.getFaces();
	int num = 0;
	int faceI = 0;
	list<TxsFace*>::iterator fit;
	for(fit = faces.begin(); fit != faces.end(); ++fit) {
		TxsFace* f = *fit;
		num += f->accumPixelLocns(locsOut, level);

		// some progress output
		if(hb.OK_to_beat()) {
			__STAT__("accum'ing pixel locns " << (double)faceI/(double)faces.size()*100 << "% done");
		}

		faceI++;
	}

	__STAT__("accum'd " << num << " pixel locs for level " << level);

	return num;
}

/*
 * This assumes that all faces have the same number of levels
 */
int getFinestLevel(TxsMesh &mesh) {
	TxsFace *f = *(mesh.getFaces().begin());
	return f->getPyramid()->finestLevel();
}

int accumPixelLocsForAllLevels(vector<vector<pixelLocn>*> &pixLocLevels, TxsMesh &mesh)
{
	int finestLev = getFinestLevel(mesh);
	int lev = 0;
	int numPixLocs = 0;
	for(lev = 0; lev <= finestLev; ++lev) {
		vector<pixelLocn> *pixLocs = new vector<pixelLocn>;
		int numPixThisLevel = accumPixelLocs(*pixLocs, mesh, lev);
		numPixLocs += numPixThisLevel;

		__LOG__(numPixThisLevel << " level " << lev << " pix locations");

		// add it
		pixLocLevels.push_back(pixLocs);
	}

	return numPixLocs;
}


/*
 * The extrapolate-and-relax method
 * Basically, to my knowledge, the one that Wei&Levoy used
 */
void synthTexturesMultiRes(TxsMesh &mesh, TxsPyramid &inPy)
{
	__STAT__("BEGIN extrapolate-and-relax method");
	// front matter..
	unsigned int seed = (unsigned int)time(NULL);
	srand(seed);
	__LOG__("Rand seed: " << seed);

	// accumulate pixel locs for each level
	vector<vector<pixelLocn>*> pixLocLevels;
	int numPixLocs = accumPixelLocsForAllLevels(pixLocLevels, mesh);

	// fill the lowest with random pixels from the full-res image
	fillPixelsWithRandomImageRGBs(*pixLocLevels[0], *inPy.getFinestLevel());

	// now, go through each non-0 level and run both synthesis passes!
	__STAT__("Begin synthesis for " << numPixLocs << " pixels..");
	double startTime = clock();
	Heartbeat hb(3);
	int counter = 0;
	int finestLev = getFinestLevel(mesh);
	int lev = 0;

	for(lev = 1; lev <= finestLev; ++lev) {
		__STAT__("synth'ing level " << lev);
		vector<pixelLocn> *pixLocs = pixLocLevels[lev];

		// synth the pixels in this level
		vector<pixelLocn>::iterator pit;

		__STAT__("  phase 1 (extrap) start");

		// first, shuffle them
		knuth_shuffle<pixelLocn>(*pixLocs);

		// phase 1 - extrapolate
		for(pit = pixLocs->begin(); pit != pixLocs->end(); ++pit) {
			pixelLocn loc = *pit;
			extrapPixel(loc, inPy);

			if(hb.OK_to_beat()) {
				__STAT__("synthesized " << counter << " of " << numPixLocs << " pixels.  " << (double)counter/(double)numPixLocs*100 << "% done.");
				__STAT__("  elapsed seconds: " << (clock() - startTime)/1000);
			}
		}

		/*
		 * we may or may not need this extra pass
		__STAT__("  phase 2 start");
		
		knuth_shuffle<pixelLocn>(*pixLocs);

		for(pit = pixLocs->begin(); pit != pixLocs->end(); ++pit) {
			pixelLocn loc = *pit;
			setToClosest2LevPyramid(loc, inPy);

			// pass 2 takes longer, so use this to count
			if(hb.OK_to_beat()) {
				__STAT__("synthesized " << counter << " of " << numPixLocs << " pixels.  " << (double)counter/(double)numPixLocs*100 << "% done.");
				__STAT__("  elapsed seconds: " << (clock() - startTime)/1000);
			}
		}
		*/

		__STAT__("  phase 2 (relax) start");
		knuth_shuffle<pixelLocn>(*pixLocs);
		for(pit = pixLocs->begin(); pit != pixLocs->end(); ++pit) {
			pixelLocn loc = *pit;
			setToClosest2LevPyramid(loc, inPy);

			// pass 2 takes longer, so use this to count
			if(hb.OK_to_beat()) {
				statProgress(counter, numPixLocs, startTime);
			}
			counter++;
		}

		__STAT__("  level " << lev << "done");

		// write it out
		SlideWriter slfer;
		char outdir[100];
		sprintf(outdir, "out_mr%d", lev);
		slfer.write(outdir, &mesh, lev);
	}

	// done!
	__STAT__("DONE synthesis..");
	double endTime = clock();
	__STAT__("total synthesis time: " << (endTime - startTime)/1000.0 << " seconds for " << numPixLocs << " pixel locations.");
	__STAT__("number of nhood rays that didn't hit anything: " << g_numNhoodRaysMissed);

	// All done.  Delete the pixel location lists
	vector<vector<pixelLocn>*>::iterator itr;
	for(itr = pixLocLevels.begin(); itr != pixLocLevels.end(); ++itr) {
		vector<pixelLocn>* pixLocs = *itr;
		delete pixLocs;
	}
}