/* Three dimensional lattice Boltzmann code for the 370 class */

#include <unistd.h>
#include <math.h>
#include <mygraph.h>


#define xdim 50
#define ydim 50
#define zdim 50

#define D 3
#define V 15


int Xdim=xdim,Ydim=ydim;
int v[V][3]={{0,0,0},
	     {1,0,0},{-1,0,0},{0,1,0},{0,-1,0},{0,0,1},{0,0,-1},
	     {1,1,1},{-1,1,1},{1,-1,1},{-1,-1,1},{1,1,-1},{-1,1,-1},
	     {1,-1,-1},{-1,-1,-1}};
double w[V] = {2./9.,
	       1./9.,1./9.,1./9.,1./9.,1./9.,1./9.,
	       1./72.,1./72.,1./72.,1./72.,1./72.,1./72.,1./72.,1./72.};
double tet = 1./3.;
typedef double dens[V][xdim][ydim][zdim];
dens f1,f2,*f,*ff;

double tau=1;

// Variables for graphical representation
int ZZ=0;
double nxy[xdim][ydim],uxy[xdim][ydim][2];
int nxyreq=0,uxyreq=0;

void init(dens f){
  for (int x=0;x<xdim;x++)
    for (int y=0; y<ydim; y++)
      for (int z=0; z<zdim; z++){
	double n=1;
	if (pow(x-xdim/2,2)+pow(y-ydim/2,2)+pow(z-zdim/2,2)<5) n=0.8;
	double u[D]={0,0,0};
	if (y>ydim/2)
	  u[0]=0.1;
	else u[0]=-0.1;
	for (int i=0;i<V;i++){
	  double vu=0,uu=0;
	  for (int a=0;a<D;a++){
	    vu+=v[i][a]*u[a];
	    uu+=u[a]*u[a];
	  }
	  f[i][x][y][z]=n*w[i]*(1+vu/tet+0.5*vu*vu/tet/tet-0.5*uu/tet);
	}
      }
}

void LB(dens f,dens ff,double tau){
  for (int x=0;x<xdim;x++)
    for (int y=0;y<ydim;y++)
      for (int z=0;z<zdim;z++){
	double n=0;
	double u[D]={0,0,0};
	for (int i=0;i<V;i++){
	  n+=f[i][x][y][z];
	  for (int a=0;a<D;a++)
	    u[a]+=f[i][x][y][z]*v[i][a];
	}
	for (int a=0;a<D;a++) u[a]/=n;
	for (int i=0;i<V;i++){
	  double vu=0,uu=0;
	  for (int a=0;a<D;a++){
	    vu+=v[i][a]*u[a];
	    uu+=u[a]*u[a];
	  }
	  double f0=n*w[i]*(1+vu/tet+0.5*vu*vu/tet/tet-0.5*uu/tet);
	  int xn=(x+v[i][0]+xdim)%xdim;
	  int yn=(y+v[i][1]+ydim)%ydim;
	  int zn=(z+v[i][2]+zdim)%zdim;
	  ff[i][xn][yn][zn]= f[i][x][y][z]+1/tau*(f0-f[i][x][y][z]);
	}
      }
}

void GetGraph(dens f){
  if (uxyreq){
    nxyreq=0;
    uxyreq=0;
    for (int x=0;x<xdim;x++)
      for (int y=0;y<ydim;y++){
	nxy[x][y]=0;
	uxy[x][y][0]=0;
	uxy[x][y][1]=0;
	for (int i=0;i<V;i++){
	  nxy[x][y]+=f[i][x][y][ZZ];
	  uxy[x][y][0]+=f[i][x][y][ZZ]*v[i][0];
	  uxy[x][y][1]+=f[i][x][y][ZZ]*v[i][1];
	}
	uxy[x][y][0]/=nxy[x][y];
	uxy[x][y][1]/=nxy[x][y];
      }
  }
  if (nxyreq){
    nxyreq=0;
    for (int x=0;x<xdim;x++)
      for (int y=0;y<ydim;y++){
	nxy[x][y]=0;
	for (int i=0;i<V;i++)
	  nxy[x][y]+=f[i][x][y][ZZ];
      }
  }
} 

void initgr(){
  init(f1);
  init(f2);
}

int main(){
  int done=0,sstep=0,cont=0,Repeat=1;
  f=&f1;
  ff=&f2;
  init(*f);
  
  DefineGraphNxN_R("nxy",&nxy[0][0],&Xdim,&Ydim,&nxyreq);
  DefineGraphNxN_RxR("uxy",&uxy[0][0][0],&Xdim,&Ydim,&uxyreq);
  StartMenu("LB 3d",1);
  DefineInt("ZZ",&ZZ);
  DefineGraph(contour2d_,"Sections");
  DefineFunction("init",initgr);
  DefineDouble("tau",&tau);
  DefineInt("Repeat",&Repeat);
  DefineBool("sstep",&sstep);
  DefineBool("cont",&cont);
  DefineBool("done",&done);
  EndMenu();
  
  while (!done){
    Events(1);
    GetGraph(*f);
    DrawGraphs();
    if (cont || sstep){
      sstep=0;
      for (int i=0;i<Repeat;i++){
	LB(*f,*ff,tau);
	if (f==&f1) {f=&f2;ff=&f1;}
	else {f=&f1;ff=&f2;}
      }
    }
    else sleep(1);
  }

 }