#include "grid.h"
#include "argparser.h"
#include "boundingbox.h"
#include "vectors.h"
#include "matrix.h"

// Included files for OpenGL Rendering
#include <GL/gl.h>
#include <GL/glut.h>

#define BETA_0 1.7
#define EPSILON 0.0001

void CubePaint();

// ==============================================================
// ==============================================================
// ==============================================================

Grid::Grid(ArgParser *_args) {
  args = _args;
  dx = dy = dz = 1;
  dt = args->timestep;
  viscosity = args->viscosity;
  density = args->density;
  gravity = args->gravity * Vec3f(0,-1,0);

  if (args->input_file != NULL) Load();
  else Initialize();

  SetEmptySurfaceFull();
}

// ==============================================================

void Grid::Initialize() {
  assert (args->input_file == NULL);
  nx = args->nx;
  ny = args->ny;
  nz = args->nz;
  cells = new Cell[nx*ny*nz];

  // *********************************************************************  
  // ASSIGNMENT:
  //
  // You may want to uncomment this and play with random initial velocities.
  //
  /*
  int i,j,k;
  for (i = 0; i < nx; i++) {
    for (j = 0; j < ny; j++) {
      for (k = 0; k < nz; k++) {
	getCell(i,j,k)->set_u_plus(2*drand48()-1);
	getCell(i,j,k)->set_v_plus(2*drand48()-1);
	getCell(i,j,k)->set_w_plus(2*drand48()-1);
      }
    }
  }
  */
  // *********************************************************************

  // initialize particles on a third of the space...
  for (int n = 0; n < nx*ny*nz*density/3.0; n++) {
    Vec3f pos = Vec3f(drand48()*nx/3.0,drand48()*ny,drand48()*nz);
    int i,j,k;
    Cell *cell = getCell(int(pos.x()),int(pos.y()),int(pos.z()));
    Particle *p = new Particle();
    p->setPosition(pos);
    cell->addParticle(p);
  }  

}

// ==============================================================

void Grid::Load() {    

  // load some velocities from a file
  assert (args->input_file != NULL);

  fscanf (args->input_file,"%d %d %d\n",&nx,&ny,&nz);
  assert (nx > 0 && ny > 0 && nz > 0);
  cells = new Cell[nx*ny*nz];
  char c;
  int i,j,k;
  float f;
  while(fscanf (args->input_file,"%c %d %d %d %f\n", &c,&i,&j,&k,&f) > 0) {
    assert(i >= 0 && i < nx);
    assert(j >= 0 && j < ny);
    assert(k >= 0 && k < nz);
    if      (c == 'u') getCell(i,j,k)->set_u_plus(f);
    else if (c == 'v') getCell(i,j,k)->set_v_plus(f);
    else if (c == 'w') getCell(i,j,k)->set_w_plus(f);
    else assert(0);
  }
  
  // initialize particles uniformly at half density
  for (int n = 0; n < nx*ny*nz*density/2.0; n++) {
    Vec3f pos = Vec3f(drand48()*nx,drand48()*ny,drand48()*nz);
    int i,j,k;
    Cell *cell = getCell(int(pos.x()),int(pos.y()),int(pos.z()));
    Particle *p = new Particle();
    p->setPosition(pos);
    cell->addParticle(p);
  }

  // *********************************************************************  
  // ASSIGNMENT:
  //
  // You may want to expand the input format to specify initial particle densities
  //
  // *********************************************************************  

}

// ==============================================================

Grid::~Grid() {
  delete [] cells;
}

// ==============================================================
// ==============================================================
// ==============================================================

void Grid::Animate() {

  // the animation manager:  this is what gets done each timestep!

  ComputeNewVelocities();
  
  int iters = 20;
  float divergence = 1;
  while (iters > 0 && divergence > EPSILON) {
    iters--;
    divergence = AdjustForIncompressibility();
  }

  CopyVelocities();

  MoveParticles();

  ReassignParticles();

  SetEmptySurfaceFull();
}

// ==============================================================
// ==============================================================
// ==============================================================

