#include <windows.h>
#include <math.h>

#include <math.h>
#include <rfftw.h>

#include "Pattern.h"





//---- Render buffer

PATTERN_RGB buf;






//---- GLOBAL DATA -----------------------------------------------------------//

/*
 * Size of the simulation grid, i.e. the number of grid
 * points are (DIM x DIM).
 *
 */
#define DIM 128
#define ARRAYSIZE (DIM * 2*(DIM/2+1))


 
 /* 
 * Simulation time step: you may need to change this to suit
 * your particular machine speed.
 *
 */
double dt = 0.4;




int    winWidth, winHeight;
int    drawVelocities = 0;
int    wireframe = 0;
int    lmx, lmy;



double t = 0;



/* 
 * See Jos Stam, "A Simple Fluid Solver Based on the FFT",
 * http://www.dgp.utoronto.ca/people/stam/reality/Research/pub.html
 * for more details on the following code.
 *
 */

fftw_real u[DIM * 2*(DIM/2+1)];
fftw_real v[DIM * 2*(DIM/2+1)];

fftw_real u0[DIM * 2*(DIM/2+1)];
fftw_real v0[DIM * 2*(DIM/2+1)];

fftw_real rho[DIM * DIM];
fftw_real rho0[DIM * DIM];  /* Smoke density */

fftw_real u_u0[DIM * DIM];
fftw_real u_v0[DIM * DIM];  /* User-induced forces */

static rfftwnd_plan plan_rc, plan_cr;







//---- init_FFT

//  Sets up the FFT engine

void init_FFT(int n)
{
	plan_rc = rfftw2d_create_plan(n, n, FFTW_REAL_TO_COMPLEX, FFTW_IN_PLACE);
	plan_cr = rfftw2d_create_plan(n, n, FFTW_COMPLEX_TO_REAL, FFTW_IN_PLACE);
}

#define FFT(s,u)\
	if(s==1) rfftwnd_one_real_to_complex(plan_rc,(fftw_real *)u,(fftw_complex*)u);\
	else rfftwnd_one_complex_to_real(plan_cr,(fftw_complex *)u,(fftw_real *)u)

#define floor(x) ((x)>=0.0?((int)(x)):(-((int)(1-(x)))))






