#include <stdlib.h> 
#include <stdio.h>
#include <cmath>
#include <iostream>
#include <fstream>
#include "vision_utilities.h"

using namespace std;

int myround(double x);

int p5(char *gradient_fname,char *seed_fname,char *mask_fname,char *depthMap_fname)
{

	Image *mask = new Image;
	readImage(mask, mask_fname);

	int NRows = getNRows(mask);//# of rows
	int NCols = getNCols(mask);//# of cols

	Gradient **grd = new Gradient*[NRows];
	for (int i = 0; i < NRows; i++)
		grd[i] = new Gradient[NCols];//gradient map

	ifstream fgrd(gradient_fname);
	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j < NCols; j++)
		{
			fgrd.read((char *)(&grd[i][j]), sizeof(grd[i][j]));
		}
	}
	fgrd.close();
	//cout<<grd[0][0].p<<","<<grd[0][0].q<<"\n"<<grd[240][240].p<<","<<grd[240][240].q<<"\n";//for test

	ifstream fseed(seed_fname);
	int numSeed = 0;
	int maxNumSeed = 10;
	int* seed_x = new int[maxNumSeed];
	int* seed_y = new int[maxNumSeed];
	while(!fseed.eof())
	{
		fseed>>seed_x[numSeed];
		fseed>>seed_y[numSeed];
		numSeed++;
	}
	numSeed--;
	fseed.close();
	//cout<<"numSeed = "<<numSeed<<"\n";//for test
	//cout<<seed_x[0]<<" "<<seed_y[0]<<"\n";//for test

	Image *depthMap = new Image;
	setSize(depthMap,NRows,NCols);
	setColors(depthMap,255);

	double **sumDepth = new double*[NRows];
	for (int i = 0; i < NRows; i++)
		sumDepth[i] = new double[NCols];

	for (int s = 0; s < numSeed; s++)
	{
		double **currentDepth = new double*[NRows];
		for (int i = 0; i < NRows; i++)
			currentDepth[i] = new double[NCols];
		currentDepth[seed_y[s]][seed_x[s]] = 0;//seed point

		/* up */
		for (int idx_up = 1; idx_up < NRows; idx_up++)
		{
			if (getPixel(mask,seed_y[s]-idx_up,seed_x[s])==1)
				currentDepth[seed_y[s]-idx_up][seed_x[s]] = currentDepth[seed_y[s]-idx_up+1][seed_x[s]] - grd[seed_y[s]-idx_up][seed_x[s]].q*(-1);
		}

		/* down */
		for (int idx_down = 1; idx_down < NRows; idx_down++)
		{
			if (getPixel(mask,seed_y[s]+idx_down,seed_x[s])==1)
				currentDepth[seed_y[s]+idx_down][seed_x[s]] = currentDepth[seed_y[s]+idx_down-1][seed_x[s]] - grd[seed_y[s]+idx_down][seed_x[s]].q*(1);
		}

		/* left */
		for (int idx_left = 1; idx_left < NCols; idx_left++)
		{
			if (getPixel(mask,seed_y[s],seed_x[s]-idx_left)==1)
				currentDepth[seed_y[s]][seed_x[s]-idx_left] = currentDepth[seed_y[s]][seed_x[s]-idx_left+1] - grd[seed_y[s]][seed_x[s]-idx_left].p*(-1);
		}	

		/* right */	
		for (int idx_right = 1; idx_right < NCols; idx_right++)
		{
			if (getPixel(mask,seed_y[s],seed_x[s]+idx_right)==1)
				currentDepth[seed_y[s]][seed_x[s]+idx_right] = currentDepth[seed_y[s]][seed_x[s]+idx_right-1] - grd[seed_y[s]][seed_x[s]+idx_right].p*(1);
		}

		double maxDepth = 0;
		double minDepth = 0;
		for (int i = 0; i < NRows; i++)
		{
			for (int j = 0; j < NCols; j++)
			{
				if (getPixel(mask,i,j)==1)
				{
					if (currentDepth[i][j] > maxDepth) maxDepth = currentDepth[i][j];
					if (currentDepth[i][j] < minDepth) minDepth = currentDepth[i][j];
				}
			}
		}

		// down-right
		for (int i = seed_y[s]+1; i < NRows; i++)
		{
			for (int j = seed_x[s]+1; j < NCols; j++)
			{
				if (getPixel(mask,i,j)==1)
				{
					if (getPixel(mask,i-1,j)==0)
						currentDepth[i][j] = currentDepth[i-1][j] - grd[i][j].q*(1);
					else if (getPixel(mask,i,j-1)==0)
						currentDepth[i][j] = currentDepth[i][j-1] - grd[i][j].p*(1);
					else
					{
						currentDepth[i][j] = 0.5*(currentDepth[i-1][j] - grd[i][j].q*(1))
							+0.5*(currentDepth[i][j-1] - grd[i][j].p*(1));
					}						
				}
			}
		}

		// up-right
		for (int i = seed_y[s]-1; i >= 0; i--)
		{
			for (int j = seed_x[s]+1; j < NCols; j++)
			{
				if (getPixel(mask,i,j)==1)
				{
					if (getPixel(mask,i+1,j)==0)
						currentDepth[i][j] = currentDepth[i+1][j] - grd[i][j].q*(-1);
					else if (getPixel(mask,i,j-1)==0)
						currentDepth[i][j] = currentDepth[i][j-1] - grd[i][j].p*(1);
					else
					{
						currentDepth[i][j] = 0.5*(currentDepth[i+1][j] - grd[i][j].q*(-1))
							+0.5*(currentDepth[i][j-1] - grd[i][j].p*(1));
					}						
				}
			}
		}

		// up-left
		for (int i = seed_y[s]-1; i >= 0; i--)
		{
			for (int j = seed_x[s]-1; j >= 0; j--)
			{
				if (getPixel(mask,i,j)==1)
				{
					if (getPixel(mask,i+1,j)==0)
						currentDepth[i][j] = currentDepth[i+1][j] - grd[i][j].q*(-1);
					else if (getPixel(mask,i,j+1)==0)
						currentDepth[i][j] = currentDepth[i][j+1] - grd[i][j].p*(-1);
					else
					{
						currentDepth[i][j] = 0.5*(currentDepth[i+1][j] - grd[i][j].q*(-1))
							+0.5*(currentDepth[i][j+1] - grd[i][j].p*(-1));
					}						
				}
			}
		}

		// down-left
		for (int i = seed_y[s]+1; i < NRows; i++)
		{
			for (int j = seed_x[s]-1; j >= 0; j--)
			{
				if (getPixel(mask,i,j)==1)
				{
					if (getPixel(mask,i-1,j)==0)
						currentDepth[i][j] = currentDepth[i-1][j] - grd[i][j].q*(1);
					else if (getPixel(mask,i,j+1)==0)
						currentDepth[i][j] = currentDepth[i][j+1] - grd[i][j].p*(-1);
					else
					{
						currentDepth[i][j] = 0.5*(currentDepth[i-1][j] - grd[i][j].q*(1))
							+0.5*(currentDepth[i][j+1] - grd[i][j].p*(-1));
					}						
				}
			}
		}


		/* scale to [0,1]*/
/*		double maxDepth = 0;
		double minDepth = 0;
		for (int i = 0; i < NRows; i++)
		{
			for (int j = 0; j < NCols; j++)
			{
				if (getPixel(mask,i,j)==1)
				{
					if (currentDepth[i][j] > maxDepth) maxDepth = currentDepth[i][j];
					if (currentDepth[i][j] < minDepth) minDepth = currentDepth[i][j];
				}
			}
		}*/
		//cout<<"maxDepth = "<<maxDepth<<"\nminDepth = "<<minDepth<<"\n";//for test

		//ofstream ftest("ftest.txt");//for test
		for (int i = 0; i < NRows; i++)
		{
			for (int j = 0; j < NCols; j++)
			{
				if (getPixel(mask,i,j)==1)
				{
					currentDepth[i][j] = CLAMP(currentDepth[i][j],minDepth,maxDepth);
					currentDepth[i][j] = (currentDepth[i][j]-minDepth)/(maxDepth-minDepth);
					sumDepth[i][j] += currentDepth[i][j];
					//ftest<<"["<<i<<","<<j<<"]:"<<currentDepth[i][j]<<"\n";//for test
				}
			}
		}

		//ftest.close();//for test


	}//end for

	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j < NCols; j++)
		{
			if (getPixel(mask,i,j)==1) 
			{
				sumDepth[i][j] /= numSeed;
				setPixel(depthMap,i,j,myround(sumDepth[i][j]*255));
			}
		}
	}

	writeImage(depthMap,depthMap_fname);
	printf("U got it\n");
	return 0;
}

int myround(double x)
{
	if (x<floor(x)+0.50) return floor(x);
	else return (floor(x)+1);
}

int main(int argc, char **argv) 
{ 
	if (argc <= 4){ 
		printf("Usage: %s <input gradient file> <seed point file> <binary mask> <output depth map>\n", argv[0]); 
		return 0; 
	} 
	char *gradient_fname = argv[1]; 
	char *seed_fname = argv[2];
	char *mask_fname = argv[3];
	char *depthMap_fname = argv[4];
	p5(gradient_fname,seed_fname,mask_fname,depthMap_fname);
	return 0; 
}
