// 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 PyObject2rvec( PyObject *o, rvec *x, int natoms)
{
  int i,k;
  for(i=0;i<natoms;i++){
    PyObject *vec = PySequence_GetItem(o, i);
    for(k=0;k<DIM;k++){
      x[i][k] = PyFloat_AsDouble( PySequence_GetItem(vec, k) );
    }
  }
}


PyObject* rvec2PyObject( rvec *x, int natoms)
{
  int i;
  PyObject *o = PyList_New( natoms );
  for(i=0;i<natoms;i++){
    PyList_SetItem(o, i, Py_BuildValue("[ddd]", x[i][XX], x[i][YY], x[i][ZZ]));
  }
  return o;
}

PyObject* matrix2PyObject( matrix R)
{
  int i;
  PyObject *o = PyList_New( DIM );
  for(i=0;i<DIM;i++){
    PyList_SetItem(o, i, Py_BuildValue("[ddd]", R[i][XX], R[i][YY], R[i][ZZ]));
  }
  return o;
}

void PyObject2matrix( PyObject *o, matrix R)
{
  int i,k;
  for(i=0;i<DIM;i++){
    PyObject *vec = PySequence_GetItem(o, i);
    for(k=0;k<DIM;k++){
      R[i][k] = PyFloat_AsDouble( PySequence_GetItem(vec, k) );
    }
  }
}


void PyObject2real_array( PyObject *o, real *arr, int natoms)
{
  int i;
  for(i=0;i<natoms;i++){
    arr[i] = PyFloat_AsDouble( PySequence_GetItem(o, i) );
  }
}

PyObject *wrap_dist(PyObject *self,PyObject *args)
{
  PyObject *cs1, *cs2;
  PyObject *fast1, *fast2;
  if(!PyArg_ParseTuple(args,"OO",&cs1,&cs2))
    return NULL;

  if(!(fast1 = PySequence_Fast(cs1, "could not get fast sequence")))
    return NULL;
  if(!(fast2 = PySequence_Fast(cs2, "could not get fast sequence"))) {
    return NULL;
    }

  rvec r1,r2;
  int i;
  PyObject *d1,*d2;
  for(i=0;i<DIM;i++){
    d1=PySequence_Fast_GET_ITEM(fast1,i);
    d2=PySequence_Fast_GET_ITEM(fast2,i);
    r1[i] = PyFloat_AsDouble(d1);
    r2[i] = PyFloat_AsDouble(d2);
  }
  real d = dist(r1,r2);
  PyObject *ret;

  ret=Py_BuildValue("d",d);
  return ret;
}

PyObject *wrap_distance2(PyObject *self,PyObject *args)
{
  PyObject *cs1, *cs2;
  PyObject *fast1, *fast2;

  if(!PyArg_ParseTuple(args,"OO",&cs1,&cs2))
    return NULL;

  if(!(fast1 = PySequence_Fast(cs1, "could not get fast sequence")))
    return NULL;
  if(!(fast2 = PySequence_Fast(cs2, "could not get fast sequence"))) {
    return NULL;
  }
  rvec r1,r2;
  int i;
  PyObject *d1,*d2;
  for(i=0;i<DIM;i++){
    d1=PySequence_Fast_GET_ITEM(fast1,i);
    d2=PySequence_Fast_GET_ITEM(fast2,i);
    r1[i] = PyFloat_AsDouble(d1);
    r2[i] = PyFloat_AsDouble(d2);
  }
  real d = distance2(r1,r2);
  PyObject *ret;
  ret=Py_BuildValue("d",d);
  return ret;
}

PyObject *wrap_angle_ijk(PyObject *self,PyObject *args)
{
  PyObject *cs1, *cs2,*cs3;
  if(!PyArg_ParseTuple(args,"OOO",&cs1,&cs2,&cs3))
    return NULL;
  rvec r1,r2,r3;
  int i;
  for(i=0;i<DIM;i++){
    r1[i]=PyFloat_AsDouble(PySequence_GetItem(cs1,i));
    r2[i]=PyFloat_AsDouble(PySequence_GetItem(cs2,i));
    r3[i]=PyFloat_AsDouble(PySequence_GetItem(cs3,i));
  }
  real angle = angle_ijk(r1,r2,r3);
  PyObject *ret;
  ret=Py_BuildValue("d",angle);
  return ret;
}

