#include "main.h"
#include "../Util/JacobiSolver.h"
#include "../Util/MultiGridSolver.h"

#include <conio.h>
#include <iostream>
#include <math.h>

// Libra
#include <Libra.h>    // C and C++ API


#define ITERCOUNT 24
#define TIMESTEPS 1000
#define GRIDSIZE 512


void setCmdLine(int argc, char** argv, int *customArgc, char*** customArgv)
{
	int i=0;
	for (;i<argc;i++)
	{
		const size_t n = strlen(argv[i]);
		(*customArgv)[i] = new char[n+1];
		memcpy((*customArgv)[i], argv[i], n);
		(*customArgv)[i][n] = '\0';
	}

	// Set -c ClientSDK, -? for help.
	(*customArgv)[i] = "-c";
	(*customArgv)[i+1] = "ClientSDK";
	(*customArgv)[i+2] = "\0";
}

void solveSimulation();


int main(int argc, char** argv)
{
	int    customArgc = argc+3;
	char** customArgv = new char*[customArgc];
	setCmdLine(argc, argv, &customArgc, &customArgv);

	std::cout << "****************" << std::endl << std::endl;
	std::cout << "Fluid simulation" << std::endl << std::endl;
	std::cout << "****************" << std::endl << std::endl;

	if (libra_Init(customArgc, customArgv) != 0)
		return 1;

	libra_SetDefaultDataType(GFLOAT64);

	// Solve system
	solveSimulation();

	libra_Shutdown();
	printf("Press a key to exit...\n");
	_getch();

	delete[] customArgv;
}

#define IX(i,j) ((i)+(N+2)*(j))
#define IXF(i,j) ((float)((i)+(N+2)*(j)))

gVar diffMtx, projMtx;
gVar diffMtxP, projMtxP;

gVar imap, jmap;
gVar ijmap, ijmap0, ijmap1, ijmap2, ijmap3;

gVar bndMap0, bndMap1, bndMap2;
gVar bndMapCorner0, bndMapCorner1, bndMapCorner2;






void setupMatrix(float diagCoeff, float coeff, int N, gVar &A)
{
	int i;

	int size = (N+2)*(N+2);

	float *data1 = new float[size];
	float *data2 = new float[size];
	float *data3 = new float[size];

	// setDataup main diagonal-Y
	for (i = 0; i < size; i++) 
		data1[i] = coeff;

	// setDataup main diagonal-1
	for (i = 0; i < size; i++) 
		data2[i] = coeff;

	// setDataup main diagonal
	for (i = 0; i < size; i++)
		data3[i] = diagCoeff;

	const int bandCount = 5;
	int bandOffsetDatas[] = {-(N+2), -1, 0, 1, (N+2)};
	float* bandValues[] = {data1, data2, data3, data2, data1};

	// create the sparse matrix with input on band form
	A = sparseBanded(size, size, bandValues, bandOffsetDatas, bandCount);

	delete []data1;
	delete []data2;
	delete []data3;
}


void set_bnd(int N, int b, gVar& x)
{
	return;
	static bool foo = false;
	static gVar foo1, foo2, foo3, foo4;
	static gVar fum1, fum2, fum3, fum4;
	static gVar fo1, fo2, fo3;

	if (!foo)
	{
		foo = true;
		foo1 = cast(step(IXF(0,1),   IXF(0,N),   N+2.f), GINT32);
		foo2 = cast(step(IXF(N+1,1), IXF(N+1,N), N+2.f), GINT32);
		foo3 = cast(step(IXF(1,0),   IXF(N,0),     1.f), GINT32);
		foo4 = cast(step(IXF(1,N+1), IXF(N,N+1),   1.f), GINT32);
		fum1 = cast(step(IXF(1,1),   IXF(1,N),   N+2.f), GINT32);
		fum2 = cast(step(IXF(N,1),   IXF(N,N),   N+2.f), GINT32);
		fum3 = cast(step(IXF(1,1),   IXF(N,1),     1.f), GINT32);
		fum4 = cast(step(IXF(1,N),   IXF(N,N),     1.f), GINT32);

		fo1 = gVar(2, 2);
		fo2 = gVar(2, 2);
		fo3 = gVar(2, 2);

		int fd1[] = {IX(0,0), IX(0, N+1), IX(N+1,0), IX(N+1, N+1)};
		int fd2[] = {IX(1,0), IX(1, N+1), IX(N,0), IX(N, N+1)};
		int fd3[] = {IX(0,1), IX(0, N), IX(N+1,1), IX(N+1, N)};
		fo1.setData(fd1);
		fo2.setData(fd2);
		fo3.setData(fd3);
	}

	if (b == 0)
	{
		scatter(foo1, gather(fum1, x), x);
		scatter(foo2, gather(fum2, x), x);
		scatter(foo3, gather(fum3, x), x);
		scatter(foo4, gather(fum4, x), x);
	}
	if (b == 1)
	{
		scatter(foo1, -gather(fum1, x), x);
		scatter(foo2, -gather(fum2, x), x);
		scatter(foo3, gather(fum3, x), x);
		scatter(foo4, gather(fum4, x), x);
	}
	if (b == 2)
	{
		scatter(foo1, gather(fum1, x), x);
		scatter(foo2, gather(fum2, x), x);
		scatter(foo3, -gather(fum3, x), x);
		scatter(foo4, -gather(fum4, x), x);
	}

	scatter(fo1, 0.5f * (gather(fo2, x) + gather(fo3, x)), x);
}


