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

using namespace std;

typedef struct _Vector3D{
	double x;
	double y;
	double z;
} Vector3D;

int inv(double **b, double **result, int T);//inverse a T*T matrix
int myround(double x);

int p4(char *srcDir_fname,char *im1_fname,char *im2_fname,char *im3_fname,char *im4_fname,char *im5_fname,char *mask_fname,char *normalMap_fname,char *albedoMap_fname,char *gradient_fname)
{
	ifstream fin(srcDir_fname);
	Vector3D *S = new Vector3D[5];
	for (int k = 0; k < 5; k++)
	{
		fin>>S[k].x>>S[k].y>>S[k].z;
		//printf("(%f,%f,%f)\n",S[k].x,S[k].y,S[k].z);//for test
	}
	fin.close();
	ofstream ftest("ftest.txt");//for test

	Image **im = new Image*[5];
	for (int k = 0; k < 5; k++) 
		im[k] = new Image;//DO NOT FORGET THIS!!!OR U WILL GET SEGMENTATION FAULT
	readImage(im[0], im1_fname);
	readImage(im[1], im2_fname);
	readImage(im[2], im3_fname);
	readImage(im[3], im4_fname);
	readImage(im[4], im5_fname);
	int NRows = getNRows(im[0]);//# of rows
	int NCols = getNCols(im[0]);//# of cols

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

	ImageColor *normalMap = new ImageColor;
	setSizeColor(normalMap,NRows,NCols);
	setColorsColor(normalMap,255);

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

	/* Step 1. Compute surface normals. */
	//For each point (masked), find the 3 images where it appears brightest:
	//im[b[0]],im[b[1]],im[b[2]]
	Vector3D **M = new Vector3D*[NRows];
	for (int i = 0; i < NRows; i++)
		M[i] = new Vector3D[NCols];//M

	Vector3D **n = new Vector3D*[NRows];
	for (int i = 0; i < NRows; i++)
		n[i] = new Vector3D[NCols];//unit surface normal

	double **albedo = new double*[NRows];
	for (int i = 0; i < NRows; i++)
		albedo[i] = new double[NCols];//albedo
	double maxAlbedo = 0.0;
	
	//int prev_b[3];//for test
	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j < NCols; j++)
		{
			//for each pixel
			if (getPixel(mask,i,j)==1)
			{
				int b[5] = {0,1,2,3,4};//index of images, we want b[0],b[1],b[2] to be 3 brightest images
				for (int times = 0; times < 2; times++)//with iteratioin = 2, we can!
				{
					for (int k = 3; k >= 0; k--)
					{
						if (getPixel(im[b[k]],i,j)/sqrt(S[b[k]].x*S[b[k]].x+S[b[k]].y*S[b[k]].y+S[b[k]].z*S[b[k]].z) 
							< getPixel(im[b[k+1]],i,j)/sqrt(S[b[k+1]].x*S[b[k+1]].x+S[b[k+1]].y*S[b[k+1]].y+S[b[k+1]].z*S[b[k+1]].z))
						{
							int temp = b[k];
							b[k] = b[k+1];
							b[k+1] = temp;//swap b[k] and b[k+1]
						}
					}
				}
				//ftest<<b[0]<<" "<<b[1]<<" "<<b[2]<<"\n";//for test

				//Read Sx[b[0]],Sy[b[0]],Sz[b[0]]
				//Sx[b[1]],Sy[b[1]],Sz[b[1]]
				//Sx[b[2]],Sy[b[2]],Sz[b[2]] from file.
				//Compute M using M=IS^(-1)
				//n=M/|M|, albedo=|M|.
				double **tempS = new double*[3];
				for (int m = 0; m < 3; m++) tempS[m] = new double[3];
				double **result = new double*[3];
				for (int m = 0; m < 3; m++) result[m] = new double[3];
				tempS[0][0] = S[b[0]].x;tempS[0][1] = S[b[0]].y;tempS[0][2] = S[b[0]].z;
				tempS[1][0] = S[b[1]].x;tempS[1][1] = S[b[1]].y;tempS[1][2] = S[b[1]].z;
				tempS[2][0] = S[b[2]].x;tempS[2][1] = S[b[2]].y;tempS[2][2] = S[b[2]].z;
				inv(tempS,result,3);
				M[i][j].x = result[0][0]*getPixel(im[b[0]],i,j) + result[0][1]*getPixel(im[b[1]],i,j) + result[0][2]*getPixel(im[b[2]],i,j);
				M[i][j].y = result[1][0]*getPixel(im[b[0]],i,j) + result[1][1]*getPixel(im[b[1]],i,j) + result[1][2]*getPixel(im[b[2]],i,j);
				M[i][j].z = result[2][0]*getPixel(im[b[0]],i,j) + result[2][1]*getPixel(im[b[1]],i,j) + result[2][2]*getPixel(im[b[2]],i,j);

				albedo[i][j] = sqrt(M[i][j].x*M[i][j].x + M[i][j].y*M[i][j].y + M[i][j].z*M[i][j].z);
				if (albedo[i][j] > maxAlbedo) maxAlbedo = albedo[i][j];
				n[i][j].x = (albedo[i][j]>0)?(M[i][j].x/albedo[i][j]):0;
				n[i][j].y = (albedo[i][j]>0)?(M[i][j].y/albedo[i][j]):0;
				n[i][j].z = (albedo[i][j]>0)?(M[i][j].z/albedo[i][j]):0;

				//Store in a rgb image.
				setPixelColor(normalMap,i,j,myround((n[i][j].x+1)*127.5),myround((n[i][j].y+1)*127.5),myround((n[i][j].z+1)*127.5));
				//if ((i>=235)&&(i<=245)&&(j>=235)&&(j<=245))
				//{
				//	ftest<<"["<<i<<","<<j<<"]:"<<myround((n[i][j].x+1)*127.5)<<" "<<myround((n[i][j].y+1)*127.5)<<" "<<myround((n[i][j].z+1)*127.5)<<"\n";//for test
				//	cout<<"["<<i<<","<<j<<"]:"<<(n[i][j].x)<<" "<<(n[i][j].y)<<" "<<(n[i][j].z)<<"\n";
				//}
				//if ((i==240)&&(j==240)) setPixelColor(normalMap,i,j,0,0,0);//for test
				//ftest<<"["<<i<<","<<j<<"]:"<<n[i][j].x*n[i][j].x+n[i][j].y*n[i][j].y+n[i][j].z*n[i][j].z<<"\n";//for test

				//if (!((prev_b[0]==b[0])&&(prev_b[1]==b[1])&&(prev_b[2]==b[2]))) setPixelColor(normalMap,i,j,255,255,255);//for test
				//prev_b[0] = b[0];prev_b[1] = b[1];prev_b[2] = b[2];//for test
			}
			else
			{
				n[i][j].x = 0; n[i][j].y = 0; n[i][j].z = 0;
				albedo[i][j] = 0;
			}
		}
	}
	//ftest.close();//for test

	/* Step 2. Compute surface gradients(p,q). */
	//At each pixel, compute its (p,q) using (p,q,1)=l(nx,ny,nz), and save it in an array.
	Gradient **grd = new Gradient*[NRows];
	for (int i = 0; i < NRows; i++)
		grd[i] = new Gradient[NCols];//gradient map
	ofstream fout(gradient_fname);
	
	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j < NCols; j++)
		{
			if (getPixel(mask,i,j)==1)
			{
				grd[i][j].p = float(n[i][j].x/n[i][j].z);
				grd[i][j].q = float(n[i][j].y/n[i][j].z);
			}
			else
			{
				grd[i][j].p = 0;
				grd[i][j].q = 0;
			}
			fout.write((char *)(&grd[i][j]),sizeof(grd[i][j]));
			//ftest<<"["<<i<<","<<j<<"]:"<<grd[i][j].p<<" "<<grd[i][j].q<<"\n";//for test
		}
	}
	fout.close();
	//ftest.close();//for test	

	/* Step 3. Compute albedo. */
	//Scale albedo.
	for (int i = 0; i < NRows; i++)
	{
		for (int j = 0; j < NCols; j++)
		{
			setPixel(albedoMap,i,j,albedo[i][j]=albedo[i][j]*255/maxAlbedo);
		}
	}	
	//Draw the output image.
	writeImageColor(normalMap,normalMap_fname);
	writeImage(albedoMap,albedoMap_fname);

	printf("U got it\n");
	return 0;
}




