#include "dump2analysis.h"
/* DOCUMENT */
/* -m cube */
/* lammpstrjから指定したデータのグリッドデータ(cubeファイル)を作成する。 */
/* cubeファイルは、VESTAやVMDでisosurfaceやvolumes urfaceで可視化できる。 */
/* グリッドデータのパラメータは、頻度、速度、力、エネルギーを選択できる。 */
/* cubeファイル中での指定以外の原子座標は、全ステップの平均値とする */
/* --cube_type  1:指定原子の頻度 */
/* --cube_type  2:指定原子の速度ノルム */
/* --cube_type  3:指定原子の力ノルム */

void ErrorCube(){
  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("--cube_type  select cube parameter         [1]\n");
  printf("             1: density\n");
  printf("             2: velocity\n");
  printf("             3: force\n");
  printf("------------------------------------------------------------\n");  
  printf("Example:\n");
  printf("dump2analysis -m cube -x Li --cube_type 2 -i hoge.lammpstrj -o hoge.cube\n");
  exit(0);
}

void OutputCube(PARAM *param, double ***rho){
  int i, j, k, cnt=0;
  FILE *fp;
  ATOMS a;
  HEAD *head;
  int l, m;
  double **xx, **yy, **zz;
  double *x_merge, *y_merge, *z_merge;
  double x, y, z;
  double A[3][3];

  /* steps * atoms * (x, y, z)のメモリ取得 */
  head = malloc(sizeof(HEAD) * param->steps);
  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) * param->atoms);
    yy[i] = malloc(sizeof(double) * param->atoms);
    zz[i] = malloc(sizeof(double) * param->atoms);
  }
  x_merge = malloc(sizeof(double) * param->atoms);
  y_merge = malloc(sizeof(double) * param->atoms);
  z_merge = malloc(sizeof(double) * param->atoms);

  Allocate(&a, param->atoms);
  fp = fopen(param->infile, "r");
  /* 分率座標のデータを全ステップで取得する */
  printf("Calculating average position of all atoms...\n");
  fflush(stdout);
  for (i=0; i<param->steps; i++){
    GetAllData(fp, param, &head[i], &a);
    /* xx, yy, zzは[0-1]の分率座標 */
    for (j=0; j<param->atoms; j++){
      xx[i][j] = a.x[j];
      yy[i][j] = a.y[j];
      zz[i][j] = a.z[j];
    }
  }

  /* 分率座標でunwrap */
  Unwrap(param->steps, param->atoms, xx, yy, zz);
  fclose(fp);
  
  /* 時間平均を求めるための初期化 */
  for (i=0; i<param->atoms; i++){
    x_merge[i] = 0.0;
    y_merge[i] = 0.0;
    z_merge[i] = 0.0;
  }
  for (l=0; l<3; l++){
    for (m=0; m<3; m++){
      A[l][m] = 0.0;
    }
  }
  /* 時間平均を求めるための足し上げ */
  for (i=0; i<param->steps; i++){
    for (j=0; j<param->atoms; j++){
      x_merge[j] = x_merge[j] + xx[i][j];
      y_merge[j] = y_merge[j] + yy[i][j];
      z_merge[j] = z_merge[j] + zz[i][j];
    }
    for (l=0; l<3; l++){
      for (m=0; m<3; m++){
	A[l][m] = A[l][m] + head[i].A[l][m];
      }
    }
  }
  /* 時間平均する */
  for (i=0; i<param->atoms; i++){
    x_merge[i] = x_merge[i]/param->steps;
    y_merge[i] = y_merge[i]/param->steps;
    z_merge[i] = z_merge[i]/param->steps;
  }
  for (l=0; l<3; l++){
    for (m=0; m<3; m++){
      A[l][m] = A[l][m]/param->steps;
    }
  }
  fp = fopen(param->outfile, "w");
  /* cubeファイルの出力 Bohr単位にする */
  fprintf(fp, "Cubfile created from dump2analysis cube\n");
  if (CUBE_TYPE == 1) fprintf(fp, "Density: density\n");
  if (CUBE_TYPE == 2) fprintf(fp, "Density: velocity\n");
  if (CUBE_TYPE == 3) fprintf(fp, "Density: force\n");
  if (CUBE_TYPE == 4) fprintf(fp, "Density: jump\n");
  /* atoms, x_origin, y_origin, z_orgin */
  fprintf(fp, "%5d ", param->atoms);
  fprintf(fp, "%12.6f ", 0.0);
  fprintf(fp, "%12.6f ", 0.0);
  fprintf(fp, "%12.6f",  0.0);
  fprintf(fp, "\n");
  
  /* NX dx dy dz minus unit: Bohr unit */
  fprintf(fp, "%5d ", NX); /* a vector */
  fprintf(fp, "%12.6f ",  A[0][0]/NX/0.52917721067); 
  fprintf(fp, "%12.6f ",  A[0][1]/NX/0.52917721067);
  fprintf(fp, "%12.6f\n", A[0][2]/NX/0.52917721067);
  fprintf(fp, "%5d ", NY); /* b vector */  
  fprintf(fp, "%12.6f ",  A[1][0]/NY/0.52917721067);
  fprintf(fp, "%12.6f ",  A[1][1]/NY/0.52917721067);
  fprintf(fp, "%12.6f\n", A[1][2]/NY/0.52917721067);
  fprintf(fp, "%5d ", NZ); /* c vector */
  fprintf(fp, "%12.6f ",  A[2][0]/NZ/0.52917721067);
  fprintf(fp, "%12.6f ",  A[2][1]/NZ/0.52917721067);
  fprintf(fp, "%12.6f\n", A[2][2]/NZ/0.52917721067);
  /* 原子座標の出力 */
  for(i=0; i<param->atoms; i++){
    x_merge[i] = x_merge[i] - floor(x_merge[i]);
    y_merge[i] = y_merge[i] - floor(y_merge[i]);
    z_merge[i] = z_merge[i] - floor(z_merge[i]);
    x = x_merge[i]*A[0][0] + y_merge[i]*A[1][0] + z_merge[i]*A[2][0];
    y = x_merge[i]*A[0][1] + y_merge[i]*A[1][1] + z_merge[i]*A[2][1];
    z = x_merge[i]*A[0][2] + y_merge[i]*A[1][2] + z_merge[i]*A[2][2];
    x = x /0.52917721067;
    y = y /0.52917721067;
    z = z /0.52917721067;
    fprintf(fp, "%5d %12.6f ", Eleme2AtomicNumber(a.elem[i]), 0.0);
    fprintf(fp, "%12.6f %12.6f %12.6f\n", x, y, z);
  }
  for (i=0; i<NX; i++){
    for (j=0; j<NY; j++){
      for (k=0; k<NZ; k++){
  	fprintf(fp, " %13.5E", rho[k][j][i]);
  	if (cnt%6==5) fprintf(fp, "\n");
  	cnt++;
      }
    }
  }
}