void add_source(gVar& p, gVar& u, gVar& v, const gVar& p0, const gVar& u0, const gVar& v0, float dt)
{
	p += dt * p0;
	u += dt * u0;
	v += dt * v0;
}

void diffuse(int N, gVar& p, gVar& u, gVar& v, const gVar& p0, const gVar& u0, const gVar& v0, float diff, float dt)
{
	float a = dt * diff * N * N;

	if (diff == 0.f)
	{
		p = p0;
		u = u0;
		v = v0;
	}
	else
	{
		mg_solve(diffMtx, diffMtxP, p, p0, ITERCOUNT);
		mg_solve(diffMtx, diffMtxP, u, u0, ITERCOUNT);
		mg_solve(diffMtx, diffMtxP, v, v0, ITERCOUNT);
	}

	set_bnd(N, 0, p);
	set_bnd(N, 1, u);
	set_bnd(N, 2, v);
}

void advect(int N, gVar& p, gVar& u, gVar& v, const gVar& p0, const gVar& u0, const gVar& v0, float dt)
{
	float dt0 = dt*N;

	gVar i = min(max(imap - dt0 * gather(ijmap, u0), 0.5f), N + 0.5f);
	gVar j = min(max(jmap - dt0 * gather(ijmap, v0), 0.5f), N + 0.5f);

	gVar i0 = floor(i);
	gVar j0 = floor(j);
	gVar i1 = i0 + 1.f;
	gVar j1 = j0 + 1.f;

	gVar s1 = i - i0;
	gVar t1 = j - j0;
	gVar s0 = i1 - i;
	gVar t0 = j1 - j;

	gVar g0 = i0 + j0 * IXF(0,1);
	gVar g1 = i0 + j1 * IXF(0,1);
	gVar g2 = i1 + j0 * IXF(0,1);
	gVar g3 = i1 + j1 * IXF(0,1);

	scatter(ijmap,
		mul(s0, mul(t0, gather(g0, p0)) + mul(t1, gather(g1, p0))) +
		mul(s1, mul(t0, gather(g2, p0)) + mul(t1, gather(g3, p0))),
			p);

	scatter(ijmap,
		mul(s0, mul(t0, gather(g0, u0)) + mul(t1, gather(g1, u0))) +
		mul(s1, mul(t0, gather(g2, u0)) + mul(t1, gather(g3, u0))),
			u);

	scatter(ijmap,
		mul(s0, mul(t0, gather(g0, v0)) + mul(t1, gather(g1, v0))) +
		mul(s1, mul(t0, gather(g2, v0)) + mul(t1, gather(g3, v0))),
			v);

	set_bnd(N, 0, p);
	set_bnd(N, 1, u);
	set_bnd(N, 2, v);
}

void project(int N, gVar& u, gVar& v, gVar& p, gVar& div)
{
	float h;
	h = 1.0f/N;

	scatter(ijmap, -0.5f * h * (
		gather(ijmap1, u) - gather(ijmap0, u) +
		gather(ijmap3, v) - gather(ijmap2, v)), div);
	set_bnd(N, 0, div);

	p = zeros(p.rowCount(), p.columnCount());

	mg_solve(projMtx, projMtxP, p, div, ITERCOUNT);

	set_bnd(N, 0, p);

	scatter(ijmap, gather(ijmap, u) -
		(0.5f / h )* (gather(ijmap1, p) - gather(ijmap0, p)), u);

	scatter(ijmap, gather(ijmap, v) -
		(0.5f / h) * (gather(ijmap3, p) - gather(ijmap2, p)), v);

	set_bnd(N, 1, u);
	set_bnd(N, 2, v);
}

