/* n Argon atoms in 2D */

#include <stdlib.h>
#include <assert.h>
#include <math.h>
#include "force.h"


const char *SWITH_METHOD_NAME[ SM_MAX_SWITCH_METHOD ] = {
  "Discontinuous", "C0", "C1", "C2", 
};

static void DC_compute_v_vp(Data *data, double r, double *pe, double *neg_vp);
static void C0_compute_v_vp(Data *data, double r, double *pe, double *neg_vp);
static void C1_compute_v_vp(Data *data, double r, double *pe, double *neg_vp);
static void C2_compute_v_vp(Data *data, double r, double *pe, double *neg_vp);
typedef void 
(*COMPUTE_V_VP_PT) (Data *data, double r, double *pe, double *neg_vp);
static COMPUTE_V_VP_PT CVVP[ SM_MAX_SWITCH_METHOD ] = {
  DC_compute_v_vp, C0_compute_v_vp, C1_compute_v_vp,
  C2_compute_v_vp
};
static void (*compute_v_vp)(Data *data, double r, double *pe, double *neg_vp); 

static void compute_dist2(const double systemsize,
			  const double pos1[DIMENSION], 
			  const double pos2[DIMENSION],
			  double dist[DIMENSION], double *r2) ;
static void LJ_poten_force(const double dist[DIMENSION], const double inv_r, 
			   double *poten, double fij[DIMENSION]);



/***************************************************************************/

int parse_input(const int argc, char** argv, Data *data)
{
  if (7 == argc) {
    data->dt = atof(argv[1]); 
    data->nsteps = atoi(argv[2]); 
    data->cutoff = atof(argv[3]);
    data->switch_distance = atof(argv[4]);
    data->switch_method = atoi(argv[5]);
    data->outfile = argv[6];

    if (data->dt > 1.0) {
      printf("timstep may be too large.\n");
    } else if (data->dt < 0.0) {
      printf("timestep should not be negative\n");
      return 1;
    }
    if (data->switch_method < SM_DC || 
	data->switch_method > SM_MAX_SWITCH_METHOD) {
      printf("switch degree should be within %d and %d\n", SM_DC, 
	     SM_MAX_SWITCH_METHOD);
      return 1;
    }
  } else if (2 == argc) {  /* only input the switch method */
    data->dt = DEFAULT_TIMESTEP;  
    data->nsteps = DEFAULT_NSTEPS;
    data->cutoff = DEFAULT_CUTOFF;
    data->switch_distance = DEFAULT_SWITCH_DIST; 
    data->switch_method = atoi(argv[1]);
    if (data->switch_method < SM_DC || 
	data->switch_method > SM_MAX_SWITCH_METHOD) {
      printf("switch degree should be within %d and %d\n", SM_DC, 
	     SM_MAX_SWITCH_METHOD);
      return 1;
    }
  } else if (1 == argc) { /* default values */
    data->dt = DEFAULT_TIMESTEP;  
    data->nsteps = DEFAULT_NSTEPS;
    data->cutoff = DEFAULT_CUTOFF;
    data->switch_distance = DEFAULT_SWITCH_DIST; 
    data->switch_method = DEFAULT_SWTICH_METHOD;
  } else {
    fprintf(stderr,"Usage:\n");
    fprintf(stderr,"1. %s\n", argv[0]);
    fprintf(stderr,"   use default: (timestep=%f(fs), numsteps=%d,"
	           " cut_off=%f(A), switch_distance=%f(A), swith_method=%d)\n",
	           DEFAULT_TIMESTEP, DEFAULT_NSTEPS, DEFAULT_CUTOFF,
	           DEFAULT_SWITCH_DIST, SM_C2);
    fprintf(stderr,"2. %s switch_method\n", argv[0]);
    fprintf(stderr,"   others use default values\n");
    fprintf(stderr,"3. %s timestep numsteps cut_off switch_distance "
	           "swith_method\n", argv[0]);
    fprintf(stderr,"switch_distance, cut_off: Angstrom\n");
    fprintf(stderr,"time: femto-second\n");
    fprintf(stderr,"(switch_method: 0: discontinous,  1: C0, 2: C1, 3:C2)\n");
    return 1;
  }

  /*
  printf("parameter value:\n"
	 "  timestep=%f fs, nsteps=%d,\n"
	 "  cutoff=%f Angstrom, switch_distance=%f Angstrom\n" 
	 "  switch_method=%s, dimension=%d\n", 
	 data->dt, data->nsteps,  
	 data->cutoff, data->switch_distance,
	 SWITH_METHOD_NAME[data->switch_method], DIMENSION);
  */

  /* convert to internal unit */
  data->dt /= TIME_UNIT;  
  data->cutoff /= LENGTH_UNIT;
  data->switch_distance /= LENGTH_UNIT;
  /* data->nsteps = 300000; */

  return 0;
}

 

