#include "dump2analysis.h"

int GetHeader(FILE *fp, HEAD *head){
  int chk;
  char buf[1024];
  double xlo = 0.0, xhi = 0.0, lx;
  double ylo = 0.0, yhi = 0.0, ly;
  double lz;

  /* ITEM: TIMESTEP */
  if (fgets(buf, sizeof(buf), fp)==NULL) return -1;
  if (fgets(buf, sizeof(buf), fp)==NULL) return -1;
  sscanf(buf, "%d", &head->step);
  /* ITEM: NUMBER OF ATOMS */
  fgets(buf, sizeof(buf), fp);
  fgets(buf, sizeof(buf), fp);
  sscanf(buf, "%d", &head->atoms);
  /* ITEM: BOX BOUNDS xy xz yz pp pp pp */
  fgets(buf, sizeof(buf), fp);
  fgets(buf, sizeof(buf), fp);
  chk = sscanf(buf, "%lf %lf %lf",&head->x_low, &head->x_hi, &head->xy);
  if (chk == 2) {
    head->A[0][0] = head->x_hi - head->x_low;
    head->A[0][1] = 0.0;
    head->A[0][2] = 0.0;
    head->a = head->x_hi - head->x_low;
  }
  fgets(buf, sizeof(buf), fp);
  chk = sscanf(buf, "%lf %lf %lf",&head->y_low, &head->y_hi, &head->xz);
  if (chk == 2) {
    head->A[1][0] = 0.0;
    head->A[1][1] = head->y_hi - head->y_low ;
    head->A[1][2] = 0.0;
    head->b = head->y_hi - head->y_low;
  }
  fgets(buf, sizeof(buf), fp);
  chk = sscanf(buf, "%lf %lf %lf", &head->z_low, &head->z_hi, &head->yz);
  if (chk == 2) {
    head->A[2][0] = 0.0;
    head->A[2][1] = 0.0;
    head->A[2][2] = head->z_hi - head->z_low ;
    head->c = head->z_hi - head->z_low;
  }
  /* ITEM: ATOMS id element x y z vx vy vz fx fy fz */
  fgets(buf, sizeof(buf), fp);
  if (chk == 3){ /* triclinic */
    /* x */
    if (xlo > head->xy) xlo = head->xy;
    if (xlo > head->xz) xlo = head->xz;
    if (xlo > (head->xy+head->xz)) xlo = head->xy+head->xz;
    if (xhi < head->xy) xhi = head->xy;
    if (xhi < head->xz) xhi = head->xz;
    if (xhi < (head->xy+head->xz)) xhi = head->xy+head->xz;
    xlo = head->x_low - xlo;
    xhi = head->x_hi - xhi;
    lx = xhi - xlo;
    head->a = lx;
    /* y */
    if (ylo > head->yz) ylo = head->yz;
    if (yhi < head->yz) yhi = head->yz;
    ylo = head->y_low - ylo;
    yhi = head->y_hi - yhi;
    ly = yhi - ylo;
    head->b = sqrt(ly*ly + head->xy*head->xy);
    /* z */
    lz = head->z_hi - head->z_low;
    head->c = sqrt(lz*lz + head->xz*head->xz + head->yz*head->yz );
    /* angle */
    head->alpha = acos((head->yz*ly+head->xy*head->xz)/
                       (head->b*head->c));
    head->beta = acos(head->xz/head->c);
    head->gamma = acos(head->xy/head->b);
    head->A[0][0] = head->a;
    head->A[0][1] = 0;
    head->A[0][2] = 0;
    head->A[1][0] = head->b*cos(head->gamma);
    head->A[1][1] = head->b*sin(head->gamma);
    head->A[1][2] = 0;
    head->A[2][0] = head->c*cos(head->beta);
    head->A[2][1] = head->c*(
                           cos(head->alpha)-cos(head->beta)*cos(head->gamma)
                           )/sin(head->gamma);
    head->A[2][2] = head->c*sqrt(
                               1+2*cos(head->alpha)*cos(head->beta)*cos(head->gamma)
                               - cos(head->alpha)*cos(head->alpha)
                               - cos(head->beta)*cos(head->beta)
                               - cos(head->gamma)*cos(head->gamma)
                               ) /sin(head->gamma);
  }
  else{ /* 直行系 */
    head->alpha = M_PI*0.5;
    head->beta = M_PI*0.5;
    head->gamma = M_PI*0.5;
  }
  /* 逆行列の計算 */
  double det_A;
  det_A = head->A[0][0]*head->A[1][1]*head->A[2][2];
  head->B[0][0] = (head->A[1][1]*head->A[2][2] - head->A[1][2]*head->A[2][1])/det_A;
  head->B[1][0] = (head->A[1][2]*head->A[2][0] - head->A[1][0]*head->A[2][2])/det_A;
  head->B[2][0] = (head->A[1][0]*head->A[2][1] - head->A[1][1]*head->A[2][0])/det_A;
  head->B[0][1] = 0.0;
  head->B[1][1] = (head->A[0][0]*head->A[2][2] - head->A[0][2]*head->A[2][0])/det_A;
  head->B[2][1] = (head->A[0][1]*head->A[2][0] - head->A[0][0]*head->A[2][1])/det_A;
  head->B[0][2] = 0.0;
  head->B[1][2] = 0.0;
  head->B[2][2] = (head->A[0][0]*head->A[1][1] - head->A[0][1]*head->A[1][0])/det_A;
  return 1;
}