void stable_solve ( int n, fftw_real * u, fftw_real * v, fftw_real * u0, fftw_real * v0, fftw_real visc, fftw_real dt )
{
	fftw_real x, y, x0, y0, f, r, U[2], V[2], s, t;
	int i, j, i0, j0, i1, j1;

	for ( i=0 ; i<n*n ; i++ ) {
		u[i] += dt*u0[i]; 
		u0[i] = u[i];
		
		v[i] += dt*v0[i]; 
		v0[i] = v[i];
  }    

	for ( x=0.5f/n,i=0 ; i<n ; i++,x+=1.0f/n ) {
		for ( y=0.5f/n,j=0 ; j<n ; j++,y+=1.0f/n ) {
			x0 = n*(x-dt*u0[i+n*j])-0.5f; 
			y0 = n*(y-dt*v0[i+n*j])-0.5f;
			i0 = floor(x0);
			s = x0-i0;
			i0 = (n+(i0%n))%n;
			i1 = (i0+1)%n;
			j0 = floor(y0);
			t = y0-j0;
			j0 = (n+(j0%n))%n;
			j1 = (j0+1)%n;
			u[i+n*j] = (1-s)*((1-t)*u0[i0+n*j0]+t*u0[i0+n*j1])+                        
				s *((1-t)*u0[i1+n*j0]+t*u0[i1+n*j1]);
			v[i+n*j] = (1-s)*((1-t)*v0[i0+n*j0]+t*v0[i0+n*j1])+
				s *((1-t)*v0[i1+n*j0]+t*v0[i1+n*j1]);
		}    
	} 
	
	for ( i=0 ; i<n ; i++ )
		for ( j=0 ; j<n ; j++ ) { 
			u0[i+(n+2)*j] = u[i+n*j]; 
		  v0[i+(n+2)*j] = v[i+n*j];
		}

	FFT(1,u0);
	FFT(1,v0);

  /*
    Temp: display the FFT

  */
/*
  int idx;
  int c;

	for (j = 0; j < DIM - 1; j++)
  {
		for (i = 0; i < DIM - 1; i++)
    {
  		idx = ((j+n/2)%n) * (DIM+2) + ((i+n/2)%n);
						
      // x velocity

      c = (int)(u0[idx] * 128.0 + 128.0);

      if( c < 0 ) c = 0;
      if( c > 255 ) c = 255;

			buf.buf[j+DIM*2][i][0] = c;
			buf.buf[j+DIM*2][i][1] = c;
			buf.buf[j+DIM*2][i][2] = c;

      //  y velocity

      c = (int)(v0[idx] * 128.0 + 128.0);

      if( c < 0 ) c = 0;
      if( c > 255 ) c = 255;

			buf.buf[j+DIM*2][i+DIM][0] = c;
			buf.buf[j+DIM*2][i+DIM][1] = c;
			buf.buf[j+DIM*2][i+DIM][2] = c;

    }

	}
*/

	for ( i=0 ; i<=n ; i+=2 ) {
		x = 0.5f*i;
		for ( j=0 ; j<n ; j++ ) {
			y = j<=n/2 ? (fftw_real)j : (fftw_real)j-n;
			r = x*x+y*y;
			if ( r==0.0f ) continue;
			f = (fftw_real)exp(-r*dt*visc);
			U[0] = u0[i  +(n+2)*j]; V[0] = v0[i  +(n+2)*j];
			U[1] = u0[i+1+(n+2)*j]; V[1] = v0[i+1+(n+2)*j];

			u0[i  +(n+2)*j] = f*( (1-x*x/r)*U[0]     -x*y/r *V[0] );
			u0[i+1+(n+2)*j] = f*( (1-x*x/r)*U[1]     -x*y/r *V[1] );
			v0[i+  (n+2)*j] = f*(   -y*x/r *U[0] + (1-y*y/r)*V[0] );
			v0[i+1+(n+2)*j] = f*(   -y*x/r *U[1] + (1-y*y/r)*V[1] );
		}    
	}

	FFT(-1,u0); 
	FFT(-1,v0);

	f = 1.0/(n*n);
 	for ( i=0 ; i<n ; i++ )
		for ( j=0 ; j<n ; j++ ) {
			u[i+n*j] = f*u0[i+(n+2)*j]; 
			v[i+n*j] = f*v0[i+(n+2)*j]; 
		}
} 


/*
 * This function diffuses matter that has been placed
 * in the velocity field. It's almost identical to the
 * velocity diffusion step in the function above. The
 * input matter densities are in rho0 and the result
 * is written into rho.
 *
 */
void diffuse_matter(int n, fftw_real *u, fftw_real *v, fftw_real *rho, fftw_real *rho0, fftw_real dt) {
	fftw_real x, y, x0, y0, s, t;
	int i, j, i0, j0, i1, j1;

	for ( x=0.5f/n,i=0 ; i<n ; i++,x+=1.0f/n ) {
		for ( y=0.5f/n,j=0 ; j<n ; j++,y+=1.0f/n ) {
			x0 = n*(x-dt*u[i+n*j])-0.5f; 
			y0 = n*(y-dt*v[i+n*j])-0.5f;
			i0 = floor(x0);
			s = x0-i0;
			i0 = (n+(i0%n))%n;
			i1 = (i0+1)%n;
			j0 = floor(y0);
			t = y0-j0;
			j0 = (n+(j0%n))%n;
			j1 = (j0+1)%n;
			rho[i+n*j] = (1-s)*((1-t)*rho0[i0+n*j0]+t*rho0[i0+n*j1])+                        
				s *((1-t)*rho0[i1+n*j0]+t*rho0[i1+n*j1]);
		}    
	} 
}
