PyObject *wrap_dihedral(PyObject *self,PyObject *args)
{

  PyObject *cs1, *cs2,*cs3, *cs4;
  if(!PyArg_ParseTuple(args,"OOOO",&cs1,&cs2,&cs3,&cs4))
    return NULL;
  rvec r1,r2,r3,r4;
  int i;
  for(i=0;i<DIM;i++){
    r1[i]=PyFloat_AsDouble(PySequence_GetItem(cs1,i));
    r2[i]=PyFloat_AsDouble(PySequence_GetItem(cs2,i));
    r3[i]=PyFloat_AsDouble(PySequence_GetItem(cs3,i));
    r4[i]=PyFloat_AsDouble(PySequence_GetItem(cs4,i));
  }
  real d = dihedral(r1,r2,r3,r4);
  PyObject *ret;
  ret=Py_BuildValue("d",d);
  return ret;
}

PyObject *wrap_planarity(PyObject *self,PyObject *args)
{

  PyObject *cs;
  if(!PyArg_ParseTuple(args,"O",&cs))
    return NULL;
  int natoms = PySequence_Length(cs);
  rvec x[natoms];

  PyObject2rvec( cs, x, natoms);

  real p = planarity(x, natoms);
  PyObject *ret;
  ret = Py_BuildValue("d",p);
  return ret;
}

PyObject *wrap_fit(PyObject *self,PyObject *args)
{

  PyObject *cs1, *cs2, *mass;
  if(!PyArg_ParseTuple(args,"OOO",&cs1, &cs2, &mass))
    return NULL;
  int natoms1 = PySequence_Length(cs1);
  int natoms2 = PySequence_Length(cs2);
  if( natoms1 != natoms2 ) {
    Error("Cannot fit coordinate sets with different lengths");
  }
  rvec x1[natoms1];
  rvec x2[natoms1];
  real m[natoms1];
  PyObject2rvec( cs1, x1, natoms1);
  PyObject2rvec( cs2, x2, natoms2);
  PyObject2real_array(mass, m, natoms1);
  rvec cent;
  center_and_get_vec(x1, natoms1, cent);     // center x1 and get vector for back translation
  center(x2, natoms1);                                // center x2
  do_fit(natoms1, m, x1, x2);

  int i;
  for(i=0;i<natoms1;i++)                    // translate back
    rvec_add( x2[i], cent, x2[i]);

  PyObject *ret = rvec2PyObject(x2, natoms1);
  return ret;
}

PyObject *wrap_calc_fit_R(PyObject *self,PyObject *args)
{

  PyObject *cs1, *cs2, *mass;
  if(!PyArg_ParseTuple(args,"OOO",&cs1, &cs2, &mass))
    return NULL;
  int natoms1 = PySequence_Length(cs1);
  int natoms2 = PySequence_Length(cs2);
  if( natoms1 != natoms2 ) {
    Error("Cannot fit coordinate sets with different lengths");
  }
  rvec x1[natoms1];
  rvec x2[natoms1];
  real m[natoms1];
  PyObject2rvec( cs1, x1, natoms1);
  PyObject2rvec( cs2, x2, natoms2);
  PyObject2real_array(mass, m, natoms1);
  center(x1, natoms1);
  center(x2, natoms1);
  matrix R;
  clear_mat(R);
  calc_fit_R(natoms1,m,x1,x2,R);
  PyObject *ret = matrix2PyObject(R);
  return ret;
}

PyObject *wrap_center_vec( PyObject *self, PyObject *args)
{
  PyObject *cs;
  if(!PyArg_ParseTuple(args,"O",&cs))
    return NULL;
  int natoms = PySequence_Length(cs);
  rvec x[natoms];

  PyObject2rvec( cs, x, natoms);
  rvec center;
  center_vec( x, natoms, center);
  return Py_BuildValue("[ddd]",center[XX], center[YY], center[ZZ]);
}

PyObject *wrap_cryst1_to_box( PyObject *self, PyObject *args)
{
  char *line;
  if(!PyArg_ParseTuple(args,"s",&line))
    return NULL;

  matrix box;
  cryst1_to_box(line, box);
  PyObject *pBox = matrix2PyObject( box );
  return pBox;
}

