#include "dump2sq.h"

/* ####################################################################### */
/* 初期化 原子タイプ毎のパラメータ */
/* ####################################################################### */
void InitializeType(HEAD *head, DATA *data, TYPE *type){
  int i, j, k;
  int flag;
  char elem[100][8];
  int elem_n = 0;

  /* 原子タイプidの決定 */
  type->types = 0;
  for (i=0; i<head->atoms; i++){
    flag = 0;
    for (j=0; j<elem_n; j++){
      if (strcmp(elem[j], data[i].elem) == 0) flag = 1;
    }
    if (flag == 0){
      strcpy(elem[elem_n], data[i].elem);
      elem_n++;
    }
  }
  type->types = elem_n;
  type->elem = malloc(sizeof(char*) * type->types);
  for (i=0; i<type->types; i++) {
    type->elem[i] = malloc(sizeof(char) * 8);
    strcpy(type->elem[i], elem[i]);
  }

  /* 原子分率の計算 */
  /* また各原子タイプを割り当てなおす */
  /* trjで定義されているtype_idは無視。0から始まるように振り直す */
  type->ids = malloc(type->types * sizeof(int));
  for (i=0; i<type->types; i++) type->ids[i] = 0;
  for (i=0; i<head->atoms; i++) {
    for (j=0; j<type->types; j++){
      if (strcmp(type->elem[j], data[i].elem) == 0){
	type->ids[j]++;
	data[i].type_new = j;
      }
    }
  }
  /* 各原子タイプの分率 */
  type->ci = malloc(type->types * sizeof(double));
  for (i=0; i<type->types; i++)  type->ci[i] = (double)type->ids[i]/head->atoms;

  /* pairのidをセット */
  type->pairid = malloc(sizeof(int *) * type->types);
  for (i=0; i<type->types; i++)
    type->pairid[i] = malloc(sizeof(int) * type->types);
  k = 0;
  for (i=0; i<type->types; i++){
    for (j=i; j<type->types; j++){
      type->pairid[i][j] = k;
      type->pairid[j][i] = k;
      k++;
    }
  }
  /* pairのセット */
  type->pairs = type->types * (type->types + 1) * 0.5;
  type->cicj = malloc(sizeof(double) * type->pairs);
  type->typeij = malloc(sizeof(int*) * type->pairs);
  type->nij = malloc(sizeof(double*) * type->pairs);
  type->elemij = malloc(sizeof(char*) * type->pairs);
  for (i=0; i<type->pairs; i++) {
    type->typeij[i] = malloc(sizeof(int) * 2);
    type->nij[i] = malloc(sizeof(double) * 2);
    type->elemij[i] = malloc(sizeof(char) * 32);
  }
  k = 0;
  for (i=0; i<type->types; i++){
    for (j=i; j<type->types; j++){
      type->typeij[k][0] = i;
      type->typeij[k][1] = j;
      type->nij[k][0] = type->ids[i];
      type->nij[k][1] = type->ids[j];
      sprintf(type->elemij[k], "%s-%s", type->elem[i], type->elem[j]);
      type->cicj[k] = type->ci[i] * type->ci[j];
      k++;
    }
  }
  /* for (i=0; i<type->pairs; i++){ */
  /*   printf("%d ", type->typeij[i][0]); */
  /*   printf("%d ", type->typeij[i][1]); */
  /*   printf("%lf ", type->nij[i][0]); */
  /*   printf("%lf ", type->nij[i][1]); */
  /*   printf("%s ", type->elemij[i]); */
  /*   printf("%lf ", type->cicj[i]); */
  /*   printf("\n"); */
  /* } */
  /* for (i=0; i<head->atoms; i++){ */
  /*   printf("%d %s %d\n", i, data[i].elem, data[i].type_new); */
  /* } */
  /* exit(0); */
}