void Grid::ComputeNewVelocities() {

  for (int i = 0; i < nx; i++) {
    for (int j = 0; j < ny; j++) {
      for (int k = 0; k < nz; k++) {
	
	// *********************************************************************  
	// ASSIGNMENT:
	//
	// This is a hack (at least I didn't find this in any of the implementation details)
	// But I don't think empty air should have so much gravity acting on it?
	// See if you can do the right thing.
	// 
	int not_empty = getCell(i,j,k)->getStatus() != CELL_EMPTY;
	//
	// *********************************************************************  

	float new_u =
	  get_u_plus(i,j,k) +
	  dt * ( (1/dx) * (square(get_u_avg(i,j,k)) - square(get_u_avg(i+1,j,k))) +
		 (1/dy) * (get_uv_plus(i,j-1,k) - get_uv_plus(i,j,k)) +
		 (1/dz) * (get_uw_plus(i,j,k-1) - get_uw_plus(i,j,k)) +
		 gravity.x()*not_empty +
		 (1/dx) * (getPressure(i,j,k)-getPressure(i+1,j,k)) +
		 (viscosity/square(dx)) * (get_u_plus(i+1,j  ,k  ) - 2*get_u_plus(i,j,k) + get_u_plus(i-1,j  ,k  )) +
		 (viscosity/square(dy)) * (get_u_plus(i  ,j+1,k  ) - 2*get_u_plus(i,j,k) + get_u_plus(i  ,j-1,k  )) +
		 (viscosity/square(dz)) * (get_u_plus(i  ,j  ,k+1) - 2*get_u_plus(i,j,k) + get_u_plus(i  ,j  ,k-1)) );
	
	float new_v =
	  get_v_plus(i,j,k) +
	  dt * ( (1/dx) * (get_uv_plus(i-1,j,k) - get_uv_plus(i,j,k)) +
		 (1/dy) * (square(get_v_avg(i,j,k)) - square(get_v_avg(i,j+1,k))) +
		 (1/dz) * (get_vw_plus(i,j,k-1) - get_vw_plus(i,j,k)) +
		 gravity.y()*not_empty  +
		 (1/dy) * (getPressure(i,j,k)-getPressure(i,j+1,k)) +
		 (viscosity/square(dx)) * (get_v_plus(i+1,j  ,k  ) - 2*get_v_plus(i,j,k) + get_v_plus(i-1,j  ,k  )) +
		 (viscosity/square(dy)) * (get_v_plus(i  ,j+1,k  ) - 2*get_v_plus(i,j,k) + get_v_plus(i  ,j-1,k  )) +
		 (viscosity/square(dz)) * (get_v_plus(i  ,j  ,k+1) - 2*get_v_plus(i,j,k) + get_v_plus(i  ,j  ,k-1)) );
	
	float new_w =
	  get_w_plus(i,j,k) +
	  dt * ( (1/dx) * (get_uw_plus(i-1,j,k) - get_uw_plus(i,j,k)) +
		 (1/dy) * (get_vw_plus(i,j-1,k) - get_vw_plus(i,j,k)) +
		 (1/dz) * (square(get_w_avg(i,j,k)) - square(get_w_avg(i,j,k+1))) +
		 gravity.z()*not_empty  + 
		 (1/dz) * (getPressure(i,j,k)-getPressure(i,j,k+1)) +
		 (viscosity/square(dx)) * (get_w_plus(i+1,j  ,k  ) - 2*get_w_plus(i,j,k) + get_w_plus(i-1,j  ,k  )) +
		 (viscosity/square(dy)) * (get_w_plus(i  ,j+1,k  ) - 2*get_w_plus(i,j,k) + get_w_plus(i  ,j-1,k  )) +
		 (viscosity/square(dz)) * (get_w_plus(i  ,j  ,k+1) - 2*get_w_plus(i,j,k) + get_w_plus(i  ,j  ,k-1)) );

	Cell *cell = getCell(i,j,k);
	cell->set_new_u_plus(new_u);
	cell->set_new_v_plus(new_v);
	cell->set_new_w_plus(new_w);
      }
    }
  }
}

// ==============================================================
// ==============================================================

