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

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

#include "Pattern.h"



/*



  uses fft fluid solver with gravity forcefield

  todo:

  use forward projection instead of sime-lagrangian lookups


*/





//---- 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 64
#define ARRAYSIZE (DIM * 2*(DIM/2+1))


 
 /* 
 * Simulation time step: you may need to change this to suit
 * your particular machine speed.
 *
 */
double dt = 2.0;  //  default 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;

  //  apply forces

  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];
  }  
 

  //  diffuse velocity

	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] = 0.0f;
    u_u0[i] = u_v0[i] = 0.0f;
    rho[i] = rho0[i] = 0;

    int x = i % DIM;
    int y = i / DIM;

    if( x> 20 && x < DIM-20 && y>20 && y<DIM-20 ) rho[i] = rho0[i] = 0.39;

	}
}







/*
 * 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 *= 1.f / len;
    dy *= 1.f / 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;
}




/*

*/











  double dist_lookup[ DIM*2 ][ DIM*2 ];
  double ood2_lookup[ DIM*2 ][ DIM*2 ];
  double dirx_lookup[ DIM*2 ][ DIM*2 ];
  double diry_lookup[ DIM*2 ][ DIM*2 ];




void calculate_gravity_field()
{
  int i, j, x, y;

  double px, py;
  double m1;

  double oodist2;

  int idx1=0,idx2=0;

  double scale_factor = 1.0/(DIM*DIM);

  for(j=0; j<DIM; j++) for(i=0; i<DIM; i++)
  {
	  px=0;
	  py=0;

    m1 = rho[idx1];
    idx2 = 0;

    if( m1 > 0.001 )
    {
      for( y=0; y<DIM; y++ )
      {
        for( x=0; x<DIM; x++ )
        {
          oodist2 = ood2_lookup[DIM+y-j][DIM+x-i];  //  one over the distance squared from (i,j) -> (x,y)
          oodist2 *= m1+rho[idx2++];

          px += dirx_lookup[DIM+y-j][DIM+x-i] * oodist2;
          py += diry_lookup[DIM+y-j][DIM+x-i] * oodist2;
        }
      }

      u0[ idx1 ] = -px*scale_factor / m1;
      v0[ idx1 ] = -py*scale_factor / m1;

    }
    else
    {
      u0[idx1] = v0[idx1] = 0;
    }

    u0[idx1] += u_u0[idx1];
    v0[idx1] += u_v0[idx1];

    u_u0[idx1]=0;
    u_v0[idx1]=0;
    
    idx1 ++;
  }




}



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

  calculate_gravity_field();
}







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

void Pattern_Startup()
{

	  int x, y;
	
	  //  calculate the distance from the centre point (SIZE,SIZE) in (SIZE*2,SIZE*2) array
	
	  double dx, dy, d2, d;

	  
	  for( y=0; y<DIM*2; y++ )
	  {
	    for( x=0; x<DIM*2; x++ )
	    {
	      if( x != DIM || y != DIM )
	      {
	        dx = (x-DIM);
	        dy = (y-DIM);
	        
	        d2 = dx*dx + dy*dy;
	
	        dist_lookup[ y ][ x ] = (d = sqrt( d2 ));  //  distance
	        ood2_lookup[ y ][ x ] = 1.f / d2;   //  one over distance squared, for gravitational force
	        dirx_lookup[ y ][ x ] = dx / d;     //  Unit direction vector
	        diry_lookup[ y ][ x ] = dy / d;
	      }
	      else  //  Special case for centre point
	      {
	        dist_lookup[ y ][ x ] = 0;
	        ood2_lookup[ y ][ x ] = 0;
	        dirx_lookup[ y ][ x ] = 0;
	        diry_lookup[ y ][ x ] = 0;
	      }
	    }
	  }



  
  
  init();
}



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


PATTERN_RGB *Pattern_Generate()
{
  //- Calculate forces

	setForces();

  int i, j, idx, c;

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

      c = (int)(u0[idx] * 12800.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 force

      c = (int)(v0[idx] * 12800.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;

    }

	}


  //- Update fluid
	stable_solve(DIM, u, v, u0, v0, 0.0, dt);
	diffuse_matter(DIM, u, v, rho, rho0, dt);

	t += dt;

  //- Copy matter into RGB buffer


	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 -----------------------------------------------------------//