#include <iostream>
#include <algorithm>
#include <stdio.h>
#include <math.h>

#include "entity.h"
#include "triangle.h"
#include "sphere.h"
#include "color.h"
#include "disjoint_entities.h"
#include "world_objects.h"
#include "write_rgb.h"
#include "tri_mesh.h"
#include "entity_transform.h"
#include "light.h"
#include "load_scene.h"

using std::fill;

static double pi = atan(1) * 4;
static double infinity = -log(0);

struct scene_info
{
	world_objects * world;
	vector3d * base, * dx, * dy, * camera;
};

long adaptive_sample(scene_info * si, long* hc, long* vc,
					 const vector3d mybase);
static double sample_tol = 0.1;
int num_samples = 0;
int main(int argc, char**argv)
{
	world_objects world;
	scene_view sv;
	
	for(int i=1; i < argc; i++)
	  load_scene(argv[i], world, &sv);

#if 0
	Ambient_Background.rgb = vector3d(0.2, 0.2, 0);
	int adaptive_sample = 1;

	Light L0, L1, L2, L3;
	L0.cs.rgb = unit_x + unit_y; // red + green = yellow
	L1.cs.rgb = vec_one; // white
//	L2.cs.rgb = unit_z; // blue
	L2.cs.rgb = vec_one; // blue
	L3.cs.rgb = vec_one; // white
//	L0.cd.rgb = unit_x; // red
	L0.cd.rgb = vec_one; // red
	L1.cd.rgb = vec_one; // white
//	L2.cd.rgb = unit_z; // blue
	L2.cd.rgb = vec_one; // blue
	L3.cd.rgb = vec_one; // white
	L0.I = .5;
	L1.I = 1.5;
	L2.I = .5;
	L3.I = .5;
	L0.pt = vector3d(20, 10, 20);
	L1.pt = vector3d(-2, -2, 0);
	L2.pt = vector3d(20, -20, -10);
	L3.pt = vector3d(0, 0, 10);

	world.lights.push_back(L2);
	world.lights.push_back(L3);
	world.lights.push_back(L1);
	world.lights.push_back(L0);
	tri_mesh cyl("cyl_poly.iv.smf");
	cyl.ks = .6;
	cyl.kd = .4;
	cyl.n = 6;
	cyl.ref = 1;
	cyl.phong = true;
	cyl.CalcNorms();
	cyl.cs.rgb = vector3d(.5, 1, 1);
	cyl.cd.rgb = vector3d(.5, 1, 1);
	tri_mesh box("outx.smf");
	box.ks = .5;
	box.kd = .5;
	box.n = 50;
	box.ref = 0;
	box.phong = false;
	box.CalcNorms();
	box.cs.rgb = vector3d(1, 1, 0);
	box.cd.rgb = vector3d(1, 1, 0);
	Sphere sph(vector3d(0, 0, -1.5), .5);
	sph.ks = .5;
	sph.kd = .5;
	sph.n = 20;
	sph.ref = 0;
	sph.cd.rgb = vector3d(1, 1, 1);
	sph.cs.rgb = vector3d(1, 1, 1);
	Sphere sph1(vector3d(0, 0, 3.7), 1);
	sph1.ks = 1;
	sph1.kd = 0;
	sph1.n = 10;
	sph1.ref = 1;
	sph1.cs.rgb = vec_one;
	sph1.cd.rgb = vec_one;
	DisjointEntities bs(&sph, &cyl);
	Entity_Translate et0[5];
	et0[0].SetChildToParent(vector3d(1.1, 2.6, 0.0));
	et0[1].SetChildToParent(vector3d(-1.1, 2.6, 0.0));
	et0[2].SetChildToParent(vector3d(1.1, -2.6, 0.0));
	et0[3].SetChildToParent(vector3d(-1.1, -2.6, 0.0));
	et0[4].SetChildToParent(vector3d(0.0, 0.0, 1.0));
	et0[0].E = &bs;
	et0[1].E = &bs;
	et0[2].E = &bs;
	et0[3].E = &bs;
	et0[4].E = &box;
	DisjointEntities dje0(&et0[0], &et0[1]);
	DisjointEntities dje1(&et0[2], &et0[3]);
	DisjointEntities dje2(&et0[4], &dje0);
	DisjointEntities dje3(&dje1, &dje2);

	world.Insert(&dje3);
	Transform tr0;
	tr0.Rotation(unit_x, pi/6);
	tr0.Scale(.5);
	tr0.Translate(vector3d(0,0,-2));
	Entity_Transform tr1;
	tr1.E = &dje3;
	tr1.SetChildToParent(tr0);

	world.Insert(&tr1);
	world.Insert(&sph1);

	tri_mesh cow("cow.smf");
	cow.ks = .5;
	cow.kd = .5;
	cow.n = 50;
	cow.ref = 0;
	cow.phong = true;
	cow.CalcNorms();
	cow.cd.rgb = vector3d(1, 0, 0);
	cow.cs.rgb = vector3d(1, 0, 0);
	tri_mesh bunny("bunny_5k.smf");
	bunny.ks = .5;
	bunny.kd = .5;
	bunny.n = 50;
	bunny.ref = 0;
	bunny.phong = true;
	bunny.CalcNorms();
	bunny.cs.rgb = vector3d(0, 0, 1);
	bunny.cd.rgb = vector3d(0, 0, 1);

	Transform trcow;
	trcow.Rotation(unit_x, -pi/2);
	trcow.Scale(4);
	trcow.Translate(vector3d(0,-3,3));
	Transform trbunny;
	trbunny.Rotation(unit_x, -pi/2);
	trbunny.Scale(4);
	trbunny.Translate(vector3d(0,3,3));

	Entity_Transform etcow;
	Entity_Transform etbunny;
	etcow.SetChildToParent(trcow);
	etbunny.SetChildToParent(trbunny);
	etcow.E = &cow;
	etbunny.E = &bunny;
	world.Insert(&etcow);
	world.Insert(&etbunny);

	Sphere sph2(vector3d(-8, 5, 5), 5);
	sph2.ks = .9;
	sph2.kd = .1;
	sph2.n = 10;
	sph2.ref = .9;
	sph2.cs.rgb = vec_one;
	sph2.cd.rgb = vec_one;

	Sphere sph3(sph2), sph4(sph2), sph5(sph2), sph6(sph2);
	sph3.SetCenter(vector3d(-8, 5, -5));
	sph4.SetCenter(vector3d(-8, -5, 5));
	sph5.SetCenter(vector3d(-8, -5, -5));
	sph6.SetCenter(vector3d(-8, 0, 0));
	sph6.SetRadius(2.0);

	world.Insert(&sph2);
	world.Insert(&sph3);
	world.Insert(&sph4);
	world.Insert(&sph5);
	world.Insert(&sph6);

	vector3d camera(10, 0, 0), view_pt(0, 0, 0);
#endif

	world.PathCompress();

#if 0
	int screen_w = 512, screen_h = 512;
	int sample_factor = 1;
#endif

	long ** screen = new long*[sv.screen_w*sv.sample_factor];
	*screen = new long[sv.screen_w*sv.screen_h*sv.sample_factor*sv.sample_factor];
	for(int i=1; i<sv.screen_w*sv.sample_factor; i++)
		screen[i] = screen[i-1] + sv.screen_h*sv.sample_factor;

	vector3d pixel_x(0, 10.0/(sv.screen_w*sv.sample_factor), 0);
	vector3d pixel_y(0, 0, 10.0/(sv.screen_h*sv.sample_factor));
	vector3d pixel_base(0, -5, -5);

	world.SetupBuckets();

	Ray ray;
	ray.o = sv.camera;
	vector3d pixel_loc;

	long vert_cache[5] = {-1, -1, -1, -1, -1};
	long * horiz_cache = new long[sv.screen_w * 4 + 1];
	fill(horiz_cache, horiz_cache + sv.screen_w * 4 + 1, -1);

	scene_info si = {&world, &pixel_base, &pixel_x, &pixel_y, &sv.camera};

	for(int i=0; i<sv.screen_w*sv.sample_factor; i++)
	{
		pixel_loc = pixel_base + pixel_x * i;
		for(int j=0; j<sv.screen_h*sv.sample_factor; j++)
		{
		  if(!sv.enable_adaptive_sample)
		    {
			ray.d = (pixel_loc - sv.camera).ret_unit();
			IntersectInfo ii = world.CalcRayIntersectBucket(ray, BACKFACE_CULL,
															1.0F, infinity);
			num_samples++;
			screen[i][j] = ii.col.GetRGB();
		    }
		  else
		    {
			screen[i][j] = adaptive_sample(&si, horiz_cache + 4 * i,
										   vert_cache, pixel_loc);
		    }
			pixel_loc += pixel_y;
		}
	}
	for(int i=0; i<sv.screen_w; i++)
	{
		for(int j=0; j<sv.screen_h; j++)
		{
			int totalr = 0;
			int totalg = 0;
			int totalb = 0;
			for(int a = 0; a < sv.sample_factor; a++)
				for(int b = 0; b < sv.sample_factor; b++)
				{
					long x = screen[i*sv.sample_factor+a][j*sv.sample_factor+b];
					totalr += x & 0xFF0000;
					totalg += x & 0x00FF00;
					totalb += x & 0x0000FF;
				}
			int div = sv.sample_factor * sv.sample_factor;
			totalr = (totalr / div) & 0xFF0000;
			totalg = (totalg / div) & 0x00FF00;
			totalb = (totalb / div) & 0x0000FF;
			screen[i][j] = totalr | totalg | totalb;
		}
	}

	save_rgb("output.rgb", sv.screen_h, sv.screen_w, screen);
	printf("samples: %i\n", num_samples);
	delete [] *screen;
	delete [] screen;
	return 0;
}

