// pmx  Copyright Notice
// ============================
//
// The pmx source code is copyrighted, but you can freely use and
// copy it as long as you don't change or remove any of the copyright
// notices.
//
// ----------------------------------------------------------------------
// pmx is Copyright (C) 2006-2011 by Daniel Seeliger
//
//                        All Rights Reserved
//
// Permission to use, copy, modify, distribute, and distribute modified
// versions of this software and its documentation for any purpose and
// without fee is hereby granted, provided that the above copyright
// notice appear in all copies and that both the copyright notice and
// this permission notice appear in supporting documentation, and that
// the name of Daniel Seeliger not be used in advertising or publicity
// pertaining to distribution of the software without specific, written
// prior permission.
//
// DANIEL SEELIGER DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS
// SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
// FITNESS.  IN NO EVENT SHALL DANIEL SEELIGER BE LIABLE FOR ANY
// SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER
// RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF
// CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
// ----------------------------------------------------------------------

#include<Geometry.h>

void rotate_rvec(int n,rvec x[],matrix trans)
{
  real   xt,yt,zt;
  int    i;

  for(i=0; (i<n); i++) {
    xt=x[i][XX];
    yt=x[i][YY];
    zt=x[i][ZZ];
    x[i][XX]=trans[XX][XX]*xt+trans[XX][YY]*yt+trans[XX][ZZ]*zt;
    x[i][YY]=trans[YY][XX]*xt+trans[YY][YY]*yt+trans[YY][ZZ]*zt;
    x[i][ZZ]=trans[ZZ][XX]*xt+trans[ZZ][YY]*yt+trans[ZZ][ZZ]*zt;
  }
}



void calc_fit_R(int natoms,real *w_rls,rvec *xp,rvec *x,matrix R)
{
  int    c,r,n,j,i,irot,s;
  //double **omega,**om;
  double d[2*DIM],xnr,xpc;
  matrix vh,vk,u;
  real   mn;
  int    index;
  real   max_d;
  double *omega[2*DIM];
  double *om[2*DIM];
  for(i=0; i<2*DIM; i++) {
    omega[i] = malloc(sizeof(double)*2*DIM);
    om[i] =  malloc(sizeof(double)*2*DIM);
  }

  for(i=0; i<2*DIM; i++) {
    d[i]=0;
    for(j=0; j<2*DIM; j++) {
      omega[i][j]=0;
      om[i][j]=0;
    }
  }

  /*calculate the matrix U*/
  clear_mat(u);
  for(n=0;(n<natoms);n++)
    if ((mn = w_rls[n]) != 0.0)
      for(c=0; (c<DIM); c++) {
	xpc=xp[n][c];
	for(r=0; (r<DIM); r++) {
	  xnr=x[n][r];
	  u[c][r]+=mn*xnr*xpc;
	}
      }

  /*construct omega*/
  /*omega is symmetric -> omega==omega' */
  for(r=0; r<2*DIM; r++)
    for(c=0; c<=r; c++)
      if (r>=DIM && c<DIM) {
        omega[r][c]=u[r-DIM][c];
        omega[c][r]=u[r-DIM][c];
      } else {
        omega[r][c]=0;
        omega[c][r]=0;
      }

  /*determine h and k*/

  jacobi6(omega,d,om,&irot);

  /*real   **omega = input matrix a[0..n-1][0..n-1] must be symmetric
   *int     natoms = number of rows and columns
   *real      NULL = d[0]..d[n-1] are the eigenvalues of a[][]
   *real       **v = v[0..n-1][0..n-1] contains the vectors in columns
   *int      *irot = number of jacobi rotations
   */


  index=0; /* For the compiler only */

  /* Copy only the first DIM-1 eigenvectors */
  for(j=0; j<DIM-1; j++) {
    max_d=-1000;
    for(i=0; i<2*DIM; i++)
      if (d[i]>max_d) {
        max_d=d[i];
        index=i;
      }
    d[index]=-10000;
    for(i=0; i<DIM; i++) {
      vh[j][i]=M_SQRT2*om[i][index];
      vk[j][i]=M_SQRT2*om[i+DIM][index];
    }
  }
  /* Calculate the last eigenvector as the outer-product of the first two.
   * This insures that the conformation is not mirrored and
   * prevents problems with completely flat reference structures.
   */
  cprod(vh[0],vh[1],vh[2]);
  cprod(vk[0],vk[1],vk[2]);


  /* determine R */
  clear_mat(R);
  for(r=0; r<DIM; r++)
    for(c=0; c<DIM; c++)
      for(s=0; s<DIM; s++)
	R[r][c] += vk[s][r]*vh[s][c];
  for(r=DIM; r<DIM; r++)
    R[r][r] = 1;

  for(i=0; i<2*DIM; i++) {
    free(omega[i]);
    free(om[i]);
  }
#if PY_MAJOR_VERSION >= 3
    /* this solves a problem in py3 in which this func raises the error:
     * SystemError: <built-in function calc_fit_R> returned a result with an error set
     * however, this fix is 100% empirical - it works but there may be better
     * ways to do this
     */
    PyErr_Clear();
#endif
}


