#include "dump2analysis.h"
/* DOCUMENT */
/* -m cube_jump */
/* lammpstrjから指定したjump距離のグリッドデータ(cubeファイル)を作成する。 */
/* cubeファイルは、VESTAやVMDでisosurfaceやvolumes urfaceで可視化できる。 */
/* cubeファイルは、原子座標は全ステップの平均座標としている。 */
/* jumpを評価するステップはtauで与える */
/* x, y, zのvexcel分割数はnx, ny, nzで与える */

void ErrorCube_Jump(){
  printf("Required arguments atoms A\n");
  printf("(-a, -x, -s) atom A\n");
  printf("------------------------------------------------------------\n");
  printf("Optional arguments for cube analysis:\n");
  printf("--nx         numebr boxcel on X            [100]\n");
  printf("--ny         numebr boxcel on X            [100]\n");
  printf("--nz         numebr boxcel on X            [100]\n");
  printf("--tau        jump period                   [1]\n");
  printf("------------------------------------------------------------\n");  
  printf("Example:\n");
  printf("dump2analysis -m cube -x Li --cube_type 2 -i hoge.lammpstrj -o hoge.msd\n");
  exit(0);
}



void EstimateCube_Jump(PARAM *param){
  int i, j, k;
  int ix, iy, iz;
  FILE *f;
  HEAD *head;
  ATOMS a;
  int ***rho;
  double ***jump;
  double da, db, dc;
  double dx, dy, dz, r;
  double **xx, **yy, **zz;
  double **fx, **fy, **fz;

  /* メモリ確保 */
  rho = malloc(sizeof(int**) * NZ);
  jump = malloc(sizeof(double**) * NZ);
  for (i=0; i<NZ; i++){
    rho[i] = malloc(sizeof(int*) * NY);
    jump[i] = malloc(sizeof(double*) * NY);
  }
  for (i=0; i<NZ; i++){
    for (j=0; j<NY; j++){
      rho[i][j] = malloc(sizeof(int*) * NX);
      jump[i][j] = malloc(sizeof(double*) * NX);
    }
  }
  /* 初期化 */
  for (i=0; i<NZ; i++){
    for (j=0; j<NY; j++){
      for (k=0; k<NX; k++){
	rho[i][j][k] = 0;
	jump[i][j][k] = 0.0;
      }
    }
  }

  /* -1の場合は全ステップでMSDを計算する */
  if (TAU == -1) TAU = 1;

  /* エラーのチェック */
  if (ArgCheckInputOutput(param) != 0) ErrorCube_Jump();
  SetAtomsSteps(param);
  
  /* 選択元素が変わるとおかしくなるので元素選択は全てid */
  SetIDfromElemTtype(param);
  Allocate(&a, param->Aid[0]-1); /* 現在のステップ */

  /* steps * atoms * (x, y, z)のメモリ取得 */
  xx = malloc(sizeof(double*) * param->steps);
  yy = malloc(sizeof(double*) * param->steps);
  zz = malloc(sizeof(double*) * param->steps);
  fx = malloc(sizeof(double*) * param->steps);
  fy = malloc(sizeof(double*) * param->steps);
  fz = 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);
    fx[i] = malloc(sizeof(double) * a.atoms);
    fy[i] = malloc(sizeof(double) * a.atoms);
    fz[i] = malloc(sizeof(double) * a.atoms);
  }
  /* ヘッダ用のメモリ取得 */
  head = malloc(sizeof(HEAD) * param->steps);

  /* ファイルオープン */
  f = fopen(param->infile, "r");

  /* 分率座標を全ステップで取得する xx, yy, zz*/
  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++){
      /* 分率座標 fx, fy, fz */
      fx[i][j] = a.x[j] - floor(a.x[j]);
      fy[i][j] = a.y[j] - floor(a.y[j]);
      fz[i][j] = a.z[j] - floor(a.z[j]);
    }
  }
  
  /* fx, fy, fzをunwrapする */
  Unwrap(param->steps, a.atoms, fx, fy, fz);
  /* fx, fy, fzから絶対座標を生成 */
  for (i=0; i<param->steps; i++){
    for (j=0; j<a.atoms; j++){
      xx[i][j] = fx[i][j]*head[i].A[0][0] + fy[i][j]*head[i].A[1][0] + fz[i][j]*head[i].A[2][0];
      yy[i][j] = fx[i][j]*head[i].A[0][1] + fy[i][j]*head[i].A[1][1] + fz[i][j]*head[i].A[2][1];
      zz[i][j] = fx[i][j]*head[i].A[0][2] + fy[i][j]*head[i].A[1][2] + fz[i][j]*head[i].A[2][2];
      fx[i][j] = fx[i][j] - floor(fx[i][j]);
      fy[i][j] = fy[i][j] - floor(fy[i][j]); 
      fz[i][j] = fz[i][j] - floor(fz[i][j]);
    }
  }
  
  /* 全ステップのデータを解析する */
  da = 1.0/(double)NX;
  db = 1.0/(double)NY;
  dc = 1.0/(double)NZ;
  for (i=0; i<param->steps - TAU; i++){
    printf("\rCalculating step: %d/%d", i, param->steps);
    fflush(stdout);
    for (j=0; j<a.atoms; j++){
      ix = (int)floor(fx[i][j]/da);
      iy = (int)floor(fy[i][j]/db);
      iz = (int)floor(fz[i][j]/dc);
      dx = xx[i+TAU][j] - xx[i][j];
      dy = yy[i+TAU][j] - yy[i][j];
      dz = zz[i+TAU][j] - zz[i][j];
      r = sqrt(dx*dx + dy*dy + dz*dz);
      rho[iz][iy][ix]++;
      jump[iz][iy][ix] = jump[iz][iy][ix] + r;
    }
  }
  printf("\n");
  
  for (i=0; i<NZ; i++){
    for (j=0; j<NY; j++){
      for (k=0; k<NX; k++){
	/* 訪れた回数で規格化 */
	if (rho[i][j][k] != 0){
	  jump[i][j][k] = jump[i][j][k]/(double)rho[i][j][k];
	}
      }
    }
  }
  fclose(f);
  CUBE_TYPE = 4;
  OutputCube(param, jump);
}