PyObject *wrap_box_to_cryst1( PyObject *self, PyObject *args)
{
  PyObject *pBox;
  if(!PyArg_ParseTuple(args,"O",&pBox))
    return NULL;

  matrix box;
  PyObject2matrix(pBox, box);
  char line[256];
  box_to_cryst1(box, line);
  return PyString_FromString( line );
}

void Pyvec2rvec( PyObject *Ox, rvec x)
{
  x[XX] = PyFloat_AsDouble( PySequence_Fast_GET_ITEM(Ox, XX) );
  x[YY] = PyFloat_AsDouble( PySequence_Fast_GET_ITEM(Ox, YY) );
  x[ZZ] = PyFloat_AsDouble( PySequence_Fast_GET_ITEM(Ox, ZZ) );
}

real get_bond_contribution(PyObject *atom)
{
  char *elem = PyString_AsString( PyObject_GetAttrString(atom, "symbol") );
  if(strcmp(elem,"C") == 0 ) return CCONTR;
  else if( strcmp(elem,"H") == 0) return HCONTR;
  else if( strcmp(elem,"N") == 0) return NCONTR;
  else if( strcmp(elem,"O") == 0) return OCONTR;
  else if( strcmp(elem,"S") == 0) return SCONTR;
  else if( strcmp(elem,"P") == 0) return PCONTR;
  else if( strcmp(elem,"F") == 0) return FCONTR;
  else if( strcmp(elem,"CL") == 0) return CLCONTR;
  else if( strcmp(elem,"BR") == 0) return BRCONTR;
  else if( strcmp(elem,"I") == 0) return ICONTR;
  else return MCONTR; // default works fine for most ions
}

void atom_ids_from_py_atomlist( PyObject *list, int *atom_ids)
{
  int nat = PySequence_Length(list);
  int i;
  for(i=0;i<nat;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(list, i);
    int atom_id =  PyInt_AsLong (PyObject_GetAttrString(atom,"id") ) - 1;
    atom_ids[i] = atom_id;
  }
}

bool atom_in_py_atomlist( PyObject *list, int id)
{
  int i;
  bool atom_in_list = FALSE;
  int nat = PySequence_Length(list);
  int atom_ids[nat];
  atom_ids_from_py_atomlist(list, atom_ids);
  for(i=0;i<nat;i++){
    if (id == atom_ids[i] ) {
      atom_in_list = TRUE;
      break;
    }
  }
  return atom_in_list;
}

bool is_bound( PyObject *atom, int id)
{
  // check if atom with id id is in the bonded lists of atom
  PyObject *bonds = PyObject_GetAttrString(atom,"bonds");
  if( atom_in_py_atomlist( bonds, id) ) return TRUE;
  PyObject *b13 = PyObject_GetAttrString(atom,"b13");
  if( atom_in_py_atomlist( b13, id) ) return TRUE;
  PyObject *b14 = PyObject_GetAttrString(atom,"b14");
  if( atom_in_py_atomlist( b14, id) ) return TRUE;
  return FALSE;
}

void clear_PySequence(PyObject *list )
{
  int i;
  int n = PySequence_Length(list);
  for(i=0;i<n;i++){
    PySequence_DelItem(list,0);
  }
}


void reset_bond_lists( PyObject *atomlist )
{
  int i;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);

    PyObject *bonds = PyObject_GetAttrString(atom,"bonds");
    clear_PySequence(bonds);
    PyObject *b13 = PyObject_GetAttrString(atom,"b13");
    clear_PySequence(b13);
    PyObject *b14 = PyObject_GetAttrString(atom,"b14");
    clear_PySequence(b14);

  }
}

real dist2_from_atoms( PyObject *atom1, PyObject *atom2)
{
  PyObject *Ox1 = PyObject_GetAttrString(atom1, "x");
  PyObject *Ox2 = PyObject_GetAttrString(atom2, "x");
  rvec x1, x2;
  Pyvec2rvec( Ox1, x1);
  Pyvec2rvec( Ox2, x2);
  return  distance2(x1, x2 );
}