long get_sample(long grid[][5], int x, int y, scene_info * si,
				const vector3d& mybase)
{
	if(grid[x][y] >= 0)
		return grid[x][y];

	Ray ray;
	ray.o = *si->camera;
	vector3d pl = mybase + (x/4.0) * *si->dx + (y/4.0) * *si->dy;
	ray.d = (pl - ray.o).ret_unit();
	IntersectInfo ii = si->world->CalcRayIntersectBucket(ray, BACKFACE_CULL,
														 1.0F, infinity);
	num_samples++;
	return grid[x][y] = ii.col.GetRGB();
}

float color_diff(long a, long b)
{
	int ar = (int)(a >> 16);
	int ag = (int)((a >> 8) & 0xFF);
	int ab = (int)(a & 0xFF);
	int br = (int)(b >> 16);
	int bg = (int)((b >> 8) & 0xFF);
	int bb = (int)(b & 0xFF);

	float dr = fabs(ar - br) / 0x100;
	float dg = fabs(ag - bg) / 0x100;
	float db = fabs(ab - bb) / 0x100;
	return (dr+dg+db) / 3;
}

int fits_toler(long grid[][5], int x, int y, int dx, int dy)
{
	if(color_diff(grid[x][y], grid[x][y+dy]) > sample_tol) return 0;
	if(color_diff(grid[x][y], grid[x+dx][y]) > sample_tol) return 0;
	if(color_diff(grid[x+dx][y+dy], grid[x][y+dy]) > sample_tol) return 0;
	if(color_diff(grid[x+dx][y+dy], grid[x+dx][y]) > sample_tol) return 0;
	return 1;
}