void MakeDataSet(double xmin, double xmax, double dx, int m, DATASET *d,
		 double shift){
  
  int i, j;

  d->dx = dx;
  d->m = m; /* 列数 */
  d->n = (int)floor((xmax-xmin)/dx) + 1;
  d->y = malloc(d->n * sizeof(double*));
  for (i=0; i<d->n; i++) d->y[i] = malloc(d->m * sizeof(double));
  /* 初期化 */
  d->x = malloc(d->n * sizeof(double));
  d->x[0] = 0;
  for (i=0; i<d->n; i++) {
    d->x[i] = xmin + i * dx + shift;
  }
  for (i=0; i<d->n; i++){
    for (j=0; j<d->m; j++) d->y[i][j] = 0;
  }
  d->xmin = xmin;
  d->xmax = xmin + (d->n - 1) * dx + 2*shift;
}



void CalcGr(DATA *data, HEAD *head, TYPE *type, DATASET *cn, DATASET *gr,
	    double **cn_, DATASET *gr_){
  int i, j;
  double rho;
  double ni, nj, bunbo;
  
  for (i=0; i<cn->n; i++){
    for (j=0; j<type->pairs; j++){
      if (type->typeij[j][0] == type->typeij[j][1]){
	ni = 2;
	nj = 1;
      }
      else {
	ni = 1;
	nj = 0;
      }
      cn->y[i][j] = cn->y[i][j] + ni * cn_[i][j]/type->nij[j][0];
      rho = (type->nij[j][1] - nj) / head->volume;
      bunbo = 4.0 * M_PI * cn->x[i] * cn->x[i] * cn->dx * rho * type->nij[j][0];
      gr_->y[i][j] = ni * cn_[i][j]/bunbo;
      gr->y[i][j] = gr->y[i][j] + gr_->y[i][j];
    }
    /* total */
    rho = (head->atoms - 1) / head->volume;
    bunbo = 4.0 * M_PI * cn->x[i] * cn->x[i] * cn->dx * rho * head->atoms;
    cn->y[i][cn->m-1] = cn->y[i][cn->m-1] + 2 * cn_[i][cn->m-1]/head->atoms;
    gr->y[i][cn->m-1] = gr->y[i][cn->m-1] + 2 * cn_[i][j]/bunbo;
  }
}

double CalcDistance(HEAD *head,
		    double v1, double v2, double v3,
		    double w1, double w2, double w3){
  double r;
  double x, y, z;
  double dx, dy, dz;

  x = v1 - w1;
  y = v2 - w2;
  z = v3 - w3;

  dx = x*head->M[0][0] + y*head->M[1][0] + z*head->M[2][0];
  dy = x*head->M[0][1] + y*head->M[1][1] + z*head->M[2][1];
  dz = x*head->M[0][2] + y*head->M[1][2] + z*head->M[2][2];
  r = sqrt(dx*dx + dy*dy + dz*dz);
  return r;
}

/* ####################################################################### */
/* CN(r)の計算 */
/* ####################################################################### */
void CalcCn(DATA *data, HEAD *head, TYPE *type, DATASET *cn,
	    DATASET *gr, DATASET *gr_){
  int i, j, k;
  double **cn_;
  int ix, iy, iz;
  int a, b, c;
  double d;
  double v1, v2, v3, w1, w2, w3;
  int p;

  ix = (int)(floor(cn->xmax / head->lx) + 1);
  iy = (int)(floor(cn->xmax / head->ly) + 1);
  iz = (int)(floor(cn->xmax / head->lz) + 1);

  cn_ = malloc(sizeof(double*) * cn->n);
  for (i=0; i<cn->n; i++) cn_[i] = malloc(sizeof(double) * cn->m);
  for (i=0; i<cn->n; i++){
    for (j=0; j<cn->m; j++) cn_[i][j] = 0.0;
  }
  #pragma omp parallel for private(j, a, b, c, v1, v2, v3, w1, w2, w3, d, p, k)
  for (i=0; i<head->atoms; i++){
    for (j=i+1; j<head->atoms; j++){
      for (a=-ix; a<=ix; a++){
	for (b=-iy; b<=iy; b++){
	  for (c=-iz; c<=iz; c++){
	    v1 = data[j].x + a;
	    v2 = data[j].y + b;
	    v3 = data[j].z + c;
	    w1 = data[i].x;
	    w2 = data[i].y;
	    w3 = data[i].z;
	    d = CalcDistance(head, v1, v2, v3, w1, w2, w3);
	    if (cn->xmin < d && d < cn->xmax){
	      p = (int)floor((d - cn->xmin)/cn->dx);
	      k = type->pairid[data[i].type_new][data[j].type_new];
	      cn_[p][k] = cn_[p][k] + 1.0;
	      cn_[p][cn->m-1] = cn_[p][cn->m-1] + 1.0;
	    }
	  }
	}
      }
    }
  }
  CalcGr(data, head, type, cn, gr, cn_, gr_);
}