#define VSWAP(x0,x) {gVar tmp=x0;x0=x;x=tmp;}

void step(int N,
		  gVar& p, gVar& p0,
		  gVar& u, gVar& u0,
		  gVar& v, gVar& v0,
			float diff, float visc,
				float dt)
{
	add_source(p, u, v, p0, u0, v0, dt);

	VSWAP(p0, p);	VSWAP(u0, u);	VSWAP(v0, v);

	diffuse(N, p, u, v, p0, u0, v0, diff, dt);
	project(N, u, v, u0, v0);

	VSWAP(p0, p);	VSWAP(u0, u);	VSWAP(v0, v);

	advect(N, p, u, v, p0, u0, v0, dt);
	project(N, u, v, u0, v0 );
}

#define SIZE ((GRIDSIZE+2)*(GRIDSIZE+2))

void get_from_UI(gVar& d, gVar& u, gVar& v)
{
	float force = 5.0f;
	float source = 70.0f;

	static double alpha = 0.f;
	static int count = 0;

	alpha += 3.1415927 / 300;
	alpha = fmod(alpha, 2 * 3.1415927);

	count++;

	static float mx = 0.f, my = 0.f, omx = 0.f, omy = 0.f;
	mx = GRIDSIZE / 2.f;
	my = GRIDSIZE / 2.f;
	omx = GRIDSIZE / 2.f-3*(float)cos(alpha);
	omy = GRIDSIZE / 2.f-3*(float)sin(alpha);

	int N = GRIDSIZE;

	int i, j, size = (N+2)*(N+2);

	u = zeros(size, 1);
	v = zeros(size, 1);
	d = zeros(size, 1);

	i = (int)mx; j = (int)my;

	scatter(IX(i,j), force * (mx-omx), u);
	scatter(IX(i,j), force * (omy-mx), v);

//	if (count < 300)
		scatter(IX(i,j), source, d);

	omx = mx;
	omy = my;
}


void solveSimulation()
{
	int N = GRIDSIZE;

	double totalTime = 0.0, iterationTime = 0.0, averageIterationTime;
	const int timeSteps = TIMESTEPS;

	double startTime = libra_GetTime();

	float dt = 0.1f;
	float diff = 0.000f;
	float visc = 0.000f;

	int size = (N+2)*(N+2);

	gVar p = zeros(size, 1);
	gVar u = zeros(size, 1);
	gVar v = zeros(size, 1);
	gVar p0 = zeros(size, 1);
	gVar u0 = zeros(size, 1);
	gVar v0 = zeros(size, 1);

	imap = cast(indexmap(N,N, 1.f, 0.f, 1.f), GINT32);
	jmap = cast(indexmap(N,N, 0.f, 1.f, 1.f), GINT32);
	ijmap = cast(indexmap(N,N, IXF(1,0), IXF(0,1), IXF(1,1)), GINT32);
	ijmap0 = cast(ijmap + IXF(-1, 0), GINT32);
	ijmap1 = cast(ijmap + IXF( 1, 0), GINT32);
	ijmap2 = cast(ijmap + IXF( 0,-1), GINT32);
	ijmap3 = cast(ijmap + IXF( 0, 1), GINT32);

	float alpha = dt * diff * N * N;

	setupMatrix(4 * alpha + 1, -alpha, N, diffMtx);
	setupMatrix(4.f, -1.f, N, projMtx);
	diffMtxP = diag(diffMtx);
	projMtxP = diag(projMtx);

	for (int i = 0; i < timeSteps; i++)
	{
		get_from_UI ( p0, u0, v0);
		step(N, p, p0, u, u0, v, v0, diff, visc, dt);
		render(p);
	}

	totalTime = libra_GetTime() - startTime;
	averageIterationTime = totalTime / timeSteps;

	std::cout << "Total time (s) : " << totalTime << std::endl;
	std::cout << "Average iteration time (ms) : " << averageIterationTime * 1000.0 << std::endl;
	std::cout << std::endl;
}