// Usage: Drag with the mouse to add smoke to the fluid. This will also move a "rotor" that disturbs 
//        the velocity field at the mouse location. Press the indicated keys to change options
//-------------------------------------------------------------------------------------------------- 

#ifdef G_OS_WIN32
#define WIN32_LEAN_AND_MEAN 1
#include <windows.h>
#endif

#include <math.h>
#include <float.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <rfftw.h>
#include <GL/gl.h>
#include <GL/glu.h>

#include "funcs.h"
#include "fluids.h"
#include "smoke.h"
#include "colormap.h"
#include "glyphs.h"
#include "heightplots.h"
#include "streamlines.h"
#include "normals.h"
#include "isolines.h"
#include "flowvis.h"

//--- SIMULATION PARAMETERS ------------------------------------------------------------------------
int var_dims = 25;
double dt = 0.4;				        //simulation time step
float visc = 0.001f;				    //fluid viscosity
fftw_real *vx, *vy;             //(vx,vy)   = velocity field at the current moment
fftw_real *vx0, *vy0;           //(vx0,vy0) = velocity field at the previous moment
fftw_real *fx, *fy;	            //(fx,fy)   = user-controlled simulation forces, steered with the mouse 
fftw_real *rho, *rho0;		    	//smoke density at the current (rho) and previous (rho0) moment 
rfftwnd_plan plan_rc, plan_cr;  //simulation domain discretization

int   winWidth, winHeight;      //size of the graphics window, in pixels

static int fluids_calculate = TRUE;
static int DIM;             //size of simulation grid
static int fluids_hisdex = 0;

//------ SIMULATION CODE STARTS HERE -----------------------------------------------------------------

//init_simulation: Initialize simulation data structures as a function of the grid size 'n'. 
//                 Although the simulation takes place on a 2D grid, we allocate all data structures as 1D arrays,
//                 for compatibility with the FFTW numerical library.
void fluids_init_simulation(int n, struct vis_data_arrays *vis_data)				
{
	int i; size_t dim1, dim2, dim;

  DIM = n;
	dim     = n * 2 *(n / 2 + 1);
	dim1    = dim *sizeof(fftw_real);        //Allocate data structures
	vx      = (fftw_real*) malloc(dim1); 
	vy      = (fftw_real*) malloc(dim1);
	vx0     = (fftw_real*) malloc(dim1); 
	vy0     = (fftw_real*) malloc(dim1);
	dim2    = n * n * sizeof(fftw_real);
	fx      = (fftw_real*) malloc(dim2);
	fy      = (fftw_real*) malloc(dim2);
	rho     = (fftw_real*) malloc(dim2); 
	rho0    = (fftw_real*) malloc(dim2);
	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);

  for (i = 0; i < HISTORY_SIZE; i++)
  {
    vis_data->history_frame[i] = (fftw_real *)malloc(dim1);
    vis_data->history_scalars[i] = (struct fftw_real_xy *)malloc(sizeof(struct fftw_real_xy));
    vis_data->history_scalars[i]->x = (fftw_real *)malloc(dim1);
    vis_data->history_scalars[i]->y = (fftw_real *)malloc(dim1);
  }
  fluids_reset_simulation();

  vis_data->rho = (fftw_real *)malloc(dim2);
  vis_data->force = (fftw_real *)malloc(dim2);
  vis_data->vel = (fftw_real *)malloc(dim1);
  vis_data->scalars.x = (fftw_real *)malloc(dim1);
  vis_data->scalars.y = (fftw_real *)malloc(dim1);
  vis_data->div_force = (fftw_real *)malloc(dim2);
  vis_data->div_vel = (fftw_real *)malloc(dim1);
  vis_data->height = (fftw_real *)malloc(dim2);
  vis_data->normals = (struct point *)malloc(dim *sizeof(struct point));
}

void fluids_reset_simulation(void) {
  int i;

	for (i = 0; i < DIM * DIM; i++)                      //Initialize data structures to 0
	{ vx[i] = vy[i] = vx0[i] = vy0[i] = fx[i] = fy[i] = rho[i] = rho0[i] = 0.0f; }
}


