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

typedef struct part {double x; double y; double z;double vx; double vy; double vz;} part;
typedef struct vect {double x; double y; double z;} vect;
#define N 100

int initialized=1;
double Lx=10,Ly=10,Lz=10;
double e,ek,ep,kT=0.1;
double dt=0.01,dx=0.1,*dd,rate;
part p[N];

double phi=0,tet=0,disp=10,alpha=1;

double (*iterate)(part p[N], double dd);
void (*Analyze)(part p[N]);

double dPE(part p[N],vect d,int n){
  double dE=0;
  for (int m=0; m<N; m++)
    if (n!=m)
      for (int x=-1;x<2; x++)
	for (int y=-1; y<2; y++)
	  for (int z=-1; z<2; z++){
	    double dist2A=  (p[n].x-(p[m].x+x*Lx))*(p[n].x-(p[m].x+x*Lx))
			       +(p[n].y-(p[m].y+y*Ly))*(p[n].y-(p[m].y+y*Ly))
			       +(p[n].z-(p[m].z+z*Lz))*(p[n].z-(p[m].z+z*Lz));
	    double dist2B=  (p[n].x+d.x-(p[m].x+x*Lx))*(p[n].x+d.x-(p[m].x+x*Lx))
			       +(p[n].y+d.y-(p[m].y+y*Ly))*(p[n].y+d.y-(p[m].y+y*Ly))
			       +(p[n].z+d.z-(p[m].z+z*Lz))*(p[n].z+d.z-(p[m].z+z*Lz));
	    double EB=pow(dist2B,-6) - pow(dist2B,-3);
	    double EA=pow(dist2A,-6) - pow(dist2A,-3);
	    dE+= EB-EA;
	  }
  return dE;
}

double PE(part p[N]){
  double E=0;
  for (int n=0; n<N; n++)
    for (int m=n+1; m<N; m++)
      for (int x=-1;x<2; x++)
	for (int y=-1; y<2; y++)
	  for (int z=-1; z<2; z++){
	    double dist2=  (p[n].x-(p[m].x+x*Lx))*(p[n].x-(p[m].x+x*Lx))
			       +(p[n].y-(p[m].y+y*Ly))*(p[n].y-(p[m].y+y*Ly))
			       +(p[n].z-(p[m].z+z*Lz))*(p[n].z-(p[m].z+z*Lz));
	    E+= pow(dist2,-6) - pow(dist2,-3);
	  }
  return E;
}

double KE(part p[N]){
  double E=0;
  for (int n=0; n<N; n++){
    E+=0.5*(p[n].vx*p[n].vx+p[n].vy*p[n].vy+p[n].vz*p[n].vz);
  }
  return E;
}

void FF(part p[N], vect F[N]){
  memset(F,0,N*sizeof(vect));
  for (int n=0; n<N; n++)
    for (int m=n+1; m<N; m++){
      for (int x=-1; x<2; x++)
	for (int y=-1; y<2; y++)
	  for (int z=-1; z<2; z++){
	    double dist= sqrt( (p[n].x-(p[m].x+x*Lx))*(p[n].x-(p[m].x+x*Lx))
			       +(p[n].y-(p[m].y+y*Ly))*(p[n].y-(p[m].y+y*Ly))
			       +(p[n].z-(p[m].z+z*Lz))*(p[n].z-(p[m].z+z*Lz)));
	  double f =-12*pow(dist,-14)+6*pow(dist,-8); 
	  F[n].x+=f*(p[m].x+x*Lx-p[n].x);
	  F[n].y+=f*(p[m].y+y*Ly-p[n].y);
	  F[n].z+=f*(p[m].z+z*Lz-p[n].z);
	  F[m].x-=f*(p[m].x+x*Lx-p[n].x);
	  F[m].y-=f*(p[m].y+y*Ly-p[n].y);
	  F[m].z-=f*(p[m].z+z*Lz-p[n].z);
	}
    }
}

