#include "dump2analysis.h"
/* DOCUMENT */
/* -m rotcorr */
/* 水分子の回転相関時間を計算する。 */
/* 日本化学会編「溶液の分子論的描像」 p.98 (1995) */
/* 室温での水の回転相関時間は1ps程度 */
/* 原子タイプで酸素原子を指定する */
/* idはsortされていてO,H,H  O,H.Hの並びになっていること */
/* 酸素原子の後に2つの水分子が続くIDを仮定している */
/* --shfitと--tauを小さくすれば統計が稼げる。ただしシフトが小さすぎる相関が影響を受ける */
/* --shiftでMSDが変わらないようパラメータを調整する。 */
/* A原子(-a, -x, -s)で酸素原子を指定する。 */
/* --shiftは速度相関関数を計算する時にシフトさせるステップ default:1 */
/* --tauは速度相関関数を計算する時間ステップ範囲。-1でMDの全ステップで計算する default:-1 */
/* --dtはMDの時間ステップでの時間(fs単位) default:1.0 */

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

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

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

}

void EstimateRotationCorr(PARAM *param){
  int i, j, p, step_counter=0;
  int matoms;
  FILE *f;
  ATOMS a;
  HEAD head;
  double **vx, **vy, **vz;   /* 水分子のベクトル */
  double vvx, vvy, vvz;
  double xa, ya, za;
  double *rcor;
  double vv, ab, norm_a, norm_b, cos2_theta;

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

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

  /* -1の場合は全ステップで計算する */  
  if (TAU == -1) {
    TAU = param->steps;
    TAU_SHIFT = param->steps;
  }

  /* 全ステップの水のベクトルを取得 */
  matoms = a.atoms/3;
  f = fopen(param->infile, "r");
  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)*matoms);
    vy[i] = malloc(sizeof(double)*matoms);
    vz[i] = malloc(sizeof(double)*matoms);
  }
  for (i=0; i<param->steps; i++){
    for (j=0; j<matoms; j++){
      vx[i][j] = 0.0;
      vy[i][j] = 0.0;
      vz[i][j] = 0.0;
    }
  }
  printf("Loading data...\n");
  for (i=0; i<param->steps; i++){
    GetData(f, param, &head, &a, NULL, NULL);
    for (j=0; j<matoms; j++){
      for (p=1; p<3; p++){
	vvx = a.x[3*j+p] - a.x[3*j];
	vvy = a.y[3*j+p] - a.y[3*j];
	vvz = a.z[3*j+p] - a.z[3*j];
	vvx = vvx - round(vvx);
	vvy = vvy - round(vvy);
	vvz = vvz - round(vvz);
	xa = vvx*head.A[0][0] + vvy*head.A[1][0] + vvz*head.A[2][0];
	ya = vvx*head.A[0][1] + vvy*head.A[1][1] + vvz*head.A[2][1];
	za = vvx*head.A[0][2] + vvy*head.A[1][2] + vvz*head.A[2][2];
	vx[i][j] = vx[i][j] + xa;
	vy[i][j] = vy[i][j] + ya;
	vz[i][j] = vz[i][j] + za;
      }
    }
  }
  fclose(f);
  
  /* 回転相関の計算 */
  rcor = malloc(sizeof(double) * TAU);
  for (i=0; i<TAU; i++) rcor[i] = 0.0;
  
  for (p=0; p+TAU <= param->steps; p=p+TAU_SHIFT){
    printf("Rotation correlation step ranging from %d to %d\n", p, p+TAU);
    for (i=0; i<TAU; i++){
      vv = 0;
      for (j=0; j<matoms; j++){
	ab = vx[p][j]*vx[p+i][j] + vy[p][j]*vy[p+i][j] + vz[p][j]*vz[p+i][j]; /* a dot b */
	norm_a = pow(vx[p][j],2) + pow(vy[p][j],2) + pow(vz[p][j],2);         /* |a|^2 */
	norm_b = pow(vx[p+i][j],2) + pow(vy[p+i][j],2) + pow(vz[p+i][j],2);   /* |b|^2 */
	cos2_theta = ab*ab/(norm_a * norm_b);   /* cos2theta */
	vv = vv + 0.5*(3*cos2_theta - 1);
      }
      rcor[i] = rcor[i] + vv/matoms;
    }
    step_counter++;
  }
  for (i=0; i<TAU; i++) rcor[i] = rcor[i]/step_counter;
  if (step_counter == 0){
    printf("--tau is too long, can not calculate Rotation correlation!!!\n");
    exit(0);
  }

  OutputRotationCorr(param, rcor);
}