//FFT: Execute the Fast Fourier Transform on the dataset 'vx'.
//     'dirfection' indicates if we do the direct (1) or inverse (-1) Fourier Transform
void FFT(int direction,void* vx)
{
	if(direction==1) rfftwnd_one_real_to_complex(plan_rc,(fftw_real*)vx,(fftw_complex*)vx);
	else             rfftwnd_one_complex_to_real(plan_cr,(fftw_complex*)vx,(fftw_real*)vx);
}

int clamp(float x) 
{
  return ((x)>=0.0?((int)(x)):(-((int)(1-(x)))));
}

//solve: Solve (compute) one step of the fluid flow simulation
void solve(int n, fftw_real* vx, fftw_real* vy, fftw_real* vx0, fftw_real* vy0, 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++) 
	{ vx[i] += dt*vx0[i]; vx0[i] = vx[i]; vy[i] += dt*vy0[i]; vy0[i] = vy[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*vx0[i+n*j])-0.5f; 
	      y0 = n*(y-dt*vy0[i+n*j])-0.5f;
	      i0 = clamp((float)x0); s = x0-i0;
	      i0 = (n+(i0%n))%n;
	      i1 = (i0+1)%n;
	      j0 = clamp((float)y0); t = y0-j0;
	      j0 = (n+(j0%n))%n;
	      j1 = (j0+1)%n;
	      vx[i+n*j] = (1-s)*((1-t)*vx0[i0+n*j0]+t*vx0[i0+n*j1])+s*((1-t)*vx0[i1+n*j0]+t*vx0[i1+n*j1]);
	      vy[i+n*j] = (1-s)*((1-t)*vy0[i0+n*j0]+t*vy0[i0+n*j1])+s*((1-t)*vy0[i1+n*j0]+t*vy0[i1+n*j1]);
	   }     
	
	for(i=0; i<n; i++)
	  for(j=0; j<n; j++) 
	  {  vx0[i+(n+2)*j] = vx[i+n*j]; vy0[i+(n+2)*j] = vy[i+n*j]; }

	FFT(1,vx0);
	FFT(1,vy0);

	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] = vx0[i  +(n+2)*j]; V[0] = vy0[i  +(n+2)*j];
	      U[1] = vx0[i+1+(n+2)*j]; V[1] = vy0[i+1+(n+2)*j];

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

	FFT(-1,vx0); 
	FFT(-1,vy0);

	f = 1.0/(n*n);
 	for (i=0;i<n;i++)
	   for (j=0;j<n;j++) 
	   { vx[i+n*j] = f*vx0[i+(n+2)*j]; vy[i+n*j] = f*vy0[i+(n+2)*j]; }
} 


// diffuse_matter: 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 *vx, fftw_real *vy, 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*vx[i+n*j])-0.5f; 
			y0 = n*(y-dt*vy[i+n*j])-0.5f;
			i0 = clamp((float)x0);
			s = x0-i0;
			i0 = (n+(i0%n))%n;
			i1 = (i0+1)%n;
			j0 = clamp((float)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]);
		}    
}

//set_forces: copy user-controlled forces to the force vectors that are sent to the solver. 
//            Also dampen forces and matter density to get a stable simulation.
void set_forces(void) 
{
	int i;
	for (i = 0; i < DIM * DIM; i++) 
	{
        rho0[i]  = 0.995 * rho[i];
        fx[i] *= 0.85; 
        fy[i] *= 0.85;
        vx0[i]    = fx[i]; 
        vy0[i]    = fy[i];
	}
}



struct point vector_normal(struct point vert1, struct point vert2, struct point vert3)
{
  struct point return_value;
  struct point v1, v2;
  float length;

  // vector v1
  v1.x = vert1.x - vert2.x;
  v1.y = vert1.y - vert2.y;
  v1.z = vert1.z - vert2.z;

  // vector v2
  v2.x = vert2.x - vert3.x;
	v2.y = vert2.y - vert3.y;
	v2.z = vert2.z - vert3.z;

  // calculate cross produkt
  return_value.x = v1.y * v2.z - v1.z * v2.y;
  return_value.y = v1.z * v2.x - v1.x * v2.z;
  return_value.z = v1.x * v2.y - v1.y * v2.x;

  length = vec_len3f(return_value.x, return_value.y, return_value.z);

  if(length == 0.0f)
  {
    length = 1.0f;
  }

  // normalize
  return_value.x /= length;
  return_value.y /= length;
  return_value.z /= length;


  return return_value;
}