int inv(double **b, double **result, int T)
{
	double t,x;
	double **a = new double*[T];
	for(int i=0;i<T;i++) a[i] = new double[2*T];// a is a T*2T matrix
	for(int i=0;i<T;i++)
	{       
		for(int j=0;j<(2*T);j++)  
        	{
                	if(j<T) a[i][j]=b[i][j];  
                	else if (j==(T+i)) a[i][j]=1.0;  
                	else a[i][j]=0.0;  
        	}
	}
        for(int i=0;i<T;i++)
        {
                for(int k=0;k<T;k++)
                {
                	if(k!=i)
                        {
                        	t=a[k][i]/a[i][i];
                        	for(int j=0;j<(2*T);j++)
                                {
                                        x=a[i][j]*t;
                                        a[k][j]=a[k][j]-x;
                                }
                        }
                }
        }
        for(int i=0;i<T;i++)
        {
                t=a[i][i];
                for(int j=0;j<(2*T);j++) a[i][j]=a[i][j]/t;
        }
    	double y=1;
    	for(int i=0;i<T;i++) y=y*a[i][i];
        if(y==0.0){printf("NOT INVERTIBLE!\n");exit(1);}
	for(int i=0;i<T;i++)  
		for(int j=0;j<T;j++) result[i][j]=a[i][j+T];
	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 <= 10){ 
		printf("Usage: %s <input directions> <image1> <image2> <image3> <image4> <image5> <mask> <normal map> <albedo map> <output gradient>\n", argv[0]); 
		return 0; 
	} 
	char *srcDir_fname = argv[1]; 
	char *im1_fname = argv[2];
	char *im2_fname = argv[3];
	char *im3_fname = argv[4];
	char *im4_fname = argv[5];
	char *im5_fname = argv[6];
	char *mask_fname = argv[7];
	char *normalMap_fname = argv[8];
	char *albedoMap_fname = argv[9];
	char *gradient_fname = argv[10];
	p4(srcDir_fname,im1_fname,im2_fname,im3_fname,im4_fname,im5_fname,mask_fname,normalMap_fname,albedoMap_fname,gradient_fname);
	return 0; 
}