void do_fit(int natoms,real *w_rls,rvec *xp,rvec *x)
{
  int    j,m,r,c;
  matrix R;
  rvec   x_old;

  /* Calculate the rotation matrix R */
  calc_fit_R(natoms,w_rls,xp,x,R);

  /*rotate X*/
  for(j=0; j<natoms; j++) {
    for(m=0; m<DIM; m++)
      x_old[m]=x[j][m];
    for(r=0; r<DIM; r++) {
      x[j][r]=0;
      for(c=0; c<DIM; c++)
        x[j][r]+=R[r][c]*x_old[c];
    }
  }
}

void center(rvec *x, int n)
{
  int i;
  rvec c;
  clear_rvec(c);
  for(i=0;i<n;i++){
    rvec_add(c,x[i],c);
  }
  c[XX]/=(real) n;
  c[YY]/=(real) n;
  c[ZZ]/=(real) n;
  for(i=0;i<n;i++){
    rvec_sub(x[i],c,x[i]);
  }
}

void center_and_get_vec(rvec *x, int n, rvec c)
{
  int i;
  clear_rvec(c);
  for(i=0;i<n;i++){
    rvec_add(c,x[i],c);
  }
  c[XX]/=(real) n;
  c[YY]/=(real) n;
  c[ZZ]/=(real) n;
  for(i=0;i<n;i++){
    rvec_sub(x[i],c,x[i]);
  }
}

void center_vec(rvec *x, int n, rvec c)
{
  int i;
  clear_rvec(c);
  for(i=0;i<n;i++){
    rvec_add(c,x[i],c);
  }
  c[XX]/=(real) n;
  c[YY]/=(real) n;
  c[ZZ]/=(real) n;
}


real planarity(rvec *x, int n)
{
  int m,k;
  matrix mat,trans;
  rvec dd;
  rvec xx[n];
  for(m=0;m<n;m++){
    copy_rvec( x[m], xx[m]);
  }
//   for(int i=0; (i<n); i++) {
//     if( x[i][XX] != x[i][XX] ||
// 	x[i][YY] != x[i][YY] ||
// 	x[i][ZZ] != x[i][ZZ] ) {
//       cout << "here" << endl;
//     }
//   }

  clear_mat(mat);
  clear_mat(trans);

  center(xx, n);

  clear_rvec(dd);
  princ_comp(n,xx,mat,dd);
  if (det(mat) < 0) {
    for(m=0; (m<DIM); m++)
      mat[ZZ][m] = -mat[ZZ][m];
  }

  rotate_rvec(n,xx,mat);
  real adev = 0.;
  for(k=0;k<n;k++){
    adev+=sqrt(sqr(xx[k][0]));
  }
  adev/=(real)n;
  return adev;
}

void set_planar(rvec *x, int n)
{
  int i,m,k;
  matrix mat,trans;
  rvec dd;
  rvec cm;
//   for(i=0; (i<n); i++) {
//     if( x[i][XX] != x[i][XX] ||
// 	x[i][YY] != x[i][YY] ||
// 	x[i][ZZ] != x[i][ZZ] ) {
//       cout << "here" << endl;
//     }
//   }

  clear_mat(mat);
  clear_mat(trans);

  center_and_get_vec(x, n, cm);

  clear_rvec(dd);
  princ_comp(n,x,mat,dd);
  if (det(mat) < 0) {
    for(m=0; (m<DIM); m++)
      mat[ZZ][m] = -mat[ZZ][m];
  }

  rotate_rvec(n,x,mat);
  real adev = 0.;
  for(k=0;k<n;k++){
    adev+=sqrt(sqr(x[k][0]));
  }
  adev/=(real)n;

  /* make it flat */
  for(k=0;k<n;k++){
    x[k][0] = 0.;
  }

  transpose(mat,trans);

  rotate_rvec(n,x,trans);

  for(i=0;i<n;i++){
    rvec_add(x[i],cm,x[i]);
  }
}