double MC(part p[N],double dx){
  vect d;
  int count=0;
  for (int n=0; n<N; n++){
    d.x = dx*((double) 2*rand()/RAND_MAX-1);
    if (p[n].x+d.x<0) d.x+=Lx;
    if (p[n].x+d.x>=Lx) d.x-=Lx;
    d.y = dx*((double) 2*rand()/RAND_MAX-1);
    if (p[n].y+d.y<0) d.y+=Ly;
    if (p[n].y+d.y>=Ly) d.y-=Ly;
    d.z = dx*((double) 2*rand()/RAND_MAX-1);
    if (p[n].z+d.z<0) d.z+=Lz;
    if (p[n].z+d.z>=Lz) d.z-=Lz;
    double dE = dPE(p,d,n);
    if (dE<=0){
      p[n].x += d.x;
      p[n].y += d.y;
      p[n].z += d.z;
      count++;
    } else if ((double) rand()/RAND_MAX<exp(-dE/kT)){
      p[n].x += d.x;
      p[n].y += d.y;
      p[n].z += d.z;
      count++;
    }
  }
  return (double) count/N;
}

double Verlet(part p[N],double dt){
  static vect F[N];

  if (!initialized) {FF(p,F); initialized =0;}

  for (int n=0; n<N; n++){
    p[n].vx+=dt*0.5*F[n].x;
    p[n].vy+=dt*0.5*F[n].y;
    p[n].vz+=dt*0.5*F[n].z;
    
    p[n].x += dt*p[n].vx;
    p[n].y += dt*p[n].vy;
    p[n].z += dt*p[n].vz;

    /*    if (p[n].x<0) p[n].x+=Lx;
	  if (p[n].x>Lx) p[n].x-=Lx;*/
    p[n].x = fmod(p[n].x+Lx,Lx);
    p[n].y = fmod(p[n].y+Ly,Ly);
    p[n].z = fmod(p[n].z+Lz,Lz);
  }
  FF(p,F);
  for (int n=0; n<N; n++){
    p[n].vx+=dt*0.5*F[n].x;
    p[n].vy+=dt*0.5*F[n].y;
    p[n].vz+=dt*0.5*F[n].z;
  }
}

void rescaleV(){
  for (int n=0; n<N; n++){
    p[n].vx*=alpha;
    p[n].vy*=alpha;
    p[n].vz*=alpha;
  }
}

void initv(double kT){
  double v0=sqrt(3*kT);
  for (int n=0;n<N;n++){
    p[n].vx= v0*((double) rand()/RAND_MAX-0.5);
    p[n].vy= v0*((double) rand()/RAND_MAX-0.5);
    p[n].vz= v0*((double) rand()/RAND_MAX-0.5);
  }
}

void init(){
  initialized=0;
  int n=0;
  int NN = ceil(pow(N,1./3.));
  for (int x=0;x<NN;x++)
    for (int y=0; y<NN; y++)
      for (int z=0; z<NN; z++){
	if (n<N){
	  p[n].x=(x+0.5)*Lx/NN;
	  p[n].y=(y+0.5)*Ly/NN;
	  p[n].z=(z+0.5)*Lz/NN;
	}
	n++;
      }
  initv(kT);
}

void xydraw(int xdim,int ydim){
  double scale;
  int xoff,yoff,r1,r2;

  if (Lx*ydim>xdim*Ly){
    scale = xdim/Lx;
    xoff=0;
    yoff= 0.5*(ydim-Ly*scale);
  }
  else {
    scale = ydim/Ly;
    xoff= 0.5*(xdim-Lx*scale);
    yoff= 0;
  }
  r1=1*scale;
  r2=pow(2,-1./6.)*scale;
  for (int n=0; n<N; n++){
    mycircle(3,xoff+p[n].x*scale,ydim-yoff-p[n].y*scale,r2);
    mycircle(2,xoff+p[n].x*scale,ydim-yoff-p[n].y*scale,r1);
  }

}

