#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <sys/time.h>

/*
 * Simple atomistic Monte Carlo update program
 * Uses metropolis update
 * The program simulates N ions of opposite charge 
 * in 2d at given temperature and volume
 * Prints out measurements of average near neighbour number
 * 
 * Needs files: 
 *   ions.c mersenne.h mersenne_inline.c gaussian_ran.c cmdline.c
 * To compile:
 *   cc -O2 -o ions ions.c mersenne_inline.c gaussian_ran.c cmdline.c -lm
 *
 * How to use: check usage-string below.
 *
 * Example:
 *   ions -s40 -n64 -t0.03 -l50000 -p100 -m100 -L
 *
 * Potential is V = q1*q2/r + 1/r^8
 */

/* define the following if you want to do plots with grace; otherwise
   comment this out */
#define GRACE

#ifdef GRACE
#include <grace_np.h>
#endif

#include "mersenne.h"

/* command line handling - see cmdline.c */
void cmdl_init( int argc, char *argv[] );
int cmdl_get_opt();
int cmdl_get_int(char *err);
double cmdl_get_double(char *err);
int cmdl_is_char(char c);
char * cmdl_get_string( char *err );
int cmdl_args_remain();

void plot_init(double scale);
void plot_data(double *x, double *y, int *q, int n, int nsets);
void plot_ready(int i);

double gaussian_ran();

char usage[] = 
  "Use: ions [options] > result\n\
 Options:\n\
 -t temperature   temperature of the system (REQUIRED)\n\
 -n N_atoms       total number of atoms (default 20)\n\
 -s size          size of the box (default 30)\n\
 -l loops         number of update loops (default 50000)\n\
 -g               use gravitational potential instead of Coulomb\n\
 -w width         width of random initial state (default 6)\n\
 -L               initial state is lattice instead of random\n\
 -S scale         Metropolis scale (default 0.07)\n\
 -m n             measure interval (default 32, value 0 : measure off)\n\
 -p n             plot with grace at interval n\n\
 Prints out in columns:\n\
  1: iteration\n\
  2: average potential energy\n\
  3: average pair distance\n\
  4: average nearest neighbour distance\n\
  5: # of neighbours at distance < 2\n\
  6: # of neighbours at distance < 2.1\n\
  7: # of neighbours at distance < 2.2\n\
  8: # of neighbours at distance < 2.4\n\
  9: Metropolis acceptance\n";

/* This routine gives the 2-particle potential */

double V(double x, double y, int q1q2)
{
  double r2,r4;
  r2 = x*x + y*y;
  r4 = r2*r2;
  if (q1q2 != 0) 
    return( q1q2/sqrt(r2) + 1.0/(r4*r4));     // coulomb
  else 
    return( -1.0/sqrt(r2) + 1.0/(r4*r4));     // gravity
}