real radius_of_gyration(rvec *x, real *mass, int n )

{
  int i,m;
  matrix mat;
  rvec dd;
  rvec comp;
  real dx2;
  real gyro;
  real tm = 0.;
  rvec xx[n];
  copy_x(x,xx,n);

  clear_rvec(dd);
  clear_rvec(comp);
  clear_mat(mat);
  center(xx, n);
  princ_comp(n,xx,mat,dd);

  if (det(mat) < 0) {
    for(m=0; (m<DIM); m++)
      mat[ZZ][m] = -mat[ZZ][m];
  }
  rotate_rvec(n,xx,mat);

  for(i=0;i<n;i++){
    tm+=mass[i];
    for(m=0;m<DIM;m++){
      dx2 = sqr(xx[i][m]);
      comp[m]+=dx2*mass[i];
    }
  }
  gyro = comp[XX]+comp[YY]+comp[ZZ];
  return sqrt(gyro/tm);
}

void cryst1_to_box( char *line, matrix box)
{
#define SG_SIZE 11
  char sa[12],sb[12],sc[12],sg[SG_SIZE+1],ident;
  double fa,fb,fc,alpha,beta,gamma,cosa,cosb,cosg,sing;
  int  syma,symb,symc;
  int  ePBC_file;
  enum {
    epbcXYZ, epbcNONE, epbcXY, epbcSCREW, epbcNR
  };


  sscanf(line,"%*s%s%s%s%lf%lf%lf",sa,sb,sc,&alpha,&beta,&gamma);

  ePBC_file = -1;
  if (strlen(line) >= 55) {
    strncpy(sg,line+55,SG_SIZE);
    sg[SG_SIZE] = '\0';
    ident = ' ';
    syma  = 0;
    symb  = 0;
    symc  = 0;
    sscanf(sg,"%c %d %d %d",&ident,&syma,&symb,&symc);
    if (ident == 'P' && syma ==  1 && symb <= 1 && symc <= 1) {
      fc = atof(sc)*0.1;
      ePBC_file = (fc > 0 ? epbcXYZ : epbcXY);
    }
    if (ident == 'P' && syma == 21 && symb == 1 && symc == 1) {
      ePBC_file = epbcSCREW;
    }
  }

  fa = atof(sa)*0.1;
  fb = atof(sb)*0.1;
  fc = atof(sc)*0.1;
  if (ePBC_file == epbcSCREW) {
    fa *= 0.5;
  }
  clear_mat(box);
  box[XX][XX] = fa;
  if ((alpha!=90.0) || (beta!=90.0) || (gamma!=90.0)) {
    if (alpha != 90.0) {
      cosa = cos(alpha*DEG2RAD);
    } else {
      cosa = 0;
    }
    if (beta != 90.0) {
      cosb = cos(beta*DEG2RAD);
    } else {
      cosb = 0;
    }
    if (gamma != 90.0) {
      cosg = cos(gamma*DEG2RAD);
      sing = sin(gamma*DEG2RAD);
    } else {
      cosg = 0;
      sing = 1;
    }
    box[YY][XX] = fb*cosg;
    box[YY][YY] = fb*sing;
    box[ZZ][XX] = fc*cosb;
    box[ZZ][YY] = fc*(cosa - cosb*cosg)/sing;
    box[ZZ][ZZ] = sqrt(fc*fc
		       - box[ZZ][XX]*box[ZZ][XX] - box[ZZ][YY]*box[ZZ][YY]);
  } else {
    box[YY][YY] = fb;
    box[ZZ][ZZ] = fc;
  }
}

void box_to_cryst1(matrix box, char *line)
{
  real alpha,beta,gamma;
  char cryst1_line[256];
  if (norm2(box[YY])*norm2(box[ZZ])!=0)
    alpha = RAD2DEG*acos(cos_angle(box[YY],box[ZZ]));
  else
    alpha = 90;
  if (norm2(box[XX])*norm2(box[ZZ])!=0)
    beta  = RAD2DEG*acos(cos_angle(box[XX],box[ZZ]));
  else
    beta  = 90;
  if (norm2(box[XX])*norm2(box[YY])!=0)
    gamma = RAD2DEG*acos(cos_angle(box[XX],box[YY]));
  else
    gamma = 90;
  sprintf(cryst1_line,"CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f %-11s%4d",
	  10*norm(box[XX]),10*norm(box[YY]),10*norm(box[ZZ]),
	  alpha,beta,gamma,"P 1",1);
  strcpy(line, cryst1_line);
}