int PartComp(vect *P1,vect *P2){
  if ((*P1).z<(*P2).z) return 1;
  else return -1;
}

void draw3d(int xdim, int ydim){
  vect box[8]={{0,0,0},{Lx,0,0},{0,Ly,0},{Lx,Ly,0},{0,0,Lz},{Lx,0,Lz},{0,Ly,Lz},{Lx,Ly,Lz}};
  vect Part[N];
  double scale;
  
  /* Translate */
  for (int n=0; n<8; n++){
    box[n].x-=Lx/2;
    box[n].y-=Ly/2;
    box[n].z-=Lz/2;
  }
  for (int n=0; n<N; n++){
    Part[n].x=p[n].x-Lx/2;
    Part[n].y=p[n].y-Ly/2;
    Part[n].z=p[n].z-Lz/2;
  }
  /* Rotate */
  double c=cos(phi);
  double s=sin(phi);
  for (int n=0;n<8;n++){
    double xt=c*box[n].x-s*box[n].z;
    box[n].z =c*box[n].z+s*box[n].x;
    box[n].x = xt;
  }
  for (int n=0;n<N;n++){
    double xt=c*Part[n].x-s*Part[n].z;
    Part[n].z =c*Part[n].z+s*Part[n].x;
    Part[n].x = xt;
  }
  c=cos(tet);
  s=sin(tet);
  for (int n=0;n<8;n++){
    double yt=c*box[n].y+s*box[n].z;
    box[n].z =c*box[n].z-s*box[n].y;
    box[n].y = yt;
  }
  for (int n=0;n<N;n++){
    double yt=c*Part[n].y+s*Part[n].z;
    Part[n].z =c*Part[n].z-s*Part[n].y;
    Part[n].y = yt;
  }
  /* Translate */
  for (int n=0; n<8; n++){
    box[n].z+=disp+Lz/2;
  }
  for (int n=0; n<N; n++){
    Part[n].z+=disp+Lz/2;
  }
  qsort(&(Part[0].x),N,sizeof(vect),PartComp);


  if (Lx*ydim>xdim*Ly){
    scale = 0.8*xdim/Lx*disp;
  }
  else {
    scale = 0.8*ydim/Ly*disp;
  }
  int x1,y1,x2,y2;
  x1 = xdim/2 + box[0].x/box[0].z*scale;
  y1 = ydim/2 - box[0].y/box[0].z*scale;
  x2 = xdim/2 + box[1].x/box[1].z*scale;
  y2 = ydim/2 - box[1].y/box[1].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[1].x/box[1].z*scale;
  y1 = ydim/2 - box[1].y/box[1].z*scale;
  x2 = xdim/2 + box[3].x/box[3].z*scale;
  y2 = ydim/2 - box[3].y/box[3].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[2].x/box[2].z*scale;
  y1 = ydim/2 - box[2].y/box[2].z*scale;
  x2 = xdim/2 + box[3].x/box[3].z*scale;
  y2 = ydim/2 - box[3].y/box[3].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[2].x/box[2].z*scale;
  y1 = ydim/2 - box[2].y/box[2].z*scale;
  x2 = xdim/2 + box[0].x/box[0].z*scale;
  y2 = ydim/2 - box[0].y/box[0].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[0].x/box[0].z*scale;
  y1 = ydim/2 - box[0].y/box[0].z*scale;
  x2 = xdim/2 + box[4].x/box[4].z*scale;
  y2 = ydim/2 - box[4].y/box[4].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[1].x/box[1].z*scale;
  y1 = ydim/2 - box[1].y/box[1].z*scale;
  x2 = xdim/2 + box[5].x/box[5].z*scale;
  y2 = ydim/2 - box[5].y/box[5].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[2].x/box[2].z*scale;
  y1 = ydim/2 - box[2].y/box[2].z*scale;
  x2 = xdim/2 + box[6].x/box[6].z*scale;
  y2 = ydim/2 - box[6].y/box[6].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[3].x/box[3].z*scale;
  y1 = ydim/2 - box[3].y/box[3].z*scale;
  x2 = xdim/2 + box[7].x/box[7].z*scale;
  y2 = ydim/2 - box[7].y/box[7].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[4].x/box[4].z*scale;
  y1 = ydim/2 - box[4].y/box[4].z*scale;
  x2 = xdim/2 + box[5].x/box[5].z*scale;
  y2 = ydim/2 - box[5].y/box[5].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[5].x/box[5].z*scale;
  y1 = ydim/2 - box[5].y/box[5].z*scale;
  x2 = xdim/2 + box[7].x/box[7].z*scale;
  y2 = ydim/2 - box[7].y/box[7].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[6].x/box[6].z*scale;
  y1 = ydim/2 - box[6].y/box[6].z*scale;
  x2 = xdim/2 + box[7].x/box[7].z*scale;
  y2 = ydim/2 - box[7].y/box[7].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  x1 = xdim/2 +  box[6].x/box[6].z*scale;
  y1 = ydim/2 - box[6].y/box[6].z*scale;
  x2 = xdim/2 + box[4].x/box[4].z*scale;
  y2 = ydim/2 - box[4].y/box[4].z*scale;
  mydrawline(1,x1,y1,x2,y2);
  /* Sort */ 

  for (int n=0; n<N; n++){
    myfilledcircle(3,xdim/2+Part[n].x/Part[n].z*scale,ydim/2-Part[n].y/Part[n].z*scale,1.122/Part[n].z*scale);
    myfilledcircle(2,xdim/2+Part[n].x/Part[n].z*scale,ydim/2-Part[n].y/Part[n].z*scale,1/Part[n].z*scale);
  }
}