int initial_conditions(Data *data)
{    
  int dim;
  const int natom_line = 5; /* ------------------------------------ */
  data->natoms = 1;
  for (dim = 0; dim < DIMENSION; dim++) {
    data->natoms *= natom_line;
  }
  data->pos  = (double *)malloc( data->natoms*DIMENSION*sizeof(double) );
  data->vel  = (double *)malloc( data->natoms*DIMENSION*sizeof(double) );
  data->f    = (double *)calloc( (size_t) data->natoms*DIMENSION, 
				 sizeof(double) );  
  data->mass = (double *)malloc( data->natoms*DIMENSION*sizeof(double) );  
  assert(NULL != data->pos);
  assert(NULL != data->vel);
  assert(NULL != data->f);
  assert(NULL != data->mass);

  assert(data->natoms > 1);
  /* LJ potential funciton's minimum is at 1.1225 */
  data->systemsize = ((double) natom_line) * 1.3 + 0.1; 
  if (data->systemsize < 2.0 * data->cutoff) {
    data->systemsize = 2.0 * data->cutoff + 0.1;
  }

  /* INITIAL CONFIGURATIONS */
  {
    double  (*pos) = data->pos;
    double  *vel;
    double *mass;
    const double systemsize = data->systemsize;
    const double delta = systemsize / (double) natom_line - 0.1;
    int i;

    /* srandom(1); */  /* explicitly set up the random seed */
    if (1 == DIMENSION) {
      int ix;
      for (ix = 0; ix < natom_line; ix++) {
	pos[0] = -0.5*systemsize + (ix + 0.5) * delta;
	pos++;
      }
    } else if (2 == DIMENSION) {
      int ix, iy;
      double y;
      for (iy = 0; iy < natom_line; iy++) {
	y = -0.5*systemsize + (iy+ 0.5) * delta;
	for (ix = 0; ix < natom_line; ix++) {
	  pos[0] = -0.5*systemsize + (ix + 0.5) * delta;
	  pos[1] = y;
	  pos+=DIMENSION;
	}
      }
    } else if (3 == DIMENSION) {
      int ix, iy, iz;
      double y,z;
      for (iz = 0; iz < natom_line; iz++) {
	z = -0.5*systemsize + (iz + 0.5) * delta;
	for (iy = 0; iy < natom_line; iy++) {
	  y = -0.5*systemsize + (iy+ 0.5) * delta;
	  for (ix = 0; ix < natom_line; ix++) {
	    pos[0] = -0.5*systemsize + (ix + 0.5) * delta;
	    pos[1] = y;
	    pos[2] = z;	    
	    pos+=DIMENSION;
	  }
	}
      }
    } else {
      fprintf(stderr, "wrong dimension %d\n", DIMENSION);
      return 1;
    }

    vel = data->vel;
    mass = data->mass;
    for (i = 0; i < data->natoms; i++) {
      for (dim = 0; dim < DIMENSION; dim++) {
	vel[dim] = 0.0;
	mass[dim] = 1.0;
      }
      vel += DIMENSION;
      mass += DIMENSION;
    }
    if (2==data->natoms && 1==DIMENSION) {/* to compare with argon.c result */
      data->pos[0] = 0.0; data->pos[1] = 1.0;
    }
  }

  compute_v_vp = CVVP[data->switch_method];
  
  /* compute_v_vp = C0_compute_v_vp; */
  
  if (SM_DC == data->switch_method) {
    data->switch_distance = data->cutoff;
  }

  {
    double inv_r = 1.0 / data->switch_distance;
    double inv_r2 = inv_r * inv_r;
    double inv_r6 = inv_r2 * inv_r2 * inv_r2;
    data->LJs    =  4.0 * inv_r6 * (inv_r6 - 1.0);
    data->LJs_p  = -4.0 * inv_r6 * inv_r  * ( 12.0*inv_r6 -  6.0);
    data->LJs_pp =  4.0 * inv_r6 * inv_r2 * (156.0*inv_r6 - 42.0);
  }

  return 0;
}