void max_crd(rvec *x, int n, matrix max)
{
  /* matrix X stores minimum
     matrix Z stores maximu
     this max[YY][XX] is minimum y value
  */

  int i;
  for(i=0;i<DIM;i++){
    max[i][XX] = 99999;
    max[i][YY] = 0.;
    max[i][ZZ] = -99999.;
  }
  for(i=0;i<n;i++){
    if(x[i][XX] < max[XX][XX]) max[XX][XX] = x[i][XX];
    if(x[i][XX] > max[XX][ZZ]) max[XX][ZZ] = x[i][XX];

    if(x[i][YY] < max[YY][XX]) max[YY][XX] = x[i][YY];
    if(x[i][YY] > max[YY][ZZ]) max[YY][ZZ] = x[i][YY];

    if(x[i][ZZ] < max[ZZ][XX]) max[ZZ][XX] = x[i][ZZ];
    if(x[i][ZZ] > max[ZZ][ZZ]) max[ZZ][ZZ] = x[i][ZZ];
  }
}


int gridp(real x, real origin, real inv_spacing, int max)
{
  real n;
  int point;
  n = (x-origin)*inv_spacing;
  point = (int) n;
  return point;
}

typedef struct GridMap GridMap;
struct GridMap
{
  rvec center;
  rvec origin;
  int nelem;
  ivec n;
  real spacing;
  real inv_spacing;
  int **cell;
  int *natom;

};


GridMap *grid_init(void)
{
  int i;
  GridMap *map = NULL;
  map = malloc(sizeof(GridMap));

  for(i=0;i<DIM;i++){
    map->n[i] = 0;
    map->origin[i] = 0.;
    map->center[i] = 0.;
  }
  map->nelem = 0;
  map->spacing = 0.;
  map->inv_spacing = 0.;
  map->natom = NULL;
  map->cell = NULL;
  return map;
}

void free_grid(GridMap *gp)
{
  int i;
  for(i=0;i<gp->nelem;i++){
    if( gp->cell[i] != NULL )
      free(gp->cell[i]);
  }
  free(gp->cell);
  free(gp->natom);
  free(gp);
}

GridMap *spread_atoms_on_grid(rvec *x, ivec *cells, int natoms, GridMap *gp, real cutoff)
{
  int i;
  matrix max;
  int xdim, ydim, zdim;
  int gpx, gpy, gpz;
  int nx, ny, nz;
  int idx;

  max_crd(x,natoms,max);
  gp->spacing = cutoff;
  gp->inv_spacing = 1./cutoff;

  for(i=0;i<DIM;i++){
    gp->origin[i] = max[i][XX]-gp->spacing/2.;
  }

  max[XX][XX]-=gp->spacing/2.;
  max[YY][XX]-=gp->spacing/2.;
  max[ZZ][XX]-=gp->spacing/2.;
  max[XX][ZZ]+=gp->spacing/2.;
  max[YY][ZZ]+=gp->spacing/2.;
  max[ZZ][ZZ]+=gp->spacing/2.;

  xdim = (int) (max[XX][ZZ] - max[XX][XX])*gp->inv_spacing +1;
  ydim = (int) (max[YY][ZZ] - max[YY][XX])*gp->inv_spacing +1;
  zdim = (int) (max[ZZ][ZZ] - max[ZZ][XX])*gp->inv_spacing +1;

  gp->n[XX] = xdim;
  gp->n[YY] = ydim;
  gp->n[ZZ] = zdim;
  nx = gp->n[XX];
  ny = gp->n[YY];
  nz = gp->n[ZZ];
  gp->nelem = gp->n[XX]*gp->n[YY]*gp->n[ZZ];
  gp->cell = malloc(sizeof(int*)*gp->nelem);
  gp->natom = malloc(sizeof(int)*gp->nelem);

  for(i=0;i<gp->nelem;i++){
    gp->cell[i] = NULL;
    gp->natom[i] = 0;
  }
  for(i=0;i<natoms;i++){
    gpx = gridp(x[i][XX],gp->origin[XX],gp->inv_spacing, xdim);
    gpy = gridp(x[i][YY],gp->origin[YY],gp->inv_spacing, ydim);
    gpz = gridp(x[i][ZZ],gp->origin[ZZ],gp->inv_spacing, zdim);
    cells[i][XX] = gpx;
    cells[i][YY] = gpy;
    cells[i][ZZ] = gpz;
    idx = gpz*nx*ny + gpy*nx + gpx;
    gp->natom[idx]+=1;
    gp->cell[idx] = realloc(gp->cell[idx], sizeof(int)*gp->natom[idx] );
    gp->cell[idx][gp->natom[idx]-1] = i;
  }
  return gp;
}