int main(int argc, char* argv[])
{
  struct timeval tv;
  struct timezone tz;
  int i,j,loop, *q;
  int accept=0,try=0;
  double *x, *y;

  /* defaults for options */
  int n_atoms = 20;  
  int n_loops = 50000;
  double size = 30;
  double beta = 0;            /* this one is required! */
  double initial_width = 6.0; /* default initial width for update */
  double scale = 0.07;        /* default Metropolis scale, seems to work */
  int gravity = 0;
  int plot_interval = 0;      /* plotting off */
  int measure_interval = 32;
  int lattice = 0;

  /* command line arguments - init */
  cmdl_init( argc, argv );

  /* get options */
  while (i = cmdl_get_opt()) {
    switch(i) {

    case 'n': n_atoms = cmdl_get_int(usage);
      break;
    case 't': beta = 1.0/cmdl_get_double(usage);
      break;
    case 's': size = cmdl_get_int(usage);
      break;
    case 'l': n_loops = cmdl_get_int(usage);
      break;
    case 'p': plot_interval = cmdl_get_int(usage);
      break;
    case 'g': gravity = 1;
      break;
    case 'm': measure_interval = cmdl_get_int(usage);
      break;
    case 'w': initial_width = cmdl_get_double(usage);
      break;
    case 'L': lattice = 1;
      break;
    default: fprintf(stderr,usage);
      exit(0);
    }
  } /* while */

  /* require beta = 1/T */
  if (beta == 0) {
    fprintf(stderr," T must be given!\n");
    exit(0);
  }

  /* allocate location and charge arrays */
  x = (double *)malloc( n_atoms * sizeof(double));
  y = (double *)malloc( n_atoms * sizeof(double));
  q = (int *)malloc( n_atoms * sizeof(int));

  /* seed the random nubers, using current time */
  gettimeofday( &tv, &tz );
  seed_mersenne( tv.tv_sec + tv.tv_usec );

  /* Set initial locations 
   * Set also electric charges -- if charge == 0, gravitational 
   */
  if (!lattice) {
    /* set initial coordinates at random - but close to each other! */
    for (i=0; i<n_atoms; i++) {
      x[i] = size/2 + initial_width * gaussian_ran(); 
      y[i] = size/2 + initial_width * gaussian_ran();
      if (!gravity) {
	if (i%2) q[i] = 1;
	else     q[i] = -1;
      } else     q[i] = 0;
    }
  } else {
    /* set up initial regular lattice, scale chosen to 1.3 */
    int n = sqrt((double)n_atoms);
    for (i=0; i<n_atoms; i++) {
      int xx = i%n;
      int yy = i/n;
      x[i] = size/2 + 1.3*(xx - 0.5*n);
      y[i] = size/2 + 1.3*(yy - 0.5*n);
      if (!gravity) {
	if ((xx+yy)%2) q[i] = +1; /* quarantees charge checkerboard */
	else           q[i] = -1;
      } else           q[i] = 0;
      /* this quarantees chekerboard for charges */
    }
  }
  for (i=0; i<n_atoms; i++) {
    /* ensure these are within the volume - 
     *   if not, put at random location 
     */
    if (x[i] < 0 || x[i] > size) x[i] = mersenne()*size;
    if (y[i] < 0 || y[i] > size) y[i] = mersenne()*size;
  }

  /* and start updating */

#ifdef GRACE
  if (plot_interval) plot_init(size);
#endif

  for (loop=0; loop<n_loops; loop++){
    accept = try = 0;
    /* modify the location, acc/rej */
    for (i=0; i<n_atoms; i++) {
      double xn,yn;
      double e1 = 0, e2 = 0;

      /* calculate the potential energy for i */
      for (j=0; j<n_atoms; j++) if (j != i) 
	e1 += V( x[i]-x[j], y[i]-y[j], q[i]*q[j] );
      
      /* update position -- keep within box! */
      do xn = x[i] + scale * gaussian_ran(); while (xn < 0 || xn > size);
      do yn = y[i] + scale * gaussian_ran(); while (yn < 0 || yn > size);
      
      /* and calculate new potential energy */
      for (j=0; j<n_atoms; j++) if (j != i) 
	e2 += V( xn-x[j], yn-y[j], q[i]*q[j] );
      
      /* now Metropolis accept/reject */
      if ( exp( (e1-e2)*beta ) > mersenne() ) {
	accept++;
	x[i] = xn;
	y[i] = yn;
      }
      try++;
    }
    
    if (measure_interval && loop % measure_interval == 0) {

      /* print meas every measure_interval iteration 
       * measure the average distances and the mean number of neighbours
       * inside distance 
       */
      double pot = 0;
      double dist = 0, mindist = 0, md;
      int nn1 = 0, nn2 = 0, nn3 = 0, nn4 = 0;

      for (i=0; i<n_atoms; i++) {
	md = 10*size;
	for (j=0; j<n_atoms; j++) if (j!=i) {
	  double xd = x[i]-x[j];
	  double yd = y[i]-y[j];
	  double d = sqrt(xd*xd + yd*yd);

	  pot += V(xd,yd,q[i]*q[j]);

	  dist += d;
	  if (d < md) md = d;     /* min. dist to atom i */
	  if (d <= 2.0) nn1 += 1;
	  if (d <= 2.1) nn2 += 1;
	  if (d <= 2.2) nn3 += 1;
	  if (d <= 2.4) nn4 += 1;
	}
	mindist += md;
      }

      printf("%d %g %g %g %g %g %g %g %g\n",
	     loop,
	     pot/(2.0*n_atoms), /* avg. potential - pairs counted twice! */
	     dist/(n_atoms*(n_atoms-1)),  /* avg. dist */
	     mindist/n_atoms,             /* avg. nearest atom */
	     (double)nn1/(n_atoms),       /* avg. number of atoms at < 2 */
	     (double)nn2/(n_atoms),       /* avg. number of atoms at < 2.1 */
	     (double)nn3/(n_atoms),       /* < 2.2 */
	     (double)nn4/(n_atoms),       /* < 2.4 */
	     (double)accept/try);         /* metropolis acceptance */
      
    }
#ifdef GRACE
    if (plot_interval && loop % plot_interval == 0) {
      /* do only every 16th plot with grace */

      if (!gravity) 
	plot_data(x,y,q,n_atoms,2);
      else 
	plot_data(x,y,q,n_atoms,1);
      
      plot_ready(loop);
    }
#endif
      
  }
  exit(1);
}