struct point calculate_normal_vector(fftw_real *height, int index, int i, int j)
{
  float px, py;
  struct point p1, p2, p3;

  fftw_real wn = (fftw_real)winWidth  / (fftw_real)(DIM + 1); // Grid cell width
	fftw_real hn = (fftw_real)winHeight / (fftw_real)(DIM + 1); // Grid cell height

  px = wn + (fftw_real)i * wn;
  py = hn + (fftw_real)(j + 1) * hn;

  p1.x = px;
  p1.y = py;
  p1.z = height[index];

  p2.x = px;
  p2.y = py + hn;
  p2.z = height[index + DIM];

  p3.x = px + wn;
  p3.y = py;
  p3.z = height[index + 1];

  return vector_normal(p2, p1, p3);
}

fftw_real calculate_height_plot(int dataset, int index)
{
  fftw_real return_value;

  switch(dataset)
  {
    default:
    case DATASET_RHO:
      return_value = rho[index];
    break;
    case DATASET_VEL:
      return_value = vec_len2f(vx[index], vy[index]);
    break;
    case DATASET_FORCE:
      return_value = vec_len2f(fx[index], fy[index]);
    break;
    case DATASET_DIVV:
    case DATASET_DIVF:
      return_value = 0.0f;
    break;
  }

  return return_value;
}


void copy_frames(fftw_real *dataset)
{
  memcpy(smoke_get_frame(), dataset, DIM * 2 * (DIM / 2 + 1) * sizeof(fftw_real));
}


float par_der(float vp1, float v, fftw_real cell)
{
  return (vp1 - v / cell);
}

struct fftw_real_xy get_scalar_frames(struct vis_data_arrays *vis_data, int dataset)
{
  struct fftw_real_xy return_value;

  switch(dataset)
  {
    default:
    case DATASET_VEL:
      return_value.x = vx;
      return_value.y = vy;
    break;
    case DATASET_FORCE:
      return_value.x = fx;
      return_value.y = fy;
    break;
  }

  return return_value;
}

fftw_real *get_vector_frame(struct vis_data_arrays *vis_data, int dataset)
{
  fftw_real *return_value;

  switch(dataset)
  {
    default:
		case DATASET_RHO:
      return_value = vis_data->rho;
		break;
		case DATASET_VEL:
      return_value = vis_data->vel;
		break;
		case DATASET_FORCE:
      return_value = vis_data->force;
		break;
		case DATASET_DIVV:
      return_value = vis_data->div_vel;
		break;
		case DATASET_DIVF:
      return_value = vis_data->div_force;
		break;
    case DATASET_HIST:
//      return_value = vis_data->history_frame;
    break;
	}
  
  return return_value;
}


void add_history_frames(struct vis_data_arrays *vis_data, int dataset)
{
  fftw_real *frame;


  if (dataset != DATASET_HIST) {
    frame = get_vector_frame(vis_data, dataset);
    
    memcpy(vis_data->history_frame[fluids_hisdex], frame, DIM * 2 * (DIM / 2 + 1) * sizeof(fftw_real));
  }
}


