#include "dump2analysis.h"
/* DOCUMENT */
/* -m msd */
/* lammpstrjから指定した原子の平均二乗変位(MSD)を時間シフトしながら計算する */
/* 原子タイプで指定しても最初のステップのidで原子を選択する */
/* --shfitと--tauを小さくすればMSDの統計が稼げる。ただしシフトが小さすぎるMSDが変化する。 */
/* --shiftでMSDが変わらないようパラメータを調整する。 */
/* 教科書Pを参照 */
/* A原子(-a, -x, -s)を指定する。 */
/* --shiftはMSDを計算する時にシフトさせるステップ default:1 */
/* --tauはMSDを計算する時間ステップ範囲。-1でMDの全ステップでMSDを計算する default:-1 */
/* --dtはMDの時間ステップでの時間(fs単位) default:1.0 */

void ErrorMSD(){
  printf("Required arguments atoms A\n");
  printf("(-a, -x, -s) atom A\n");
  printf("------------------------------------------------------------\n");
  printf("Optional arguments for MSD analysis:\n");
  printf("--shift time shift for MSD                [1]\n");
  printf("--tau time period for MSD (-1: all steps  [-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 OutputMSD(PARAM *p, double **msd, ATOMS *a, int step){
  int i;
  double t;
  FILE *f;

  f = fopen(p->outfile, "w");
  fprintf(f, "# Atomic list: ");
  for (i=0; i<a->atoms; i++) fprintf(f, "%s(%d) ", a->elem[i], a->id[i]);
  fprintf(f, "\n");
  fprintf(f, "# TAU: %d\n", TAU);
  fprintf(f, "# TAU_SHIFT: %d\n", TAU_SHIFT);
  fprintf(f, "# total TAU step: %d\n", step);
  fprintf(f, "# atoms: %d\n", a->atoms);
  fprintf(f, "# %12s ", "time/fs");
  fprintf(f, "%14s ", "msd/Ang.2");
  fprintf(f, "%14s ", "a/Ang.2");
  fprintf(f, "%14s ", "b/Ang.2");
  fprintf(f, "%14s ", "c/Ang.2");
  fprintf(f, "\n");
  for (i=0; i<TAU; i++){
    t = i * (double)DT;
    fprintf(f, "%14lf ", t);
    fprintf(f, "%14e ", msd[0][i]);
    fprintf(f, "%14e ", msd[1][i]);
    fprintf(f, "%14e ", msd[2][i]);
    fprintf(f, "%14e ", msd[3][i]);
    fprintf(f, "\n");
  }
  fclose(f);

}

void EstimateMSD(PARAM *param){
  int i, j, k, p, step_counter=0;
  FILE *f;
  HEAD *head;
  double A[3][3];
  ATOMS a;
  double **xx, **yy, **zz;
  double **msd; /* r, a, b, c */
  double xp, yp, zp;
  double r, sum[4], dx, dy, dz;

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

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

  /* -1の場合は全ステップでMSDを計算する */
  if (TAU == -1) {
    TAU = param->steps;
    TAU_SHIFT = param->steps;
  }
  
  /* steps * atoms * (x, y, z)のメモリ取得 */
  xx = malloc(sizeof(double*) * param->steps);
  yy = malloc(sizeof(double*) * param->steps);
  zz = malloc(sizeof(double*) * param->steps);
  for(i=0; i<param->steps; i++){
    xx[i] = malloc(sizeof(double) * a.atoms);
    yy[i] = malloc(sizeof(double) * a.atoms);
    zz[i] = malloc(sizeof(double) * a.atoms);
  }

  /* 分率座標のデータを全ステップで取得する */
  printf("Loading data...\n");
  fflush(stdout);
  for (i=0; i<param->steps; i++){
    GetData(f, param, &head[i], &a, NULL, NULL);
    /* xx, yy, zzは分率座標 */
    for (j=0; j<a.atoms; j++){
      xx[i][j] = a.x[j];
      yy[i][j] = a.y[j];
      zz[i][j] = a.z[j];
    }
  }
  /* 分率座標のwarapを解く */
  Unwrap(param->steps, a.atoms, xx, yy, zz);
  
  /* msdの計算 r, a, b, c*/
  /* メモリ確保と初期化 */
  msd = malloc(sizeof(double*) * 4);
  for(i=0; i<4; i++) msd[i] = malloc(sizeof(double) * TAU);
  for (k=0; k<4; k++)
    for (i=0; i<TAU; i++) msd[k][i] = 0.0;
  
  for (p=0; p+TAU <= param->steps; p=p+TAU_SHIFT){
    printf("MSD step range: from %d to %d\n", p, p+TAU);
    fflush(stdout);
    for (i=0; i<TAU; i++){
      for (k=0; k<4; k++) sum[k] = 0.0;
      for (j=0; j<a.atoms; j++){
	dx = xx[p+i][j] - xx[p][j];
	dy = yy[p+i][j] - yy[p][j];
	dz = zz[p+i][j] - zz[p][j];
	/* nptのための処理。正しいか検討すること */
	A[0][0] = (head[p+i].A[0][0] + head[p].A[0][0])*0.5;
	A[0][1] = (head[p+i].A[0][1] + head[p].A[0][1])*0.5;
	A[0][2] = (head[p+i].A[0][2] + head[p].A[0][2])*0.5;
	A[1][0] = (head[p+i].A[1][0] + head[p].A[1][0])*0.5;
	A[1][1] = (head[p+i].A[1][1] + head[p].A[1][1])*0.5;
	A[1][2] = (head[p+i].A[1][2] + head[p].A[1][2])*0.5;
	A[2][0] = (head[p+i].A[2][0] + head[p].A[2][0])*0.5;
	A[2][1] = (head[p+i].A[2][1] + head[p].A[2][1])*0.5;
	A[2][2] = (head[p+i].A[2][2] + head[p].A[2][2])*0.5;
	
	xp = dx*A[0][0] + dy*A[1][0] + dz*A[2][0];
	yp = dx*A[0][1] + dy*A[1][1] + dz*A[2][1];
	zp = dx*A[0][2] + dy*A[1][2] + dz*A[2][2];
	
	r = pow(xp, 2) + pow(yp, 2) + pow(zp, 2);
	sum[0] = sum[0] + r;
	sum[1] = sum[1] + dx*dx*head[p+i].a*head[p+i].a;
	sum[2] = sum[2] + dy*dy*head[p+i].b*head[p+i].b;
	sum[3] = sum[3] + dz*dz*head[p+i].c*head[p+i].c;
      }
      msd[0][i] = msd[0][i] + sum[0]/(double)a.atoms;
      msd[1][i] = msd[1][i] + sum[1]/(double)a.atoms;
      msd[2][i] = msd[2][i] + sum[2]/(double)a.atoms;
      msd[3][i] = msd[3][i] + sum[3]/(double)a.atoms;
    }
    step_counter++;
  }

  if (step_counter == 0){
    printf("--tau is too long, can not calculate MSD!!!\n");
    exit(0);
  }

  for (k=0; k<4; k++)
    for (i=0; i<TAU; i++) msd[k][i] = msd[k][i]/(double)step_counter;
  OutputMSD(param, msd, &a, step_counter);
  free(head);
  for(i=0; i<param->steps; i++){
    free(xx[i]);
    free(yy[i]);
    free(zz[i]);
  }
  free(xx);
  free(yy);
  free(zz);
  for(i=0; i<4; i++) free(msd[i]);
  free(msd);
}