#include "dump2analysis.h"
/* -m template_msd */
/* lammpstrjから指定した原子のvan Hove関数を時間シフトしながら計算する */
/* 原子タイプで指定しても最初のステップのidで原子を選択する */
/* --tauでジャンプさせるステップ数を指定する */
/* --shfitを指定すれば、複数のtau期間をサンプリングしvan Hoveを計算する */
/* --shfitを指定しなければ、tauが小さくても1セットのtauだけで計算する */
/* --shiftを指定する時は、--shftで関数が変わらないよう調整する。 */
/* --drでrの刻みを指定する default=0.01 */
/* A原子(-a, -x, -s)を指定する。 */
/* --shiftはMSDを計算する時にシフトさせるステップ default:1 */
/* --tauはval Hove関数のt。defaultでtotal step */
/* --dtはMDの時間ステップでの時間(fs単位) default:1.0 */

void ErrorVanHove(){
  printf("Required arguments atoms A\n");
  printf("(-a, -x, -s) atom A\n");
  printf("------------------------------------------------------------\n");
  printf("Optional arguments for van Hove analysis:\n");
  printf("--tau time period for van Hove function (-1: total steps  [-1]\n");
  printf("--shift time shift for van Hove function             [1]\n");
  printf("--dt time period for one step (fs)        [1.0]\n");
  printf("------------------------------------------------------------\n");  
  printf("Example:\n");
  printf("dump2analysis -m msd -x Li --dt 1.2 --tau 2500 --shift 100 -i hoge.lammpstrj -o hoge.msd\n");
  exit(0);
}

void OutputVanHove(PARAM *p, double *g){
  int i, RDATA;
  double r;
  FILE *f;

  RDATA = floor((RMAX-RMIN)/DR);

  f = fopen(p->outfile, "w");
  fprintf(f, "# step: %d\n", TAU);
  fprintf(f, "# time/fs: %lf\n", TAU*DT);
  for (i=0; i<RDATA; i++){
    r = RMIN + (double)i*DR + 0.5*DR;
    fprintf(f, "%lf %lf\n", r, g[i]);
  }
  fclose(f);

}

void EstimateVanHove(PARAM *param){
  int i, j, p, rp, step_counter=0;
  FILE *f;
  HEAD *head;
  ATOMS a;
  double **x, **y, **z;
  int RDATA;
  double *g;
  double xa, ya, za, dx, dy, dz, r;

  if (ArgCheckInputOutput(param) != 0) ErrorVanHove();
  SetAtomsSteps(param);
  if (CheckArgAtomSelect(param) != 1) ErrorVanHove();

  /* G(r,TAU)のメモリ確保 */
  RDATA = floor((RMAX-RMIN)/DR);
  g = malloc(sizeof(double)*RDATA);
  for (i=0; i<RDATA; i++) g[i] = 0.0;


  /* 選択元素が時間ステップで変わるとおかしくなるので元素選択は全てid */
  SetIDfromElemTtype(param);
  Allocate(&a, param->Aid[0]-1);
  f = fopen(param->infile, "r");

  /* メモリー確保 */
  x = malloc(sizeof(double*) * param->steps);
  y = malloc(sizeof(double*) * param->steps);
  z = malloc(sizeof(double*) * param->steps);
  for (i=0; i<param->steps; i++){
    x[i] = malloc(sizeof(double)*a.atoms);
    y[i] = malloc(sizeof(double)*a.atoms);
    z[i] = malloc(sizeof(double)*a.atoms);
  }
  head = malloc(sizeof(HEAD) * param->steps);

  /* -1の場合は全ステップのtでvan Hove関数を計算する */  
  if (TAU == -1) TAU = param->steps;
  if (TAU_SHIFT == -1) TAU_SHIFT = param->steps;
  
  /* 分率座標のデータを全ステップで取得する */
  printf("Loading data...\n");
  for (i=0; i<param->steps; i++){
    GetData(f, param, &head[i], &a, NULL, NULL);
    for (j=0; j<a.atoms; j++){
      x[i][j] = a.x[j] - floor(a.x[j]);
      y[i][j] = a.y[j] - floor(a.y[j]);
      z[i][j] = a.z[j] - floor(a.z[j]);
    }
  }
  Unwrap(param->steps, a.atoms, x, y, z);
  

  /* 絶対座標に変換 */  
  for (i=0; i<param->steps; i++){
    for (j=0; j<a.atoms; j++){
      xa = x[i][j]*head[i].A[0][0] + y[i][j]*head[i].A[1][0] + z[i][j]*head[i].A[2][0];
      ya = x[i][j]*head[i].A[0][1] + y[i][j]*head[i].A[1][1] + z[i][j]*head[i].A[2][1];
      za = x[i][j]*head[i].A[0][2] + y[i][j]*head[i].A[1][2] + z[i][j]*head[i].A[2][2];
      x[i][j] = xa;
      y[i][j] = ya;
      z[i][j] = za;
    }
  }

  /* van Hove関数の計算 r, a, b, c*/
  step_counter = 0;
  for (p=0; p+TAU+TAU_SHIFT < param->steps; p=p+TAU_SHIFT){
    printf("VanHove step range: from %d to %d\r", p, p+TAU);
    fflush(stdout);
    for (j=0; j<a.atoms; j++){
      dx = x[p+TAU][j] - x[p][j];
      dy = y[p+TAU][j] - y[p][j];
      dz = z[p+TAU][j] - z[p][j];
      r = sqrt(dx*dx + dy*dy + dz*dz);
      if (RMIN < r && r < RMAX) {
      	rp = (int)floor((r-RMIN)/DR);
      	g[rp] = g[rp] + 1.0;
      }
    }
    step_counter++;
  }
  for (i=0; i<RDATA; i++) g[i] = g[i]/step_counter/a.atoms;
  
  if (step_counter == 0){
    printf("--tau is too long, can not calculate VanHove!!!\n");
    exit(0);
  }
  
  OutputVanHove(param, g);
}