void setup_arrays(struct vis_data_arrays *vis_data)
{
  int idx, i, j;
  float scale_min, scale_max;
  fftw_real *frame_smoke, *frame_streamlines;
  struct fftw_real_xy scalar_frames_glyphs, scalar_frames_streamlines;
  int dataset;
  
  fftw_real wn = (fftw_real)winWidth  / (fftw_real)(DIM + 1); // Grid cell width
	fftw_real hn = (fftw_real)winHeight / (fftw_real)(DIM + 1); // Grid cell height
  
  frame_smoke = get_vector_frame(vis_data, smoke_get_dataset());
  frame_streamlines = get_vector_frame(vis_data, streamlines_get_dataset());
  scale_min = scale_max = frame_smoke[0];

  dataset = glyphs_get_dataset_direction();
  scalar_frames_glyphs = get_scalar_frames(vis_data, dataset);
  dataset = streamlines_get_dataset();
  scalar_frames_streamlines = get_scalar_frames(vis_data, dataset);

	for (j = 0; j < DIM; j++)
	{
    for (i = 0; i < DIM; i++)
	  {
      idx = (j * DIM) + i;

      vis_data->rho[idx]       = rho[idx];
      vis_data->vel[idx]       = vec_len2f(vx[idx], vy[idx]);
      vis_data->force[idx]     = vec_len2f(fx[idx], fy[idx]);
      vis_data->scalars.x[idx] = scalar_frames_glyphs.x[idx];
      vis_data->scalars.y[idx] = scalar_frames_glyphs.y[idx];
      vis_data->div_vel[idx]   = par_der(vx[idx + 1], vx[idx], wn) + par_der(vy[idx + 1], vy[idx], wn);
      vis_data->div_force[idx] = par_der(fx[idx + 1], fx[idx], hn) + par_der(fy[idx + 1], fy[idx], hn);
      vis_data->history_frame[fluids_hisdex][idx] = frame_streamlines[idx];
      vis_data->history_scalars[fluids_hisdex]->x[idx] = scalar_frames_streamlines.x[idx];
      vis_data->history_scalars[fluids_hisdex]->y[idx] = scalar_frames_streamlines.y[idx];
      
      vis_data->height[idx]    = calculate_height_plot(heightplots_get_dataset(), idx);
      vis_data->normals[idx]   = calculate_normal_vector(vis_data->height, idx, i, j);

      if (frame_smoke[idx] < scale_min && frame_smoke[idx] > FLT_MIN) scale_min = frame_smoke[idx];
      if (frame_smoke[idx] > scale_max) scale_max = frame_smoke[idx];
    }
  }


  if (colormap_get_autoscaling())
  {
    colormap_set_scale_min(scale_min);
    colormap_set_scale_max(scale_max);
  }
}

//do_one_simulation_step: Do one complete cycle of the simulation:
//      - set_forces:       
//      - solve:            read forces from the user
//      - diffuse_matter:   compute a new set of velocities
//      - gluPostRedisplay: draw a new visualization frame


void fluids_calculate_one_simulation_step(struct vis_data_arrays *vis_data)
{
  int dataset;
  fftw_real *frame;

  if (fluids_calculate)
  {
      set_forces();
      solve(DIM, vx, vy, vx0, vy0, visc, dt);
      diffuse_matter(DIM, vx, vy, rho, rho0, dt);

      fluids_hisdex = (fluids_hisdex >= HISTORY_SIZE -1) ? 0 : fluids_hisdex +1;
  }
  setup_arrays(vis_data);

  dataset = smoke_get_dataset();
  frame = get_vector_frame(vis_data, dataset);
  smoke_set_frame(frame);

  heightplots_set_frame(vis_data->height);
  normals_set_frame(vis_data->normals);

  dataset = glyphs_get_dataset_color();
  frame = get_vector_frame(vis_data, dataset);
  glyphs_set_frame_color(frame);
  glyphs_set_frames_direction(vis_data->scalars);

  dataset = isolines_get_dataset();
  frame = get_vector_frame(vis_data, dataset);
  isolines_set_frame(frame);

  flowvis_set_history(vis_data->history_frame);

  streamlines_set_history_scalars(vis_data->history_scalars);
}

//------ VISUALIZATION CODE STARTS HERE -----------------------------------------------------------------
void fluids_insert_smoke(int x, int y)
{
    int xi,yi,X,Y; double  dx, dy, len;
		static int lx=0,ly=0;				//remembers last mouse location

		/* Density calculations etc are only computed on area's below the slider */

    // Compute the array index that corresponds to the cursor location 
    xi = (int)clamp((double)(DIM + 1) * ((double)x / (double)winWidth));
    yi = (int)clamp((double)(DIM + 1) * ((double)(winHeight - y) / (double)winHeight));

    X = xi; Y = yi;

    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 
    y = winHeight - y;
    dx = x - lx; dy = y - ly;
    len = vec_len2f(dx, dy);
    if (len != 0.0) {  dx *= 0.1 / len; dy *= 0.1 / len; }
    fx[Y * DIM + X] += dx; 
    fy[Y * DIM + X] += dy;
    rho[Y * DIM + X] = 10.0f;
    lx = x; ly = y;
}

// getters and setters

int fluids_get_var_dim(void)
{
  return var_dims;
}

void fluids_set_var_dim(int dims)
{
  var_dims = dims;
}

void fluids_set_calculate(int calculate)
{
  fluids_calculate = calculate;
}

int fluids_get_calculate(void)
{
  return fluids_calculate;
}

int fluids_get_dim(void)
{
  return DIM;
}

int fluids_get_hisdex(void)
{
  return fluids_hisdex;
}