int force_poten_eval(Data *data)
{
  double systemsize = data->systemsize;
  const double cutoff2 = data->cutoff * data->cutoff;
  const double switch_dist2 = data->switch_distance * data->switch_distance;
  double * const force = data->f;
  const double * const pos = data->pos;
  double dist[DIMENSION], fij[DIMENSION];
  const double *pi, *pj;
  double *fi, *fj;
  double esum, peij, neg_vp;
  double r, r2, inv_r;
  const int natoms = data->natoms;
  int i, j;
  int dim;

  /* clear */
  memset(force, 0, data->natoms * DIMENSION * sizeof(double)); 
  esum = 0.0;
  for (i = 1; i < natoms; i++) {
    fi = force + i*DIMENSION;
    pi = pos   + i*DIMENSION;
    for (j = 0; j < i; j++) {
      fj = force + j*DIMENSION;
      pj = pos   + j*DIMENSION;
      compute_dist2(systemsize, pi, pj, dist, &r2);
      if (r2 < cutoff2) {
	r = sqrt(r2);
	inv_r = 1.0 / r;
	if (r2 < switch_dist2) {  /* normal */
	  LJ_poten_force(dist, inv_r, &peij, fij);
	} else {  /* switching region */
	  compute_v_vp(data, r, &peij, &neg_vp);
	  for (dim = 0; dim < DIMENSION; dim++) {
	    fij[dim] =  neg_vp * dist[dim] * inv_r;
	  }
	}
	esum += peij;
        for (dim = 0; dim < DIMENSION; dim++) {
	  fi[dim] += fij[dim];
	  fj[dim] -= fij[dim];
        }
      } else {
	; /* do nothing */
      }
    }    
  }

  data->pe = esum;
  return 0;
}


static void compute_dist2(const double systemsize,
			  const double pos1[DIMENSION], 
			  const double pos2[DIMENSION],
			  double dist[DIMENSION], double *r2) 
{
  int dim;
  double rsqr = 0.0;
  for (dim = 0; dim < DIMENSION; dim++) {
    dist[dim] = pos1[dim] - pos2[dim];
    BOUND(dist[dim], systemsize);
    rsqr += dist[dim] * dist[dim];
  }
  *r2 = rsqr;
}  

static void LJ_poten_force(const double dist[DIMENSION], const double inv_r, 
			   double *poten, double fij[DIMENSION])
{
  double inv_r2 = inv_r * inv_r;
  double inv_r6 = inv_r2 * inv_r2 * inv_r2;
  double neg_vp = 4.0 * inv_r6 * inv_r * (12.0*inv_r6 - 6.0);
  int dim;
  for (dim = 0; dim < DIMENSION; dim++) {
    fij[dim] =  neg_vp * dist[dim] * inv_r;
  }
  *poten =  4.0 * inv_r6 * (inv_r6 - 1.0);
}


static void DC_compute_v_vp(Data *data, double r, double *pe, double *neg_vp)
{
  assert(data && pe && neg_vp); r = 0; /* stupid thing to fool compiler */
  fprintf(stderr, "should not reach here\n");
}


static void C0_compute_v_vp(Data *data, double r, double *pe, double *neg_vp)
{
  register double switch_distance = data->switch_distance;
  register double cutoff = data->cutoff;
  *pe = data->LJs * (cutoff - r) / (cutoff - switch_distance);
  *neg_vp = data->LJs / (cutoff - switch_distance); 
}

static void C1_compute_v_vp(Data *data, double r, double *pe, double *neg_vp)
{
  const double switch_distance = data->switch_distance;
  const double cutoff = data->cutoff;
  double rleft = r - switch_distance;
  double rright = cutoff - r;
  double inv_rd = 1.0 / (cutoff - switch_distance);
  double w = data->LJs * (3.0*rleft + rright) * inv_rd + data->LJs_p * rleft;
  double wp = data->LJs * 2.0 * inv_rd + data->LJs_p;
  *pe = (rright*rright)*(inv_rd*inv_rd) * w;
  *neg_vp = (w*2.0 - wp*rright) * rright * (inv_rd*inv_rd);
}

static void C2_compute_v_vp(Data *data, double r, double *pe, double *neg_vp)
{
  const double switch_distance = data->switch_distance;
  const double cutoff = data->cutoff;
  double rleft = r - switch_distance;
  double rright = cutoff - r;
  double inv_rd = 1.0 / (cutoff - switch_distance);
  double lratio = rleft  * inv_rd;
  double rratio = rright * inv_rd;
  double w  = data->LJs    * (1.0 + 3.0 * lratio * (1.0 + 2.0*lratio)) + 
              data->LJs_p  * (1.0 + 3.0 * lratio) * rleft + 
              data->LJs_pp * 0.5 * rleft * rleft;
  double wp = data->LJs    * (3.0 + 12.0 * lratio) * inv_rd +
              data->LJs_p  * (1.0 +  6.0 * lratio) +
              data->LJs_pp * rleft;
  *pe = rratio*rratio*rratio * w;
  *neg_vp = (w*3.0*inv_rd - wp*rratio) * rratio * rratio;
}