void EstimateCube(PARAM *param){
  int i, j, k;
  int ix, iy, iz;
  FILE *f;
  HEAD head;
  ATOMS a;
  double ***rho;
  double ***vel;
  double ***force;
  double ***jump;
  double vel_norm, force_norm;
  double da, db, dc;

  /* メモリ確保 */
  rho = malloc(sizeof(double**) * NZ);
  vel = malloc(sizeof(double**) * NZ);
  force= malloc(sizeof(double**) * NZ);
  jump = malloc(sizeof(double**) * NZ);
  for (i=0; i<NZ; i++){
    rho[i] = malloc(sizeof(double*) * NY);
    vel[i] = malloc(sizeof(double*) * NY);
    force[i] = malloc(sizeof(double*) * NY);
    jump[i] = malloc(sizeof(double*) * NY);
  }
  for (i=0; i<NZ; i++){
    for (j=0; j<NY; j++){
      rho[i][j] = malloc(sizeof(double*) * NX);
      vel[i][j] = malloc(sizeof(double*) * NX);
      force[i][j] = malloc(sizeof(double*) * 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.0;
	vel[i][j][k] = 0.0;
	force[i][j][k] = 0.0;
	jump[i][j][k] = 0.0;
      }
    }
  }
  
  /* エラーのチェック */
  if (ArgCheckInputOutput(param) != 0) ErrorCube();
  SetAtomsSteps(param);
  /* 1: A原子のみを指定している */
  if (CheckArgAtomSelect(param) != 1) ErrorCube();

  /* 原子数を与えての構造体のメモリ確保 */
  Allocate(&a, param->atoms);

  /* ファイルオープン */
  f = fopen(param->infile, "r");
  
  /* ステップ毎にデータを取得する */
  da = 1.0/(double)NX;
  db = 1.0/(double)NY;
  dc = 1.0/(double)NZ;
  for (i=0; i<param->steps; i++){
    GetData(f, param, &head, &a, NULL, NULL);
    for (j=0; j<a.atoms; j++){
      a.x[j] = a.x[j] - floor(a.x[j]);
      a.y[j] = a.y[j] - floor(a.y[j]);
      a.z[j] = a.z[j] - floor(a.z[j]);
      ix = (int)floor(a.x[j]/da);
      iy = (int)floor(a.y[j]/db);
      iz = (int)floor(a.z[j]/dc);
      vel_norm = sqrt(a.vx[j]*a.vx[j] + a.vy[j]*a.vy[j] + a.vz[j]*a.vz[j]);
      force_norm = sqrt(a.fx[j]*a.fx[j] + a.fy[j]*a.fy[j] + a.fz[j]*a.fz[j]);
      rho[iz][iy][ix] = rho[iz][iy][ix] + 1.0;
      vel[iz][iy][ix] = rho[iz][iy][ix] + vel_norm;
      force[iz][iy][ix] = rho[iz][iy][ix] + force_norm;
    }
  }
  fclose(f);
  
  /* 訪れた回数で規格化 */
  for (i=0; i<NZ; i++){
    for (j=0; j<NY; j++){
      for (k=0; k<NX; k++){
	if (rho[i][j][k] == 0) continue;
	/* 訪れた回数で規格化 */
	vel[i][j][k] = vel[i][j][k]/rho[i][j][k];
	force[i][j][k] = force[i][j][k]/rho[i][j][k];
      }
    }
  }

  if (CUBE_TYPE == 1) OutputCube(param, rho);
  if (CUBE_TYPE == 2) OutputCube(param, vel);
  if (CUBE_TYPE == 3) OutputCube(param, force);
}