float Grid::AdjustForIncompressibility() {

  // *********************************************************************  
  // ASSIGNMENT:
  //
  // This is not a complete implementation of the Marker and Cell (MAC) method.
  // Additional boundary velocities should be equalized as described in the references
  // depending on whether the boundaries are free-slip or no-slip.
  //
  // Also play around with compressible flow!
  //
  // *********************************************************************    

  float answer = 0;
  for (int i = 0; i < nx; i++) {
    for (int j = 0; j < ny; j++) {
      for (int k = 0; k < nz; k++) {

	Cell *cell = getCell(i,j,k);

	float diff = 
	  - ( (1/dx) * (get_new_u_plus(i,j,k) - get_new_u_plus(i-1,j,k)) +
	      (1/dy) * (get_new_v_plus(i,j,k) - get_new_v_plus(i,j-1,k)) +
	      (1/dz) * (get_new_w_plus(i,j,k) - get_new_w_plus(i,j,k-1)) );
	
	float beta = BETA_0/((2*dt) * (1/square(dx) + 1/square(dy) + 1/square(dz)));
	float dp = beta*diff;
	
	cell->setPressure(cell->getPressure()+beta*diff);
  
	if (!u_boundary(i,j,k)) adjust_new_u_plus(i,j,k,dt/dx*dp);
	if (!v_boundary(i,j,k)) adjust_new_v_plus(i,j,k,dt/dy*dp);
	if (!w_boundary(i,j,k)) adjust_new_w_plus(i,j,k,dt/dz*dp);
	if (!u_boundary(i-1,j,k)) adjust_new_u_plus(i-1,j,k,-dt/dx*dp);
	if (!v_boundary(i,j-1,k)) adjust_new_v_plus(i,j-1,k,-dt/dx*dp);
	if (!w_boundary(i,j,k-1)) adjust_new_w_plus(i,j,k-1,-dt/dx*dp);
  
	answer = max2(answer,fabs(diff));
      }
    }
  }
  return answer;
}

// ==============================================================
// ==============================================================

void Grid::CopyVelocities() {
  for (int i = 0; i < nx; i++) {
    for (int j = 0; j < ny; j++) {
      for (int k = 0; k < nz; k++) {
	getCell(i,j,k)->copyVelocity();
      }
    }
  }
}

// ==============================================================
// ==============================================================

void Grid::MoveParticles() {
  for (int i = 0; i < nx; i++) {
    for (int j = 0; j < ny; j++) {
      for (int k = 0; k < nz; k++) {
	
	Cell *cell = getCell(i,j,k);
	Iterator<Particle*> *iter = cell->getParticles()->StartIteration();
	while (Particle *p = iter->GetNext()) {
	  Vec3f pos = p->getPosition();

	  // *********************************************************************  
	  // ASSIGNMENT:
	  //
	  // I've intentionally reverted to the "dumb" velocity interpolation.
	  // Do it right, as described in the papers.
	  //
	  Vec3f vel = Vec3f(get_u_avg(i,j,k),get_v_avg(i,j,k),get_w_avg(i,j,k));
	  //
	  // *********************************************************************  

	  pos += vel * dt;
	  p->setPosition(pos);
	}
	cell->getParticles()->EndIteration(iter);
      }
    }
  }
}

void Grid::ReassignParticles() {
  for (int i = 0; i < nx; i++) {
    for (int j = 0; j < ny; j++) {
      for (int k = 0; k < nz; k++) {
	Cell *cell = getCell(i,j,k);
	Iterator<Particle*> *iter = cell->getParticles()->StartIteration();
	while (Particle *p = iter->GetNext()) {
	  Vec3f pos = p->getPosition();
	  int x = min2(nx-1,max2(0,(int)floor(pos.x())));
	  int y = min2(ny-1,max2(0,(int)floor(pos.y())));
	  int z = min2(nz-1,max2(0,(int)floor(pos.z())));
	  if (i != x || j != y || k != z) {
	    cell->removeParticle(p);
	    getCell(x,y,z)->addParticle(p);
	  }
	}
	cell->getParticles()->EndIteration(iter);
      }
    }
  }
}

void Grid::SetEmptySurfaceFull() {
  int i,j,k;
  for (i = 0; i < nx; i++) {
    for (j = 0; j < ny; j++) {
      for (k = 0; k < nz; k++) {
        Cell *cell = getCell(i,j,k);
        if (cell->numParticles() == 0)
	  cell->setStatus(CELL_EMPTY);
        else 
          cell->setStatus(CELL_FULL);
      }
    }
  }

  // pick out the boundary cells
  for (i = 0; i < nx; i++) {
    for (j = 0; j < ny; j++) {
      for (k = 0; k < nz; k++) {
        Cell *cell = getCell(i,j,k);
        if (cell->getStatus() == CELL_FULL &&
            (isEmpty(i-1,j,k) || isEmpty(i+1,j,k) ||
             isEmpty(i,j-1,k) || isEmpty(i,j+1,k) ||
             isEmpty(i,j,k-1) || isEmpty(i,j,k+1)))
          cell->setStatus(CELL_SURFACE);
      }
    }
  }
}