void build_bonds_by_distance( PyObject *atomlist, int **nlist, int *nat)
{
  int i,k;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_GetItem(atomlist, i);
    PyObject *bonds = PyObject_GetAttrString(atom,"bonds");
    real bcontr1 = get_bond_contribution( atom );
    for(k=0;k<nat[i];k++){
      int atom_id = nlist[i][k];
      if( atom_id > i ) {
	PyObject *atom2 = PySequence_GetItem(atomlist, atom_id);
	real bcontr2 = get_bond_contribution( atom2 );
	real d = dist2_from_atoms( atom, atom2 );
	if( d < sqr( bcontr1+bcontr2 ) ){
	  PyObject *bonds2 = PyObject_GetAttrString(atom2,"bonds");
	  PyList_Append(bonds, atom2);
	  PyList_Append(bonds2, atom);
	}
      }
    }
  }
}

void build_b13_from_bonds( PyObject *atomlist)
{
  int i,k,l;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_GetItem(atomlist, i);
    PyObject *b13 = PyObject_GetAttrString(atom,"b13");
    int atom_id = PyInt_AsLong (PyObject_GetAttrString(atom,"id") );

    PyObject *bonds = PyObject_GetAttrString(atom,"bonds");
    int len_bonds = PySequence_Length(bonds);
    for(k=0;k<len_bonds;k++){
      PyObject *bb = PySequence_GetItem(bonds, k);
      PyObject *bb_bonds = PyObject_GetAttrString(bb,"bonds");
      int len_bbonds = PySequence_Length(bb_bonds);
      for(l=0;l<len_bbonds;l++){
	PyObject *bond = PySequence_GetItem(bb_bonds, l);
	int atom_id2 = PyInt_AsLong (PyObject_GetAttrString(bond,"id") );
	if( atom_id < atom_id2 ) {
	  PyObject *bb13 = PyObject_GetAttrString(bond,"b13");
	  PyList_Append(b13, bond );
	  PyList_Append(bb13, atom );
	}
      }
    }
  }
}

void build_b14_from_bonds(PyObject *atomlist )
{

  int i,k, l;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_GetItem(atomlist, i);
    PyObject *b14 = PyObject_GetAttrString(atom,"b14");
    int atom_id = PyInt_AsLong (PyObject_GetAttrString(atom,"id") );

    PyObject *b13 = PyObject_GetAttrString(atom,"b13");
    int len_b13 = PySequence_Length(b13);
    for(k=0;k<len_b13;k++){
      PyObject *bb = PySequence_GetItem(b13, k);
      PyObject *bb_bonds = PyObject_GetAttrString(bb,"bonds");
      int len_bonds = PySequence_Length(bb_bonds);
      for(l=0;l<len_bonds;l++){
	PyObject *bond = PySequence_GetItem(bb_bonds, l);
	int atom_id2 = PyInt_AsLong (PyObject_GetAttrString(bond,"id") );
	if( atom_id < atom_id2  && ! is_bound(atom, atom_id2 -1 )) {
	  PyObject *bb14 = PyObject_GetAttrString(bond,"b14");
	  PyList_Append(b14, bond );
	  PyList_Append(bb14, atom );
	}
      }
    }
  }
}


con_table *build_table( PyObject *atomlist)
{
  int i,k;
  int natoms = PySequence_Length(atomlist);
  con_table *t = malloc(sizeof(con_table) );
  t->natoms = natoms;
  t->ncon = malloc(sizeof(int)*natoms);
  t->con = malloc(sizeof(int*)*natoms);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    PyObject *bonds = PyObject_GetAttrString(atom,"bonds");
    PyObject *b13 = PyObject_GetAttrString(atom,"b13");
    PyObject *b14 = PyObject_GetAttrString(atom,"b14");
    int nbonds = PySequence_Length(bonds);
    int nb13 = PySequence_Length(b13);
    int nb14 = PySequence_Length(b14);
    int size = nbonds+nb13+nb14;
    t->ncon[i] = size;
    t->con[i] = malloc(sizeof(int)*size);
    size = 0;
    for(k=0;k<nbonds;k++){
      PyObject *a = PySequence_Fast_GET_ITEM( bonds, k );
      int id = PyInt_AsLong( PyObject_GetAttrString(a,"id") ) - 1;
      t->con[i][size] = id;
      size++;
    }
    for(k=0;k<nb13;k++){
      PyObject *a = PySequence_Fast_GET_ITEM( b13, k );
      int id = PyInt_AsLong( PyObject_GetAttrString(a,"id") ) - 1;
      t->con[i][size] = id;
      size++;
    }
    for(k=0;k<nb14;k++){
      PyObject *a = PySequence_Fast_GET_ITEM( b14, k );
      int id = PyInt_AsLong( PyObject_GetAttrString(a,"id") ) - 1;
      t->con[i][size] = id;
      size++;
    }
  }
  return t;
}