long ave_color(long grid[][5], int x, int y, int dx, int dy)
{
	long r = 0, g = 0, b = 0;
	for(int i=0; i<2; i++)
		for(int j=0; j<2; j++)
		{
			long z = grid[x+i*dx][y+j*dy];
			r += z & 0xFF0000;
			g += z & 0x00FF00;
			b += z & 0x0000FF;
		}
	r = r / 4 & 0xFF0000;
	g = g / 4 & 0x00FF00;
	b = b / 4 & 0x0000FF;
	return r | g | b;
}

long super_sample(long grid[][5], int x, int y, int dx, int dy,
				  scene_info* si, const vector3d& mybase);

long adaptive_sample(scene_info * si, long* hc, long* vc,
					 const vector3d mybase)
{
	long grid[5][5];
	for(int i=0; i<5; i++)
	{
		grid[0][i] = vc[i];
		grid[i][0] = hc[i];
	}

	long vals[2][2];
	for(int i=0; i<4; i++)
		for(int j=0; j<4; j++)
			vals[i][j] = grid[i+1][j+1] = -1;

	get_sample(grid, 0, 0, si, mybase);	
	get_sample(grid, 0, 4, si, mybase);	
	get_sample(grid, 4, 0, si, mybase);	
	get_sample(grid, 4, 4, si, mybase);	

	long ret;
	if(!fits_toler(grid, 0, 0, 4, 4))
		ret = super_sample(grid, 0, 0, 2, 2, si, mybase);
	else ret = ave_color(grid, 0, 0, 4, 4);

	for(int i=0; i<5; i++)
	{
		vc[i] = grid[4][i];
		hc[i] = grid[i][4];
	}
	return ret;
}

long super_sample(long grid[][5], int x, int y, int dx, int dy,
				  scene_info* si, const vector3d& mybase)
{
	for(int i=0; i<3; i++)
		for(int j=0; j<3; j++)
			get_sample(grid, x+dx*i, y+dy*j, si, mybase);

	long vals[2][5];

	for(int i=0; i<2; i++)
		for(int j=0; j<2; j++)
			if(dx > 1 && dy > 1 && !fits_toler(grid, x+dx*i, y+dy*j, dx, dy))
				vals[i][j] = super_sample(grid, x+dx*i, y+dy*j,
										  dx/2, dy/2, si, mybase);
			else
				vals[i][j] = ave_color(grid, x+dx*i, y+dy*j, dx, dy);

	return ave_color(vals, 0, 0, 1, 1);
}