void search_neighbors(rvec *x, int natoms, real cutoff, int **nlist, int *nat)
{

  int i,k;

  GridMap *gp = grid_init();


  int nx, ny, nz;
  int idx;
  int gpx, gpy, gpz;
  int xx, yy, zz;
  int atom_id;
  rvec diff;
  real d;
  real cut2;

  ivec cells[natoms];

  xx = -1;
  yy = -1;
  zz = -1;

  cut2 = sqr(cutoff);
  gp = spread_atoms_on_grid(x,cells,natoms,gp,cutoff);
  nx = gp->n[XX];
  ny = gp->n[YY];
  nz = gp->n[ZZ];



  for(i=0;i<natoms;i++){
    xx = -1;
    while(xx<2){
      yy = -1;
      while(yy<2){
	zz = -1;
	while(zz<2){

	  gpx = xx+cells[i][XX];
	  gpy = yy+cells[i][YY];
	  gpz = zz+cells[i][ZZ];
	  if( (gpx < gp->n[XX]) && (gpy < gp->n[YY]) && (gpz < gp->n[ZZ])
	      && (gpx >=0) && (gpy >= 0) && (gpz >= 0))
	    {
	      idx = gpz*nx*ny + gpy*nx + gpx;
	      for(k=0;k<gp->natom[idx];k++){
		atom_id = gp->cell[idx][k];
		if(atom_id > i){
		  rvec_sub(x[i],x[atom_id],diff);
		  d = norm2(diff);
		  if(d < cut2){
		    nat[i]+=1;
		    nlist[i] = realloc(nlist[i], sizeof(int)*(nat[i]));
		    nlist[i][nat[i]-1] = atom_id;
		  }
		}
	      }
	    }
	  zz++;
	}
	yy++;
      }
      xx++;
    }
  }
  free_grid(gp);
}


real distance_from_atoms( PyObject *atom1, PyObject *atom2)
{
  PyObject *cs1 = PyObject_GetAttrString(atom1,"x");
  PyObject *cs2 = PyObject_GetAttrString(atom2,"x");
  rvec x1;
  rvec x2;
  Pyvec2rvec(cs1,x1);
  Pyvec2rvec(cs2,x2);
  return dist( x1, x2 );
}

real distance2_from_atoms( PyObject *atom1, PyObject *atom2)
{
  PyObject *cs1 = PyObject_GetAttrString(atom1,"x");
  PyObject *cs2 = PyObject_GetAttrString(atom2,"x");
  rvec x1;
  rvec x2;
  Pyvec2rvec(cs1,x1);
  Pyvec2rvec(cs2,x2);
  return distance2( x1, x2 );
}

real angle_from_atoms( PyObject *atom1, PyObject *atom2, PyObject *atom3)
{
  PyObject *cs1 = PyObject_GetAttrString(atom1,"x");
  PyObject *cs2 = PyObject_GetAttrString(atom2,"x");
  PyObject *cs3 = PyObject_GetAttrString(atom3,"x");
  rvec x1;
  rvec x2;
  rvec x3;
  Pyvec2rvec(cs1,x1);
  Pyvec2rvec(cs2,x2);
  Pyvec2rvec(cs3,x3);
  return angle_ijk(x1,x2,x3);
}

real dihedral_from_atoms( PyObject *atom1, PyObject *atom2, PyObject *atom3, PyObject *atom4)
{
  PyObject *cs1 = PyObject_GetAttrString(atom1,"x");
  PyObject *cs2 = PyObject_GetAttrString(atom2,"x");
  PyObject *cs3 = PyObject_GetAttrString(atom3,"x");
  PyObject *cs4 = PyObject_GetAttrString(atom4,"x");
  rvec x1;
  rvec x2;
  rvec x3;
  rvec x4;
  Pyvec2rvec(cs1,x1);
  Pyvec2rvec(cs2,x2);
  Pyvec2rvec(cs3,x3);
  Pyvec2rvec(cs4,x4);
  return dihedral(x1,x2,x3,x4);
}