bool id_in_con_table( con_table *t, int atom_id, int check)
{
  int i;
  for(i=0;i<t->ncon[atom_id];i++){
    if( t->con[atom_id][i] == check ) return TRUE;
  }
  return FALSE;
}

void delete_con_table( con_table *t )
{
  int i;
  for(i=0;i<t->natoms;i++){
    free( t->con[i] );
  }
  free( t->con );
  free( t->ncon );
  free(t);
}


PyObject* wrap_search_neighbors(PyObject *self, PyObject *args)
{
  PyObject *atomlist;
  real cutoff;
  bool build_bonds;
  if(!PyArg_ParseTuple(args,"Odi",&atomlist,&cutoff,&build_bonds))
    return NULL;

  int natoms = PySequence_Length(atomlist);
  rvec x[natoms];
  int *nlist[natoms];
  int nat[natoms];


  int i, k;
  for(i=0;i<natoms;i++){
    nat[i] = 0;
    nlist[i] = malloc(sizeof(int));
    PyObject *py_atom = PySequence_GetItem(atomlist, i);
    PyObject *Ox = PyObject_GetAttrString(py_atom,"x");
    Pyvec2rvec(Ox, x[i] );
  }

  search_neighbors(x, natoms, cutoff, nlist, nat);

  if( build_bonds ) {
    reset_bond_lists( atomlist );
  }


  // reset neighborlists
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    PyObject *neighbor_list = PyObject_GetAttrString(atom,"neighbors");
    int list_size = PySequence_Length( neighbor_list );
    for(k=0;k<list_size;k++){
      PySequence_DelItem(neighbor_list,0);
    }
  }

  if( build_bonds ) {
    build_bonds_by_distance( atomlist, nlist, nat );
    build_b13_from_bonds( atomlist );
    build_b14_from_bonds( atomlist );
  }

  con_table *t = build_table( atomlist );


  for(i=0;i<natoms;i++){
    //    PyObject *atom = PySequence_GetItem(atomlist, i);
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    int atom_id = PyInt_AsLong(PyObject_GetAttrString(atom,"id")) - 1;
    PyObject *neighbor_list = PyObject_GetAttrString(atom,"neighbors");
    for(k=0;k<nat[i];k++){
      int id = nlist[i][k];
      if( ! id_in_con_table( t, atom_id, id ) ){
	PyObject *neighbor_atom = PySequence_Fast_GET_ITEM(atomlist, id);
	PyObject *n_neighborlist = PyObject_GetAttrString(neighbor_atom,"neighbors");
	PyList_Append( neighbor_list, neighbor_atom);
	PyList_Append( n_neighborlist, atom);
      }
    }
  }
  delete_con_table(t);
  for(i=0;i<natoms;i++){
    free(nlist[i]);
  }

  return Py_BuildValue("i",1); //atomlist;
}




PyObject *apply_rotation( PyObject *self, PyObject *args)
{
  PyObject *Rotation;
  PyObject *py_v;
  real phi;
  if(!PyArg_ParseTuple(args,"OOd",&Rotation,&py_v, &phi))
    return NULL;

  PyObject *rm1 = PyObject_GetAttrString(Rotation,"m1");
  PyObject *rm2 = PyObject_GetAttrString(Rotation,"m2");

  matrix m1, m2;
  PyObject2matrix(rm1, m1);
  PyObject2matrix(rm2, m2);
  rvec v, v2;

  PyObject *py_v2 = PyObject_GetAttrString(Rotation,"v2");

  Pyvec2rvec(py_v, v);
  Pyvec2rvec(py_v2, v2);

  rvec vec;
  rvec_sub(v, v2, vec );

  rvec b, d, a, c, e;
  mvmul(m1, vec, b);
  mvmul(m2, vec, d);
  real cc = cos(phi);
  svmul( cc, vec, a);
  svmul( -cc, b, c);
  svmul( sin(phi), d, e);

  clear_rvec( vec );
  rvec_add( a, b, vec);
  rvec_add( vec, c, vec );
  rvec_add( vec, e, vec );

  clear_rvec(v);
  rvec_add( v2, vec, v);
  return Py_BuildValue("[ddd]", v[XX], v[YY], v[ZZ] );
 }