// ==============================================================
// ==============================================================

void Grid::Paint() const {
  glMatrixMode(GL_MODELVIEW);

  // center the volume in the window
  Matrix m;
  m.SetToIdentity();
  m *= Matrix::MakeTranslation(Vec3f(-0.5,-0.5,-0.5));
  int max = max3(nx,ny,nz);
  m *= Matrix::MakeTranslation(Vec3f((max-nx)/float(2*max),(max-ny)/float(2*max),(max-nz)/float(2*max)));
  m *= Matrix::MakeScale(1/float(max3(nx,ny,nz)));
  glMultMatrixf(m.glGet());

  // paint the boundaries of the volume
  BoundingBox box(Vec3f(0,0,0), Vec3f(nx,ny,nz));
  box.Paint();

  // =====================================================================================
  // render the markers/particles
  // =====================================================================================
  if (args->markers) {
    glDisable(GL_LIGHTING);
    glColor3f(0,0,0);
    glPointSize(3);
    glBegin(GL_POINTS);
    for (int x = 0; x < nx; x++) {
      for (int y = 0; y < ny; y++) {
	for (int z = 0; z < nz; z++) {
	  Cell *cell = getCell(x,y,z);
	  Iterator<Particle*> *iter = cell->getParticles()->StartIteration();
	  while (Particle *p = iter->GetNext()) {
	    Vec3f v = p->getPosition();
	    glVertex3f(v.x(),v.y(),v.z());
	  }
	  cell->getParticles()->EndIteration(iter);
  	} 
      }
    }
    glEnd();
    glEnable(GL_LIGHTING);
  } 

  // =====================================================================================
  // render stubby triangles to visualize the u, v, and w velocities between cell faces
  // =====================================================================================
  if (args->edge_velocity) {
    glDisable(GL_LIGHTING);
    glLineWidth(3);
    glBegin(GL_TRIANGLES);
    for (int x = 0; x < nx; x++) {
      for (int y = 0; y < ny; y++) {
        for (int z = 0; z < nz; z++) {
          
          glColor3f(1,0,0);
          glVertex3f(x,y+0.6,z+0.4);
          glVertex3f(x,y+0.4,z+0.6);
          glVertex3f(x+get_u_plus(x-1,y,z),y+0.5,z+0.5);

          glColor3f(0,1,0);
          glVertex3f(x+0.4,y+0.1,z+0.6);
          glVertex3f(x+0.6,y-0.1,z+0.4);
          glVertex3f(x+0.5,y+get_v_plus(x,y-1,z),z+0.5);

          glColor3f(0,0,1);
          glVertex3f(x+0.6,y+0.4,z);
          glVertex3f(x+0.4,y+0.6,z);
          glVertex3f(x+0.5,y+0.5,z+get_w_plus(x,y,z-1));
        } 
      }
    }
    glEnd();
    glEnable(GL_LIGHTING);
  }

  // =====================================================================================
  // visualize the average cell velocity
  // =====================================================================================
  if (args->velocity) {
    glDisable(GL_LIGHTING);
    glColor3f(1,0,0);
    glLineWidth(3);
    glBegin(GL_LINES);
    for (int x = 0; x < nx; x++) {
      for (int y = 0; y < ny; y++) {
	for (int z = 0; z < nz; z++) {
	  glVertex3f(x+0.5,y+0.5,z+0.5);
	  glVertex3f(x+0.5+get_u_avg(x,y,z),y+0.5+get_v_avg(x,y,z),z+0.5+get_w_avg(x,y,z));
  	} 
      }
    }
    glEnd();
    glEnable(GL_LIGHTING);
  } 

  // =====================================================================================
  // render a marching cubes representation of the surface
  //    note: make sure you set the Grid::getIsovalue() to what you want
  // =====================================================================================
  if (args->surface) {
    glPushMatrix();
    m.SetToIdentity();
    m *= Matrix::MakeTranslation(Vec3f(0.5,0.5,0.5));
    glMultMatrixf(m.glGet());
    glBegin(GL_TRIANGLES);
    for (int x = 0; x < nx; x++) {
      for (int y = 0; y < ny; y++) {
	for (int z = 0; z < nz; z++) {
	  MarchingCubes(x,y,z,args->wireframe);
	} 
      }
    }
    glEnd();
    glPopMatrix();
  } 

  // =====================================================================================
  // render the MAC cells (FULL, SURFACE, or EMPTY)
  // =====================================================================================
  if (args->cubes) {
    for (int x = 0; x < nx; x++) {
      for (int y = 0; y < ny; y++) {
        for (int z = 0; z < nz; z++) {
          glPushMatrix();
          m.SetToIdentity();
          m *= Matrix::MakeTranslation(Vec3f(x,y,z));
          glMultMatrixf(m.glGet());
          Cell *cell = getCell(x,y,z);
          if (cell->getStatus() == CELL_FULL) {
            glColor3f(1,0,0);
            CubePaint(); 
          } else if (cell->getStatus() == CELL_SURFACE) {
            glColor3f(0,0,1);
            CubePaint(); 
          }
          glPopMatrix();
        } 
      }
    }
  }
  
  // =====================================================================================
  // render a visualization of the pressure
  // =====================================================================================
  if (args->pressure) {
    float min_p = getCell(0,0,0)->getPressure()-0.00001;
    float max_p = getCell(0,0,0)->getPressure()+0.00001;

    for (int x = 0; x < nx; x++) {
      for (int y = 0; y < ny; y++) {
        for (int z = 0; z < nz; z++) {
	  float p = getCell(x,y,z)->getPressure();;
          min_p = min2(min_p,p);
          max_p = max2(max_p,p);
	}
      }
    }

    for (int x = 0; x < nx; x++) {
      for (int y = 0; y < ny; y++) {
        for (int z = 0; z < nz; z++) {
          glPushMatrix();
          m.SetToIdentity();
          m *= Matrix::MakeTranslation(Vec3f(x,y,z));
          glMultMatrixf(m.glGet());
	  float p = getCell(x,y,z)->getPressure();
          p = (p-min_p)/float(max_p-min_p);
	  assert(p >= 0 && p <= 1);
          insertColor(p);
          CubePaint(); 
          glPopMatrix();
        } 
      }
    }
  } 

  // *********************************************************************  
  // ASSIGNMENT:
  //
  // For extra credit, add additional visualization modes.  Render smoke, etc.!
  //
  // *********************************************************************    

}