/* ####################################################################### */
/* CN(r)の計算 */
/* ####################################################################### */
void CalcCn_pp(DATA *data, HEAD *head, TYPE *type, DATASET *cn,
	       DATASET *gr, DATASET *gr_){
  int i, j, k, p=0;
  double dx, dy, dz;
  double v1, v2, v3;
  double d;
  double **cn_;

  cn_ = malloc(sizeof(double*) * cn->n);
  for (i=0; i<cn->n; i++) cn_[i] = malloc(sizeof(double) * cn->m);
  for (i=0; i<cn->n; i++){
    for (j=0; j<cn->m; j++) cn_[i][j] = 0.0;
  }
  #pragma omp parallel for private(j, dx, dy, dz, v1, v2, v3, d, p, k)
  for (i=0; i<head->atoms; i++){
    for (j=i+1; j<head->atoms; j++){
      dx = data[j].x - data[i].x;
      dy = data[j].y - data[i].y;
      dz = data[j].z - data[i].z;
      v1 = dx - rint(dx);
      v2 = dy - rint(dy);
      v3 = dz - rint(dz);
      d = CalcDistance(head, v1, v2, v3, 0.0, 0.0, 0.0);
      if (cn->xmin < d && d < cn->xmax){
	p = (int)floor((d - cn->xmin)/cn->dx);
	k = type->pairid[data[i].type_new][data[j].type_new];
	cn_[p][k] = cn_[p][k] + 1.0;
	cn_[p][cn->m-1] = cn_[p][cn->m-1] + 1.0;
      }
    }
  }
  CalcGr(data, head, type, cn, gr, cn_, gr_);
}

/* ####################################################################### */
/* partial Sq(q)の計算 */
/* ####################################################################### */
void CalcSq(HEAD *head, TYPE *type, DATASET *gr, DATASET *sq){
  int i, j, k;
  double keisuu, sekibun;
  double **sq_;

  sq_ = malloc(sizeof(double*) * sq->n);
  for (i=0; i<sq->n; i++) sq_[i] = malloc(sizeof(double) * sq->m);
  
  for (i=0; i<sq->n; i++)
    for (j=0; j<sq->m; j++ )sq_[i][j] = 0.0;
  
  /* X線・中性子の散乱理論入門 pp. 143-144 */
  for (i=0; i<sq->n; i++){
    keisuu = 4.0 * M_PI * head->rho / sq->x[i];
    for (j=0; j<type->pairs; j++){
      sekibun = 0.0;
      for (k=0; k<gr->n; k++){  /* rで積分 */
        sekibun = sekibun + gr->x[k] * (gr->y[k][j] - 1.0)*
	  sin(sq->x[i] * gr->x[k]) * gr->dx;
      }
      sq_[i][j] = keisuu * sekibun;
    }
  }
  for (i=0; i<sq->n; i++){
    for (j=0; j<sq->m; j++){
      sq->y[i][j] = sq->y[i][j] + sq_[i][j];
    }
  }
}

void TimeAverage(DATASET *d, int step){
  int i, j;

  for (i=0; i<d->n; i++){
    for (j=0; j<d->m; j++){
      d->y[i][j] = d->y[i][j]/step;
    }
  }
}

