#include "dump2analysis.h"
/* DOCUMENT */
/* -m gr */
/* lammpstrjから2原子間の動径分布関数を計算する */
/* またカットオフ距離内の平均結合距離を求める */
/* A原子(-a, -x, -s)とB原子(-c, -z, -u)を指定する。 */
/* --rmaxは動径分布関数を計算する最大距離 default:10.005 */
/* --rminは動径分布関数を計算する最小距離 default:0.01 */
/* --drは動径分布関数を計算する刻み距離 default:0.01 */
/* --rcutは結合距離を計算するカットオフ default:2.3 */

void ErrorGr(){
  printf("Required arguments atoms A and atoms B\n");
  printf("(-a, -x, -s) atom A\n");
  printf("(-b, -y, -t) atom B\n");
  printf("------------------------------------------------------------\n");
  printf("Optional arguments for Gr analysis:\n");
  printf("--dr distance step for g(r)                     [0.01]\n");
  printf("--rmax maximum distance for g(r)                [10.005]\n");
  printf("--rmin minimum distance for g(r)                [0.005]\n");
  printf("--rcut cut-off for alculation of mean distance  [2.3]\n");
  printf("------------------------------------------------------------\n");  
  printf("Example:\n");
  printf("dump2analysis -m gr -x Si -y O -i hoge.lammpstrj -o hoge.gr\n");
  exit(0);
}

void OutputGr(PARAM *p, double *gr, double *cn, ATOMS *a, ATOMS *b,
	      double ave, double std, int pairs){
  int i, RDATA;
  double rp, rcn=0;
  FILE *fp;

  RDATA = floor((RMAX-RMIN)/DR);
  fp = fopen(p->outfile, "w");
  fprintf(fp, "# Total step: %d\n", p->steps);
  /* fprintf(fp, "# atom-A:"); */
  /* for(i=0; i<a->atoms; i++) */
  /*   fprintf(fp, " %s(%d)", a->elem[i], a->id[i]); */
  /* fprintf(fp, "\n"); */
  /* fprintf(fp, "# atom-B:"); */
  /* for(i=0; i<b->atoms; i++) */
  /*   fprintf(fp, " %s(%d)", b->elem[i], b->id[i]); */
  /* fprintf(fp, "\n"); */
  fprintf(fp, "# (Rmin, Rmax, dr)=(%g, %g, %g)\n", RMIN, RMAX, DR);
  fprintf(fp, "# RDATA=%d\n", RDATA);
  fprintf(fp, "\n");
  fprintf(fp, "# Total pairs: %d\n", pairs);
  fprintf(fp, "# Rcut : %lf\n", RCUT);
  fprintf(fp, "# distance: %lf +/- %lf\n", ave, std);
  
  fprintf(fp, "# %10s %12s %12s\n", "r", "g(r)", "cn(r)");
  fprintf(fp, "%12g %12g %12g\n", 0.0, 0.0, 0.0);
  for(i=0; i<RDATA; i++){
    rp = RMIN + (double)i*DR + 0.5*DR;
    rcn = rcn + (double)cn[i];
    fprintf(fp, "%12g ", rp);
    fprintf(fp, "%12g ", gr[i]/(double)p->steps);
    fprintf(fp, "%12g ", rcn/(double)p->steps);
    fprintf(fp, "\n");
  }
  fclose(fp);
}