#ifdef GRACE
/* The rest is included only if GRACE is defined */

static double symbolsize;  /* plot symbol size */

void plot_init(double scale)
{
 
  /* Start grace - free canvas (allows resize), no questions, default .eps */
  GraceOpenVA("xmgrace",2048,"-free","-noask","-hdevice","EPS",NULL);

  GracePrintf("page size 700 700");
  GracePrintf("focus off");   /* turn off the focus markers (annoying) */
  GracePrintf("g0 on");       /* Activate graph 0 */
  GracePrintf("with g0");     /* reset the current graph to graph 0 */

  GracePrintf("view 0.01, 0.05, 0.9, 0.94"); /* set the viewport - fills */  

  GracePrintf("world 0, 0, %g, %g",scale,scale); /* box size */
  GracePrintf("xaxis off");   /* no axis labels or ticks */
  GracePrintf("yaxis off"); 

  symbolsize = 45.0 / scale;
  
}


void plot_data(double *x, double *y, int *q, int n, int nsets)
{
  int i,set;

  for (set=0; set<nsets; set++) {
    GracePrintf("kill s%d",set);    /* make sure set s0 is available */
    GracePrintf("s%d on",set);      /* activate s0 */

    GracePrintf("s%d type xy",set);  /* data type (xy xye etc.) */
    GracePrintf("s%d symbol 1",set); /* circle */
    GracePrintf("s%d symbol size %g",set,symbolsize);   
    GracePrintf("s%d symbol color %d",set,1); 
    GracePrintf("s%d symbol pattern 1",set); /* solid color */
    GracePrintf("s%d symbol fill color %d",set,set+4); 
    /* colors 4,5: blue,yellow */
    GracePrintf("s%d symbol fill pattern 1",set); /* solid color */
    GracePrintf("s%d symbol linewidth 1.0",set);
    GracePrintf("s%d symbol linestyle 1",set);  /* solid line */
    GracePrintf("s%d line type 0",set);         /* but not drawn */

    for (i = 0; i < n; i++) {
      /* and print the data */
      if (nsets == 1 || 2*set == q[i]+1) {
	/* Either print all or in Coulomb case map -+1 to 0,1*/
	GracePrintf("g0.s%d point %g, %g",set,x[i],y[i]);
      }	
    }
  }
}

void plot_ready(int i)
{
  /* do the plotting */
  GracePrintf("subtitle \"%d\"",i);
  GracePrintf("redraw");
}

#endif /* ifdef GRACE */