void CalcTotalSq(TYPE *type, DATASET *sq){
  int i, j, k;
  int p;
  double n;
  double c_xrd, c_nd;
  double *coeff_xrd, coeff_nd;
  
  /* Total sqの計算 */
  /* See: X線・中性子線の散乱理論入門(森北出版) p.145 */
  coeff_xrd = malloc(sizeof(double) * sq->n);
  coeff_nd = 0.0;
  for (i=0; i<sq->n; i++) coeff_xrd[i] = 0.0;
  for (i=0; i<type->types; i++){
    coeff_nd = coeff_nd + type->ci[i] * type->NDcoeff[i];
    for (j=0; j<sq->n; j++){
      coeff_xrd[j] = coeff_xrd[j] + type->ci[i] * type->XRDcoeff[j][i];
    }
  }

  for (i=0; i<sq->n; i++){
    p = 0;
    for (j=0; j<type->types; j++){
      for (k=j; k<type->types; k++){
	if (j == k) n = 1.0;
	else n = 2.0;
	c_xrd = type->ci[j] * type->ci[k] * type->XRDcoeff[i][j] * type->XRDcoeff[i][k];
	c_xrd = n * c_xrd / coeff_xrd[i] / coeff_xrd[i];
	sq->y[i][sq->m-2] = sq->y[i][sq->m-2] + c_xrd * sq->y[i][p];
	
	c_nd = type->ci[j] * type->ci[k] * type->NDcoeff[j] * type->NDcoeff[k];
	c_nd = n * c_nd / coeff_nd / coeff_nd;
	sq->y[i][sq->m-1] = sq->y[i][sq->m-1] + c_nd * sq->y[i][p];
	p++;
      }
    }
  }
}

/* ####################################################################### */
/* メイン関数 */
/* ####################################################################### */
int main(int argn, char **argv){
  FILE *f;
  HEAD head;
  DATA *data = NULL;
  TYPE type;
  DATASET cn, gr, gr_, sq;
  COMMAND com;
  int step, npair;
  int (*readdata)(), (*readheader)();

/* 引数解析 */
  SetArgment(argn, argv, &com);
  if (strcmp(com.filetype, "lammpstrj") == 0){
    readheader = ReadHeadLammpstrj;
    readdata = ReadDataLammpstrj;
  }
  if (strcmp(com.filetype, "xyz") == 0){
    readheader = ReadHeadXYZ;
    readdata = ReadDataXYZ;
  }

  /* 初期化操作 */
  f = fopen(com.infile, "r");
  readheader(f, &head);
  rewind(f);

  if (com.GRmax == -1)  com.GRmax = head.lmax*0.5;
  printf("Maximum cell length = %lf\n", head.lmax);
  printf("Rmax = %lf\n", com.GRmax);
  data = malloc(head.atoms * sizeof(DATA));

  readheader(f, &head);
  readdata(f, &head, data);
  rewind(f);
  
  InitializeType(&head, data, &type);
  npair = (int)(type.types * (type.types + 1) * 0.5);
  MakeDataSet(com.GRmin, com.GRmax, com.GRdel, npair+1, &gr, com.GRdel*0.5);
  MakeDataSet(com.GRmin, com.GRmax, com.GRdel, npair+1, &gr_, com.GRdel*0.5);
  MakeDataSet(com.GRmin, com.GRmax, com.GRdel, npair+1, &cn, com.GRdel*0.5);
  /* X-rayとNeutronのcoeffientセット */
  MakeDataSet(com.SQmin, com.SQmax, com.SQdel, npair+2, &sq, 0.0);
  SetXRDCoeff(&type, &sq);
  SetNDCoeff(&type, &sq);
  /* ここから計算 */
  printf("Start Gr and Sq calculation...\n");
  for (step=0; readheader(f, &head) == 0; step++){
    readdata(f, &head, data);
    if (com.GRmax < head.lmax*0.5)
      CalcCn_pp(data, &head, &type, &cn, &gr, &gr_);
    else
      CalcCn(data, &head, &type, &cn, &gr, &gr_);
    CalcSq(&head, &type, &gr_, &sq);
    if (step % 10 == 0) printf("Calculation %d steps...\n", step);
  }
  printf("Total: %d steps\n", step);
  com.total_step = step;
  /* 時間平均 */
  TimeAverage(&gr, step);
  TimeAverage(&cn, step);
  TimeAverage(&sq, step);
  CalcTotalSq(&type, &sq);
  OutputCoeff(&com, &head, &type, &sq);
  OutputCN(&com, &head, &type, &cn);
  OutputGr(&com, &head, &type, &gr);
  OutputSQ(&com, &head, &type, &sq);
  fclose(f);
}