void SetData(HEAD *head, ATOMS *a, ATOMS *b){
  double x, y, z;

  x = b->x[0];
  y = b->y[0];
  z = b->z[0];
  
  strcpy(a->elem[a->atoms], b->elem[0]);
  a->id[a->atoms] = b->id[0];
  a->type[a->atoms] = b->type[0];
  /* 絶対座標を分率座標へ */
  a->x[a->atoms] = x*head->B[0][0] + y*head->B[1][0] + z*head->B[2][0];
  a->y[a->atoms] = x*head->B[0][1] + y*head->B[1][1] + z*head->B[2][1];
  a->z[a->atoms] = x*head->B[0][2] + y*head->B[1][2] + z*head->B[2][2];
  a->x[a->atoms] = a->x[a->atoms] - floor(a->x[a->atoms]);
  a->y[a->atoms] = a->y[a->atoms] - floor(a->y[a->atoms]);
  a->z[a->atoms] = a->z[a->atoms] - floor(a->z[a->atoms]);
  a->vx[a->atoms] = b->vx[0];
  a->vy[a->atoms] = b->vy[0];
  a->vz[a->atoms] = b->vz[0];
  a->fx[a->atoms] = b->fx[0];
  a->fy[a->atoms] = b->fy[0];
  a->fz[a->atoms] = b->fz[0];
  a->atoms++;
}

void GetAllData(FILE *f, PARAM* p, HEAD *head, ATOMS *a){
  int i;
  char buf[1024];
  ATOMS l;
  Allocate(&l, 1);

  GetHeader(f, head);
  a->atoms = 0;
  for(i=0; i<p->atoms; i++){
    fgets(buf, sizeof(buf), f);
    sscanf(buf, "%d %d %s %lf %lf %lf %lf %lf %lf %lf %lf %lf",
	   &l.id[0], &l.type[0], l.elem[0],
	   &l.x[0], &l.y[0], &l.z[0],
	   &l.vx[0], &l.vy[0], &l.vz[0],
	   &l.fx[0], &l.fy[0], &l.fz[0]
	   );
    SetData(head, a, &l);
  }
  head->atoms = p->atoms;
}

void GetData(FILE *f, PARAM* p, HEAD *head, ATOMS *a, ATOMS *b, ATOMS *c){
  int i, j;
  char buf[1024];
  ATOMS l;
  Allocate(&l, 1);

  if (a != NULL) a->atoms = 0;
  if (b != NULL) b->atoms = 0;
  if (c != NULL) c->atoms = 0;
  GetHeader(f, head);
  for(i=0; i<head->atoms; i++){
    fgets(buf, sizeof(buf), f);
    sscanf(buf, "%d %d %s %lf %lf %lf %lf %lf %lf %lf %lf %lf",
	   &l.id[0], &l.type[0], l.elem[0],
	   &l.x[0], &l.y[0], &l.z[0],
	   &l.vx[0], &l.vy[0], &l.vz[0],
	   &l.fx[0], &l.fy[0], &l.fz[0]
	   );
    /* idで選択 */
    for (j=1; j<p->Aid[0]; j++)
      if (l.id[0] == p->Aid[j]) SetData(head, a, &l);
    /* 元素名で選択 */
    if (strcmp(p->Aelem, l.elem[0]) == 0) SetData(head, a, &l);
    /* 元素タイプ番号で選択 */
    if (p->Atype == l.type[0]) SetData(head, a, &l);

    if (b == NULL) continue;
    /* idで選択 */
    for (j=1; j<p->Bid[0]; j++)
      if (l.id[0] == p->Bid[j]) SetData(head, b, &l);
    /* 元素名で選択 */
    if (strcmp(p->Belem, l.elem[0]) == 0) SetData(head, b, &l);
    /* 元素タイプ番号で選択 */
    if (p->Btype == l.type[0]) SetData(head, b, &l);
    
    if (c == NULL) continue;    
    /* idで選択 */
    for (j=1; j<p->Cid[0]; j++)
      if (l.id[0] == p->Cid[j]) SetData(head, c, &l);
    /* 元素名で選択 */
    if (strcmp(p->Celem, l.elem[0]) == 0) SetData(head, c, &l);
    /* 元素タイプ番号で選択 */
    if (p->Ctype == l.type[0]) SetData(head, c, &l);
  }
}