void CubePaint() {

  glBegin (GL_QUADS);

  // front 
  glNormal3f(0,0,1.0);
  glVertex3f(0.0,0.0,1.0);
  glVertex3f(1.0,0.0,1.0);
  glVertex3f(1.0,1.0,1.0);
  glVertex3f(0.0,1.0,1.0);

  // top 
  glNormal3f(0,1.0,0);
  glVertex3f(0.0,1.0,0.0);
  glVertex3f(0.0,1.0,1.0);
  glVertex3f(1.0,1.0,1.0);
  glVertex3f(1.0,1.0,0.0);

  // right
  glNormal3f(1.0,0,0);
  glVertex3f(1.0,0.0,0.0);
  glVertex3f(1.0,1.0,0.0);
  glVertex3f(1.0,1.0,1.0);
  glVertex3f(1.0,0.0,1.0);

  // back 
  glNormal3f(0,0,-1.0);
  glVertex3f(0.0,0.0,0.0);
  glVertex3f(0.0,1.0,0.0);
  glVertex3f(1.0,1.0,0.0);
  glVertex3f(1.0,0.0,0.0);

  // bottom
  glNormal3f(0,-1.0,0);
  glVertex3f(0.0,0.0,0.0);
  glVertex3f(1.0,0.0,0.0);
  glVertex3f(1.0,0.0,1.0);
  glVertex3f(0.0,0.0,1.0);

  // left 
  glNormal3f(-1.0,0,0);
  glVertex3f(0.0,0.0,0.0);
  glVertex3f(0.0,0.0,1.0);
  glVertex3f(0.0,1.0,1.0);
  glVertex3f(0.0,1.0,0.0);

  glEnd();
}


// ==============================================================
// ==============================================================
// ==============================================================