#include "dump2analysis.h"
/* DOCUMENT */
/* -m vcorr */
/* D=sum(v(t)v(0)dt)/3の関係がある。 */
/* 原子タイプで指定しても最初のステップのidで原子を選択する */
/* lammpstrjは少なくとも1.0fsで作成しないと誤差が大きい */
/* --shfitと--tauを小さくすれば統計が稼げる。ただしシフトが小さすぎる相関が影響を受ける */
/* --shiftでMSDが変わらないようパラメータを調整する。 */
/* 「コンピュータシミュレーションによる物質化学」P50を参照 */
/* A原子(-a, -x, -s)を指定する。 */
/* --shiftは速度相関関数を計算する時にシフトさせるステップ default:1 */
/* --tauは速度相関関数を計算する時間ステップ範囲。-1でMDの全ステップで計算する default:-1 */
/* --dtはMDの時間ステップでの時間(fs単位) default:1.0 */

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

void OutputVelocityCorr(PARAM *p, double **vcor){
  int i;
  double vsum = 0.0;
  double vv;
  FILE *f;

  f = fopen(p->outfile, "w");
  fprintf(f, "# %22s ", "time/s");
  fprintf(f, "%24s ", "v(t)v(0) m2/s2");
  fprintf(f, "%24s ", "sum(v(t)v(0))*dt/3 m2/s");
  fprintf(f, "\n");
  for (i=0; i<TAU; i++){
    vv = vcor[0][i] + vcor[1][i] + vcor[2][i];
    vsum = vsum + vv * DT * 1e-5;
    fprintf(f, "%24e ", (double)i*DT*1e-15);
    fprintf(f, "%24e ", vv*1e+10);
    fprintf(f, "%24e ", vsum/3.0);
    fprintf(f, "\n");
  }
  printf("%s was created.\n", p->outfile);
  fclose(f);

}

void EstimateVelocityCorr(PARAM *param){
  int i, j, p, step_counter=0;
  FILE *f;
  ATOMS a;
  HEAD head;
  double **vx, **vy, **vz;
  double **vcor;
  double vvx, vvy, vvz;

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

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

  /* -1の場合は全ステップで計算する */  
  if (TAU == -1) {
    TAU = param->steps;
    TAU_SHIFT = param->steps;
  }
  vcor = malloc(sizeof(double*) * 3);
  for (i=0; i<3; i++) vcor[i] = malloc(sizeof(double) * TAU);

  for (i=0; i<3; i++)
    for (j=0; j<TAU; j++) vcor[i][j] = 0.0;
  vx = malloc(sizeof(double*) * param->steps);
  vy = malloc(sizeof(double*) * param->steps);
  vz = malloc(sizeof(double*) * param->steps);
  for (i=0; i<param->steps; i++){
    vx[i] = malloc(sizeof(double) * a.atoms);
    vy[i] = malloc(sizeof(double) * a.atoms);
    vz[i] = malloc(sizeof(double) * a.atoms);
  }
  f = fopen(param->infile, "r");
  /* 全ステップの速度を読み込み */
  for (i=0; i<param->steps; i++){
    GetData(f, param, &head, &a, NULL, NULL);
    for (j=0; j<a.atoms; j++){
      vx[i][j] = a.vx[j];
      vy[i][j] = a.vy[j];
      vz[i][j] = a.vz[j];
    }
  }
  fclose(f);
  
  /* 速度相関の計算 */
  for (p=0; p+TAU <= param->steps; p=p+TAU_SHIFT){
    printf("Velocity correlation step ranging from %d to %d\n", p, p+TAU);
    for (i=0; i<TAU; i++){
      vvx = 0.0;
      vvy = 0.0;
      vvz = 0.0;
      for (j=0; j<a.atoms; j++){
	vvx = vvx + vx[p][j] * vx[p+i][j];
	vvy = vvy + vy[p][j] * vy[p+i][j];
	vvz = vvz + vz[p][j] * vz[p+i][j];
      }
      vcor[0][i] = vcor[0][i] + vvx/a.atoms;
      vcor[1][i] = vcor[1][i] + vvy/a.atoms;
      vcor[2][i] = vcor[2][i] + vvz/a.atoms;
    }
    step_counter++;
  }
  for (i=0; i<3; i++)
    for (j=0; j<TAU; j++) vcor[i][j] = vcor[i][j]/step_counter;
  if (step_counter == 0){
    printf("--tau is too long, can not calculate Velocity correlation!!!\n");
    exit(0);
  }

  OutputVelocityCorr(param, vcor);
}