void GetOneStepData(FILE *f, HEAD *head, ATOMS *a){
  int i;
  double x, y, z;
  char buf[1024];

  GetHeader(f, head);
  for(i=0; i<head->atoms; i++){
    fgets(buf, sizeof(buf), f);
    sscanf(buf, "%d %d %s %lf %lf %lf %lf %lf %lf %lf %lf %lf",
	   &a->id[i], &a->type[i], a->elem[i],
	   &x, &y, &z,
	   &a->vx[i], &a->vy[i], &a->vz[i],
	   &a->fx[i], &a->fy[i], &a->fz[i]);
    /* 絶対座標を分率座標へ  2017-12-27に修正 */
    a->x[i] = x*head->B[0][0] + y*head->B[1][0] + z*head->B[2][0];
    a->y[i] = x*head->B[0][1] + y*head->B[1][1] + z*head->B[2][1];
    a->z[i] = x*head->B[0][2] + y*head->B[1][2] + z*head->B[2][2];
    a->x[i] = a->x[i] - floor(a->x[i]);
    a->y[i] = a->y[i] - floor(a->y[i]);
    a->z[i] = a->z[i] - floor(a->z[i]);
  }
}

void SetIDfromElemTtype(PARAM *param){
  int i, j;
  ATOMS a;
  FILE *f;
  HEAD head;
  
  /* 構造体のメモリ確保 */
  Allocate(&a, param->atoms);
  /* ファイルオープン */
  f = fopen(param->infile, "r");
  GetOneStepData(f, &head, &a);
  /* 元素名で指定した場合はidに変換 */
  if (strcmp(param->Aelem, "") != 0){
    j = 0;
    for (i=0; i<a.atoms; i++){
      if (strcmp(a.elem[i], param->Aelem) == 0){
	param->Aid = realloc(param->Aid, sizeof(int) * (j + 2));
	param->Aid[j+1] = a.id[i];
	j++;
      }
    }
    param->Atype = -1;
    strcpy(param->Aelem, "");
    param->Aid[0] = j+1;
  }

  /* 元素タイプで指定した場合はidに変換 */
  if (param->Atype != -1){
    j = 0;
    for (i=0; i<a.atoms; i++){
      if (param->Atype == a.type[i]){
	param->Aid = realloc(param->Aid, sizeof(int) * (j + 2));
	param->Aid[j+1] = a.id[i];
	j++;
      }
    }
    param->Atype = -1;
    strcpy(param->Aelem, "");
    param->Aid[0] = j+1;
  }
  fclose(f);
}

int* SetIDfromMolecules(int *Aid, int n){
  int i, j;
  int *tmp;
  int atoms;

  /* 指定原子から連番で原子の数(n)だけ選択原子を増やす */
  atoms = (Aid[0]-1) * n;
  tmp = malloc(sizeof(int) * atoms + 1);
  tmp[0] = atoms + 1;
  for (i=1; i<Aid[0]; i++){
    tmp[(i-1)*n + 1] = Aid[i];
    for (j=0; j<n; j++){
      tmp[(i-1)*n + 1 + j + 1] = Aid[i] + j + 1;
    }
  }
  Aid = malloc(sizeof(int)*tmp[0]);
  Aid[0] = tmp[0];
  for (i=1; i<Aid[0]; i++)
    Aid[i] = tmp[i];
  return Aid;

}

void Unwrap(int steps, int atoms, double **x, double **y, double **z){
  int **nx, **ny, **nz;
  int i, j;

  nx = malloc(sizeof(int*) * steps);
  ny = malloc(sizeof(int*) * steps);
  nz = malloc(sizeof(int*) * steps);
  for (i=0; i < steps; i++){
    nx[i] = malloc(sizeof(int) * atoms);
    ny[i] = malloc(sizeof(int) * atoms);
    nz[i] = malloc(sizeof(int) * atoms);
  }
  for (i=0; i<steps; i++){
    for (j=0; j<atoms; j++){
      nx[i][j] = 0;
      ny[i][j] = 0;
      nz[i][j] = 0;
    }
  }
  for (i=1; i<steps; i++){
    for (j=0; j<atoms; j++){
      nx[i][j] = nx[i-1][j] + round(x[i-1][j] - x[i][j]);
      ny[i][j] = ny[i-1][j] + round(y[i-1][j] - y[i][j]);
      nz[i][j] = nz[i-1][j] + round(z[i-1][j] - z[i][j]);
    }
  }
  for (i=0; i<steps; i++){
    for (j=0; j<atoms; j++){
      x[i][j] = x[i][j] + (double)nx[i][j];
      y[i][j] = y[i][j] + (double)ny[i][j];
      z[i][j] = z[i][j] + (double)nz[i][j];
    }
  }

}