//
// WaterDropletsSample.h
//

#ifndef WaterDropletsSample_H
#define WaterDropletsSample_H

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

class WaterDropletsSample
{
public:
	
	// Implements LibraSample.h (Native Code Interface)
	virtual int init(int argc, char **argv);
	virtual void getRenderBufferBytes8888(int* renderBufferBytes);
	virtual void tick();
	virtual void shutdown();

	// Custom sample functions
	void setupWaveEquationSystem();
	void solveRenderWaveEquation();

private:

	double totalTime;
	double iterationTime;
	double averageIterationTime;
	const static int timeSteps = 5000;
	const static int gridSizeX = 256;
	const static int gridSizeY = 256;
	double startTime;
	int i;
	float alpha;
	gVar M, s, o;
	gVar A,x,b;
	gVar g1, g2, g3, g4;
};

#endif


//
// WaterDropletsSample.cpp
//

#include "WaterDropletsSample.h"
#include <stdio.h>
#include <iostream>
#include <string.h>

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] = (char*)"-c";
	(*customArgv)[i+1] = (char*)"ClientSDK";
	(*customArgv)[i+2] = (char*)"\0";
}

int main(int argc, char** argv)
{
	//
	WaterDropletsSample sample;

	// Init sample
	sample.init(argc, argv);

	// Process, solve sample
	sample.solveRenderWaveEquation();

	// Cleanup sample
	sample.shutdown();
}

int WaterDropletsSample::init(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 << "Water droplets simulation" << std::endl << std::endl;
	std::cout << "*************************" << std::endl << std::endl;

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

	// 
	libra_SetDefaultDataType(GFLOAT32);

	// Setup system
	setupWaveEquationSystem();

	delete[] customArgv;

	return 0;
}

void WaterDropletsSample::getRenderBufferBytes8888(int* renderBufferBytes)
{
	libra_GetRenderBufferBytes8888(renderBufferBytes);
}

void WaterDropletsSample::shutdown()
{
	libra_Shutdown();
}

void WaterDropletsSample::solveRenderWaveEquation()
{
	totalTime = 0.0, iterationTime = 0.0;
	startTime = libra_GetTime();

	// for each step in time...
	for (i=0;i<timeSteps;++i)
	{ 
		tick();
	}

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

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

void WaterDropletsSample::tick()
{
//	startTime = libra_GetTime();

	// Create new drops by scattering a random value
	// into x at the random position
	// Scatter heavy rain
	scatter(rand(15, 1) * (float)(x.rowCount() - 1.f), rand(15, 1)*0.2f-0.1f, x);

	// Setup right hand side (b) of equation
	b = alpha * (gather(g1, x) + 
				gather(g2, x) + 
				gather(g3, x) + 
				gather(g4, x)) + (2.0f - 4.0f * alpha) * x - o;

	// o <-- old solution
	o = x;

	// x <-- solve for new solution
	for (int j=0;j<3;++j)
		x = jacobiIter(A, x, b);

	// Rescale and render solution to screen
	render(log10(o + 1.0f)*15.0f, gRenderMode(GHEIGHTFIELD2D), gridSizeX);

//	totalTime = libra_GetTime() - startTime;
//	std::cout << "Total time (ms) : " << totalTime*1000.0 << std::endl;
}

void WaterDropletsSample::setupWaveEquationSystem()
{
	//
	// 2D Wave equation
	// 

	const float C = 0.5f;
	const float h = 1.f;
	const float dT = 1.0f;
	
	const int vectorSize = gridSizeX*gridSizeY;
	const int matrixSizeX = vectorSize;
	const int matrixSizeY = vectorSize;

	alpha = (dT*dT*C*C) / (2*h*h);
	float mainDiag = 4*alpha + 1;

	int i;

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

	// setup data main diagonal-Y
	for (i = 0; i < vectorSize; i++) 
		data1[i] = (i>=gridSizeX) ? -alpha : 0;

	// setup data main diagonal-1
	for (i = 0; i < vectorSize; i++) 
		data2[i] = (i%gridSizeY) ? -alpha : 0;

	// setup data main diagonal
	for (i = 0; i < vectorSize; i++)
		data3[i] = mainDiag;

	const int bandCount = 5;
	int bandOffsetDatas[] = {-gridSizeX, -1, 0, 1, gridSizeX};
	float* bandValues[] = {data1, data2, data3, data2, data1};

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

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

	gVar u = floor(step(0.f, vectorSize - 1.f, 1.f) % (float)gridSizeX);
	gVar v = floor(div(step(0.f, vectorSize - 1.f, 1.f), (float)gridSizeX));
	g1 = max(u - 1.f, 0.f) + v * (float)gridSizeX;
	g2 = min(u + 1.f, gridSizeX - 1.f) + v * (float)gridSizeX;
	g3 = u + max(v - 1.f, 0.f) * (float)gridSizeX;
	g4 = u + min(v + 1.f, gridSizeX - 1.f) * (float)gridSizeX;

	// setup initial solution, x = 0
	x = ones(vectorSize, 1)*0.0001f;

	// create drops by scattering a few random values into x
	int dropCount = int(gridSizeX  / 20.f);
	scatter(rand(dropCount, 1) * (float)(x.rowCount() - 1), rand(dropCount, 1)*0.5f, x);

	M = diag(A);

	s = gVar(gridSizeX, gridSizeY);
	o = zeros(x.rowCount(), x.columnCount());
}

//
// LibraSampleInerface.cpp
//

#include "LibraSample.h"
#include "WaterDropletsSample.h"

WaterDropletsSample* g_libraSample = 0;

////////////////////////////////////////////////////////////////////////////////////////////////
// Libra Sample exported functions
////////////////////////////////////////////////////////////////////////////////////////////////

LIBRA_API int libraSample_Init(int argc, char** argv)
{
	if (g_libraSample == 0)
		g_libraSample = new WaterDropletsSample();
	g_libraSample->init(argc, argv);

	return 0;
}

LIBRA_API void libraSample_Shutdown()
{
	if (g_libraSample)
	{
		g_libraSample->shutdown();
		delete g_libraSample;
		g_libraSample = 0;
	}
}


LIBRA_API void libraSample_GetRenderBufferBytes8888(int* renderBufferBytes)
{
	if (g_libraSample)
		g_libraSample->getRenderBufferBytes8888(renderBufferBytes);
}

LIBRA_API void libraSample_Tick()
{
	if (g_libraSample)
		g_libraSample->tick();
}