void keyboard(unsigned char key, int x, int y) {
	switch (key) {
		case 'v':
			drawVelocities = 1 - drawVelocities;
			break;
		case 'w':
			wireframe = 1 - wireframe;
			break;
		case 't':
			dt -= 0.001;
			printf("dt = %f\n", dt);
			break;
		case 'T':
			dt += 0.001;
			printf("dt = %f\n", dt);
			break;
	}
}


void init(void) {
	int i;

	init_FFT(DIM);

	for (i = 0; i < DIM * DIM; i++) {
		u[i] = v[i] = u0[i] = v0[i] = rho[i] = rho0[i] = u_u0[i] = u_v0[i] = 0.0f;
	}
}


/*
 * When the user drags with the mouse, add a force
 * that corresponds to the direction of the mouse
 * cursor movement. Also inject som new matter into
 * the field at the mouse location.
 *
 */
void drag(int mx, int my)
{
	double  dx, dy, len;
	int     X, Y;

	/* Compute the array index that corresponds to the
	 * cursor location */

	X = mx;
	Y = my;

	if (X > (DIM - 1)) {
		X = DIM - 1;
	}
	if (Y > (DIM - 1)) {
		Y = DIM - 1;
	}
	if (X < 0) {
		X = 0;
	}
	if (Y < 0) {
		Y = 0;
	}

	/* Add force at the cursor location */

//	my = winHeight - my;
	dx = mx - lmx;
	dy = my - lmy;
	len = sqrt(dx * dx + dy * dy);
  if (len != 0.0) { 
    dx *= 0.1 / len;
    dy *= 0.1 / len;
  }
	u_u0[Y * DIM + X] += dx;
	u_v0[Y * DIM + X] += dy;

	/* Increase matter density at the cursor location */

	rho[Y * DIM + X] = 10.0f;

	lmx = mx;
	lmy = my;
}




/*

*/

void setForces(void)
{
	int i, j;
	for (i = 0; i < DIM * DIM; i++)
  {
    rho0[i] = rho[i]*0.995;

		u_u0[i] *= 0.85;
		u_v0[i] *= 0.85;

		u0[i] = u_u0[i];
		v0[i] = u_v0[i];
	}
}




















//---- Pattern_Startup -------------------------------------------------------//

void Pattern_Startup()
{
  init();
}



//---- Pattern_Generate ------------------------------------------------------//


PATTERN_RGB *Pattern_Generate()
{
  //- Update fluid

	setForces();
	stable_solve(DIM, u, v, u0, v0, 0.001, dt);
	diffuse_matter(DIM, u, v, rho, rho0, dt);

	t += dt;

  //- Copy matter into RGB buffer


  int i, j, idx, c;

	for (j = 0; j < DIM - 1; j++)
  {
		idx = j * DIM;
						
		for (i = 0; i < DIM - 1; i++, idx++)
    {
      //  mass

      c = (int)(rho0[idx] * 255);

      if( c < 0 ) c = 0;
      if( c > 255 ) c = 255;

			buf.buf[j][i][0] = c;
			buf.buf[j][i][1] = c;
			buf.buf[j][i][2] = c;

      // x velocity

      c = (int)(u[idx] * 12800.0 + 128.0);

      if( c < 0 ) c = 0;
      if( c > 255 ) c = 255;

			buf.buf[j+DIM][i][0] = c;
			buf.buf[j+DIM][i][1] = c;
			buf.buf[j+DIM][i][2] = c;

      //  y velocity

      c = (int)(v[idx] * 12800.0 + 128.0);

      if( c < 0 ) c = 0;
      if( c > 255 ) c = 255;

			buf.buf[j+DIM][i+DIM][0] = c;
			buf.buf[j+DIM][i+DIM][1] = c;
			buf.buf[j+DIM][i+DIM][2] = c;

    }

	}

  //- Return buffer for rendering

  return &buf;
}








//---- END OF FILE -----------------------------------------------------------//