void EstimateGr(PARAM *param){
  FILE *f;
  HEAD head;
  int i, j, k, l, p, s, RDATA;
  double *gr, *cn, *cr;
  double volume, bunbo, sflag = 0.0;
  double xa, ya, za, xb, yb, zb, xp, yp, zp, r;
  double ave = 0.0, std = 0.0, *d = NULL;
  int rp, pairs=0;
  ATOMS a, b;

  if (ArgCheckInputOutput(param) != 0) ErrorGr();
  SetAtomsSteps(param);
  if (CheckArgAtomSelect(param) != 3) ErrorGr();

  
  Allocate(&a, param->atoms);
  Allocate(&b, param->atoms);
  f = fopen(param->infile, "r");
  RDATA = floor((RMAX-RMIN)/DR);
  gr = malloc(sizeof(double)*RDATA);
  cn = malloc(sizeof(double)*RDATA);
  cr = malloc(sizeof(double)*RDATA);
  
  /* grの計算 */
  for (i=0; i<RDATA; i++) gr[i] = 0.0;
  for (s=0; s<param->steps; s++){
    if (s%10 == 0) {
      printf("Calculating Gr step: %d\r", s);
      fflush(stdout);
    }
    /* 初期化 */
    for (i=0; i<RDATA; i++) cr[i] = 0.0;
    /* 原子の選択 */
    GetData(f, param, &head, &a, &b, NULL);
    /* 同一原子かのチェック */
    if (a.atoms == b.atoms){
      sflag = 1.0;
      for (i=0; i<a.atoms; i++)
	if (a.id[i] != b.id[i]) sflag = 0.0;
    }
    volume = head.A[0][0] * head.A[1][1] * head.A[2][2];
    for(i=0; i<a.atoms; i++){
      /* 絶対座標に変換 */
      xa = a.x[i]*head.A[0][0] + a.y[i]*head.A[1][0] + a.z[i]*head.A[2][0];
      ya = a.x[i]*head.A[0][1] + a.y[i]*head.A[1][1] + a.z[i]*head.A[2][1];
      za = a.x[i]*head.A[0][2] + a.y[i]*head.A[1][2] + a.z[i]*head.A[2][2];
      /* x-1, x+1, y-1,y+1, z-1, z+1 のb原子を作る */
      for (j=-1; j<2; j++){
	for (k=-1; k<2; k++){
	  for (l=-1; l<2; l++){
	    for (p=0; p<b.atoms; p++){
	      xp = b.x[p] + j;
	      yp = b.y[p] + k;
	      zp = b.z[p] + l;
	      xb = xp*head.A[0][0] + yp*head.A[1][0] + zp*head.A[2][0];
	      yb = xp*head.A[0][1] + yp*head.A[1][1] + zp*head.A[2][1];
	      zb = xp*head.A[0][2] + yp*head.A[1][2] + zp*head.A[2][2];
	      r = sqrt(pow((xa-xb), 2) + pow((ya-yb), 2) + pow((za-zb), 2));
	      if (RMIN < r && r < RMAX) {
		rp = (int)floor((r-RMIN)/DR);
		cr[rp]++;
		cn[rp] = cn[rp] + 1.0/(double)a.atoms;
	      }
	      if (0 < r && r < RCUT){
		d = realloc(d, sizeof(double) * (pairs+1));
		d[pairs] = r;
		pairs++;
	      }
	    }
	  }
	}
      }
    }
    /* 1ステップ毎にgrの計算 */
    for (i=0; i<RDATA; i++){
      if(cr[i] != 0){    
	r = RMIN + (double)i*DR + 0.5*DR;
	bunbo = 4.0 * M_PI * r*r * DR * (a.atoms) * (b.atoms - sflag);
	gr[i] = gr[i] + (double)cr[i] * volume / bunbo;
      }
    }
  }
  printf("\n");
  fclose(f);
  Deallocate(&a, param->atoms);
  Deallocate(&b, param->atoms);
  /* 平均値の計算 */
  for (i=0; i<pairs; i++) ave = ave + d[i];
  ave = ave/pairs;
  /* 標準偏差の計算 */
  for (i=0; i<pairs; i++) std = std + (d[i] - ave)*(d[i] - ave);
  std = sqrt(std/pairs);
  /* 出力 */
  OutputGr(param, gr, cn, &a, &b, ave, std, pairs);
  free(gr);
  free(cn);
  free(cr);
  free(d);
}