void AnalyzeMD(part p[N]){
  ek=KE(p);
  ep=PE(p);
  e =ek+ep;
  kT=2./3.*ek/N;
}
void AnalyzeMC(part p[N]){
  ep=PE(p);
}

void setMD(){
  iterate=Verlet;
  Analyze=AnalyzeMD;
  dd=&dt;
  initv(kT);
}

void setMC(){
  iterate=MC;
  Analyze=AnalyzeMC;
  dd=&dx;
}

int main(){
  int sstep=0,cont=0,done=0,repeat=1;
  init();
  setMC();
  AddFreedraw("xy",xydraw);
  AddFreedraw("3d",draw3d);

  StartMenu("LJ 3d sim",1);
  DefineDouble("Lx",&Lx);
  DefineDouble("Ly",&Ly);
  DefineDouble("Lz",&Lz);
  DefineDouble("dt",&dt);
  DefineDouble("dx",&dx);
  StartMenu("Energy",0);
  DefineDouble("Total Energy",&e);
  DefineDouble("Kinetic Energy",&ek);
  DefineDouble("Potential Energy",&ep);
  DefineDouble("kT",&kT);
  DefineDouble("MC rate",&rate);
  EndMenu();
  StartMenu("Rescale V",0);
  DefineDouble("Factor",&alpha);
  DefineFunction("Rescale",rescaleV);
  EndMenu();

  StartMenu("init",0);
  DefineDouble("kT",&kT);
  DefineFunction("init",init);
  EndMenu();
  StartMenu("Visual",0);
  DefineDouble("phi",&phi);
  DefineDouble("tet",&tet);
  DefineDouble("disp",&disp);
  EndMenu();
  DefineGraph(freedraw_,"Visualize");
  DefineFunction("Set MD",setMD);
  DefineFunction("Set MC",setMC);
  DefineBool("sstep",&sstep);
  DefineBool("cont",&cont);
  DefineInt("repeat",&repeat);
  DefineBool("done",&done);
  EndMenu();
  while (!done){
    Events(1);
    DrawGraphs();
    if (cont||sstep){
      sstep=0;
      for (int i=0; i<repeat; i++) rate=iterate(p,*dd);
      Analyze(p);
    }
    else sleep(1);
  }
}