//
// WaterDroplets.java
//

import java.applet.Applet;
import java.awt.*;

public class WaterDroplets extends Applet
{
    private gVar p1 = null;private gVar p2 = null;private gVar p3 = null;
    private gVar p4 = null;private gVar p5 = null;private gVar p6 = null;  

    static final long serialVersionUID = 989743789;

    public void init() 
    { 
        String[] args = {""};
	main(args);
    }
        
    public void stop() 
    { 
    }
    
    public void paint(Graphics g) 
    {
	String udir = System.getProperty("user.dir");
	g.drawString(udir, 20,50);
	g.drawString("Hello, world!", 20,10);
    }
	
    static public void setCmdLine(int argc, String[] argv, int customArgc, String[] customArgv)
    {
        int i=0;
        for (;i<argc;i++)
        {
	        customArgv[i] = argv[i];
	    }

        // Set -c ClientSDK, -? for help.
        customArgv[i] = "-c";
        customArgv[i+1] = "ClientSDK";
    }

    public static void main(String[] args)
    {
    	Thread.currentThread().setPriority(Thread.MAX_PRIORITY);
    	
        System.loadLibrary("jv");
    	int customArgc = args.length+2;
        String[] customArgv = new String[customArgc];

        setCmdLine(args.length, args, customArgc, customArgv);

        System.out.print("*************************\n\n");
        System.out.print("Water droplets simulation\n\n");
        System.out.print("*************************\n\n");

        if (libra.Init(customArgc, customArgv) != 0)
	        return;

        //libra.SetCurrentBackend(gBackend.CPU_BACKEND);

        new WaterDroplets().run();
        
        libra.Shutdown();     
    }

    public void run()
    {
        // 
        libra.SetDefaultDataType(gDataType.GFLOAT32);

        // Setup data.
        int gridSizeX = 512;
        int gridSizeY = 512;
        float[] alpha = new float[1];

        gVar A = new gVar(), x = new gVar(), b = new gVar();
        gVar g1 = new gVar(), g2 = new gVar(), g3 = new gVar(), g4 = new gVar();

        setupWaveEquationSystem(alpha, A, x, b,
                                g1, g2, g3, g4,
						        gridSizeX, gridSizeY);
   
        // Solve system
        solveRenderWaveEquation(alpha[0], A, x, b,
                                g1, g2, g3, g4,
						        gridSizeX, gridSizeY);
    }

    public void solveRenderWaveEquation(float alpha, gVar A, gVar x, gVar b,
						        gVar g1, gVar g2, gVar g3, gVar g4,
						         int gridSizeX, int gridSizeY)
    {
    	p1 = new gVar();p2 = new gVar();p3 = new gVar();
    	p4 = new gVar();p5 = new gVar();p6 = new gVar();  

//        gVar M = libra.diag(A);
        gVar u = libra.zeros(x.rowCount(), x.columnCount());
        b.assign(libra.zeros(x.rowCount(), x.columnCount()));
        int rowCount = x.rowCount();
        double totalTime = 0.0, averageIterationTime;
        int timeSteps = 1000;
        double startTime = libra.GetTime();
        
        // for each step in time...
        for (int i=0;i<timeSteps;++i)
        { 
	        if (i == 10)
		        startTime = libra.GetTime();        
	        
	        // Create new drops by scattering a random value
	        // into x at the random position
	        // Scatter heavy rain
	        p1.assign(libra.rand(15, 1));
	        p2.assign(libra.mul(p1,0.2f));
	        p3.assign(libra.sub(p2, 0.1f));
	        p4.assign(libra.rand(15, 1));
	        p5.assign(libra.mul(p4, (float)(rowCount - 1.0f)));
	        
	        libra.scatter(p5, p3, x);
	
	        // Setup right hand side (b) of equation	          	        
	        p6.assign(libra.gather(g1, x));
	        p5.assign(libra.gather(g2, x));	        
	        p1.assign(libra.add(p6, p5));
	        p6.assign(libra.gather(g3, x));
	        p2.assign(libra.add(p1, p6));
	        p6.assign(libra.gather(g4, x));
	        p3.assign(libra.add(p2, p6));   
	        p4.assign(libra.mul(alpha, p3));
	        p5.assign(libra.mul(x, 2.0f - 4.0f * alpha));	        
	        p6.assign(libra.add(p4, p5));	        
	        b.assign(libra.sub(p6, u));
	        
	        // u <-- old solution
	        u.assign(x);      
            
	        // x <-- solve for new solution
	        for (int j=0;j<3;++j)
		        x.assign(libra.jacobiIter(A, x, b));

	        // render solution x to screen
	    	// Rescale and render solution to screen
	        p1.assign(libra.add(u, 1.0f));
	        p2.assign(libra.log10(p1));
	        p3.assign(libra.mul(p2, 15.0f));

	        //
	        libra.render(p3, gRenderMode.GHEIGHTFIELD2D, gridSizeX);        
	        
	        // Collect variables at end of each iteration
        	System.gc();
        	System.runFinalization();
        }

        totalTime = libra.GetTime() - startTime;
        averageIterationTime = totalTime / (timeSteps-10);

        System.out.print("Total time (s) : " + totalTime + "\n");
        System.out.print("Average iteration time (ms) : " + averageIterationTime * 1000.0 + "\n");
        System.out.print("\n");
    }

    public void setupWaveEquationSystem(float[] alpha, gVar A, gVar x, gVar b,
							        gVar g1, gVar g2, gVar g3, gVar g4,
							        int gridSizeX, int gridSizeY)
    {
        //
        // 2D Wave equation
        // 

        float C = 0.5F;
        float h = 1.0F;
        float dT = 1.0F;
    	
        int vectorSize = gridSizeX*gridSizeY;
        int matrixSizeX = vectorSize;
        int matrixSizeY = vectorSize;

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

        int i, offset = 0;
        int bandCount = 5;
     
        float[] bandValues = new float[vectorSize*bandCount];
      
        // setup data main diagonal-Y
        for (i = 0; i < vectorSize; i++)
            bandValues[offset+i] = (i >= gridSizeX) ? -alpha[0] : 0;

        offset += i;

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

        offset += i;

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

        offset += i;

        for (i = 0; i < vectorSize; i++)
            bandValues[offset + i] = (i >= gridSizeX) ? -alpha[0] : 0;

        offset += i;

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

        offset += i;

        int[] bandOffsetDatas = {-gridSizeX, -1, 0, 1, gridSizeX};
        
        // create the sparse matrix with input on band form
        A.assign(libra.sparseBanded(matrixSizeX, matrixSizeY, bandValues, bandOffsetDatas, bandCount));

        gVar u = libra.floor(libra.mod(libra.step(0.0f, vectorSize - 1.0f, 1.0f), (float)gridSizeX));      
        gVar v = libra.floor(libra.div(libra.step(0.0f, vectorSize - 1.0f, 1.0f), (float)gridSizeX));
        g1.assign(libra.add(libra.max(libra.sub(u, 1.0f), 0.0f), libra.mul(v, (float)gridSizeX)));
        g2.assign(libra.add(libra.min(libra.add(u, 1.0f), gridSizeX - 1.0f), libra.mul(v, (float)gridSizeX)));
        g3.assign(libra.add(u, libra.mul(libra.max(libra.sub(v, 1.0f), 0.0f), (float)gridSizeX)));
        g4.assign(libra.add(u, libra.mul(libra.min(libra.add(v, 1.0f), gridSizeX - 1.0f), (float)gridSizeX)));

        // setup initial solution, x = 0
        x.assign(libra.zeros(vectorSize, 1));

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