// 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>


real calc_lj_energy( PyObject *atom1, PyObject *atom2, real fudgeLJ )
{
  real rij2 = distance2_from_atoms(atom1, atom2);
  real eps1 = PyFloat_AsDouble( PyObject_GetAttrString(atom1, "eps") );
  real eps2 = PyFloat_AsDouble( PyObject_GetAttrString(atom2, "eps") );
  real sigma1 = PyFloat_AsDouble( PyObject_GetAttrString(atom1, "sigma") );
  real sigma2 = PyFloat_AsDouble( PyObject_GetAttrString(atom2, "sigma") );

  real sigma_ij = 0.5* ( sigma1 + sigma2 );
  real eps_ij = sqrt (eps1*eps2 );

  real c6 = pow(sigma_ij,6)/pow(rij2,3);
  real c12 = sqr(c6);
  return 4*eps_ij*(c12-c6)*fudgeLJ;
}

real calc_coulomb_energy(PyObject *atom1, PyObject *atom2, real fudgeQQ )
{
  real rij = distance_from_atoms(atom1, atom2);
  real q1 = PyFloat_AsDouble( PyObject_GetAttrString(atom1, "q") );
  real q2 = PyFloat_AsDouble( PyObject_GetAttrString(atom2, "q") );
  real Ke = 1389.35338407;
  return Ke*q1*q2*fudgeQQ/rij;
}


real calc_bond_energy( PyObject *atom1, PyObject *atom2, real kb, real b0)
{
  real rij = distance_from_atoms(atom1, atom2);
  rij*=.1; // nm

  return 0.5*kb*sqr( rij-b0);

}

real calc_angle_energy(PyObject *atom1, PyObject *atom2, PyObject *atom3, real kb, real phi0)
{

  real angle = angle_from_atoms(atom1, atom2, atom3);
  phi0*=DEG2RAD;

  return 0.5*kb*( sqr(angle-phi0) );
}

real calc_dihedral_energy(PyObject *atom1, PyObject *atom2, PyObject *atom3, PyObject *atom4,  real *params)
{

  // n params = 6
  real dih = dihedral_from_atoms( atom1, atom2, atom3, atom4)*RAD2DEG - 180.;
  dih*=DEG2RAD;
  int i;
  real en = 0;
  for(i=0;i<6;i++){
    en+= params[i]*pow(cos(dih),i);
  }
  return en;

}

real calc_improper_energy(PyObject *atom1, PyObject *atom2, PyObject *atom3, PyObject *atom4, real kb, real phi0, int mult)
{
  real dih = dihedral_from_atoms( atom1, atom2, atom3, atom4)*RAD2DEG;

  return kb*(1+cos( (mult*dih-phi0)*DEG2RAD ) );
}


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

  PyObject *atom1, *atom2;
  real fudgeLJ;

  if(!PyArg_ParseTuple(args,"OOd",&atom1, &atom2, &fudgeLJ))
    return NULL;
  real energy = calc_lj_energy(atom1, atom2, fudgeLJ);
  return Py_BuildValue("d",energy);
}

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

  PyObject *atom1, *atom2;
  real fudgeQQ;

  if(!PyArg_ParseTuple(args,"OOd",&atom1, &atom2, &fudgeQQ))
    return NULL;
  real energy = calc_coulomb_energy(atom1, atom2, fudgeQQ);
  return Py_BuildValue("d",energy);
}


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

  PyObject *bondO; // object contains two atoms, b0 and k0

  if(!PyArg_ParseTuple(args,"O",&bondO))
    return NULL;

  real b0;
  real kb;
  PyObject *atom1 = PySequence_GetItem(bondO,0);
  PyObject *atom2 = PySequence_GetItem(bondO,1);

  b0 = PyFloat_AsDouble( PySequence_GetItem(bondO,3) );
  kb = PyFloat_AsDouble( PySequence_GetItem(bondO,4) );

  real energy = calc_bond_energy( atom1, atom2, kb, b0);
  return Py_BuildValue("d",energy);
}


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

  PyObject *bond_list; // list with bond entries

  if(!PyArg_ParseTuple(args,"O",&bond_list))
    return NULL;

  int nbonds = PySequence_Length( bond_list );
  int i;
  real energy = 0.;
  for(i=0;i<nbonds;i++){
    PyObject *bond = PySequence_Fast_GET_ITEM(bond_list, i);
    real b0;
    real kb;
    PyObject *atom1 = PySequence_GetItem(bond,0);
    PyObject *atom2 = PySequence_GetItem(bond,1);
    b0 = PyFloat_AsDouble( PySequence_GetItem(bond,3) );
    kb = PyFloat_AsDouble( PySequence_GetItem(bond,4) );
    energy += calc_bond_energy( atom1, atom2, kb, b0);
  }
  return Py_BuildValue("d",energy);
}


PyObject *wrap_calc_angle_energy(PyObject *self, PyObject *args)
{
  PyObject *angleO; // three atoms + ff
  if(!PyArg_ParseTuple(args,"O",&angleO))
    return NULL;

  real phi0;
  real kb;
  PyObject *atom1 = PySequence_GetItem(angleO,0);
  PyObject *atom2 = PySequence_GetItem(angleO,1);
  PyObject *atom3 = PySequence_GetItem(angleO,2);

  phi0 = PyFloat_AsDouble( PySequence_GetItem(angleO,4) );
  kb = PyFloat_AsDouble( PySequence_GetItem(angleO,5) );

  real energy = calc_angle_energy( atom1, atom2, atom3, kb, phi0);
  return Py_BuildValue("d",energy);
}


PyObject *wrap_total_angle_energy(PyObject *self, PyObject *args)
{
  PyObject *angle_list; // list with angle entries

  if(!PyArg_ParseTuple(args,"O",&angle_list))
    return NULL;

  int nangles = PySequence_Length( angle_list );
  int i;
  real energy = 0.;
  for(i=0;i<nangles;i++){
    PyObject *angleO = PySequence_Fast_GET_ITEM(angle_list, i);
    real phi0;
    real kb;
    PyObject *atom1 = PySequence_GetItem(angleO,0);
    PyObject *atom2 = PySequence_GetItem(angleO,1);
    PyObject *atom3 = PySequence_GetItem(angleO,2);

    phi0 = PyFloat_AsDouble( PySequence_GetItem(angleO,4) );
    kb = PyFloat_AsDouble( PySequence_GetItem(angleO,5) );

    energy += calc_angle_energy( atom1, atom2, atom3, kb, phi0);
  }
  return Py_BuildValue("d",energy);
}


PyObject *wrap_calc_dihedral_energy(PyObject *self, PyObject *args)
{
  PyObject *dihedO; // four atoms + dihed_type + 6 rb params
  if(!PyArg_ParseTuple(args,"O",&dihedO))
    return NULL;

  real rb_params[6];

  PyObject *atom1 = PySequence_GetItem(dihedO,0);
  PyObject *atom2 = PySequence_GetItem(dihedO,1);
  PyObject *atom3 = PySequence_GetItem(dihedO,2);
  PyObject *atom4 = PySequence_GetItem(dihedO,3);

  int i;
  for(i=0;i<6;i++){
    rb_params[i] =  PyFloat_AsDouble( PySequence_GetItem(dihedO,5+i) );
  }
  real energy = calc_dihedral_energy( atom1, atom2, atom3, atom4, rb_params);
  return Py_BuildValue("d",energy);
}


PyObject *wrap_total_dihedral_energy(PyObject *self, PyObject *args)
{
  PyObject *dihedral_list; // list with dihedral entries

  if(!PyArg_ParseTuple(args,"O",&dihedral_list))
    return NULL;

  int ndihedrals = PySequence_Length( dihedral_list );
  int i;
  real energy = 0.;
  for(i=0;i<ndihedrals;i++){
    real rb_params[6];
    PyObject *dihedO = PySequence_Fast_GET_ITEM(dihedral_list, i);
    PyObject *atom1 = PySequence_GetItem(dihedO,0);
    PyObject *atom2 = PySequence_GetItem(dihedO,1);
    PyObject *atom3 = PySequence_GetItem(dihedO,2);
    PyObject *atom4 = PySequence_GetItem(dihedO,3);

    int k;
    for(k=0;k<6;k++){
      rb_params[k] =  PyFloat_AsDouble( PySequence_GetItem(dihedO,5+k) );
    }
    energy += calc_dihedral_energy( atom1, atom2, atom3, atom4, rb_params);
  }
  return Py_BuildValue("d",energy);
}



PyObject *wrap_calc_improper_energy(PyObject *self, PyObject *args)
{
  PyObject *dihedO; // four atoms + dihed_type + phi0 + kb + mult
  if(!PyArg_ParseTuple(args,"O",&dihedO))
    return NULL;

  real kb;
  int mult;
  real phi0;

  PyObject *atom1 = PySequence_GetItem(dihedO,0);
  PyObject *atom2 = PySequence_GetItem(dihedO,1);
  PyObject *atom3 = PySequence_GetItem(dihedO,2);
  PyObject *atom4 = PySequence_GetItem(dihedO,3);

  phi0 = PyFloat_AsDouble( PySequence_GetItem(dihedO, 5) );
  kb = PyFloat_AsDouble( PySequence_GetItem(dihedO, 6) );
  mult = PyInt_AsLong( PySequence_GetItem(dihedO, 7) );

  real energy = calc_improper_energy( atom1, atom2, atom3, atom4, kb, phi0, mult);
  return Py_BuildValue("d",energy);
}


PyObject *wrap_total_improper_energy(PyObject *self, PyObject *args)
{
  PyObject *dihedral_list; // list with dihedral entries

  if(!PyArg_ParseTuple(args,"O",&dihedral_list))
    return NULL;

  int ndihedrals = PySequence_Length( dihedral_list );
  int i;
  real energy = 0.;
  for(i=0;i<ndihedrals;i++){
    real kb;
    int mult;
    real phi0;
    PyObject *dihedO = PySequence_Fast_GET_ITEM(dihedral_list, i);
    PyObject *atom1 = PySequence_GetItem(dihedO,0);
    PyObject *atom2 = PySequence_GetItem(dihedO,1);
    PyObject *atom3 = PySequence_GetItem(dihedO,2);
    PyObject *atom4 = PySequence_GetItem(dihedO,3);
    phi0 = PyFloat_AsDouble( PySequence_GetItem(dihedO, 5) );
    kb = PyFloat_AsDouble( PySequence_GetItem(dihedO, 6) );
    mult = PyInt_AsLong( PySequence_GetItem(dihedO, 7) );

    energy += calc_improper_energy( atom1, atom2, atom3, atom4, kb, phi0, mult);
  }
  return Py_BuildValue("d",energy);
}

real nb_lj_energy( PyObject *atomlist )
{
  int i, k;
  real energy = 0.;
  real fudgeLJ = 1.;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    PyObject *nb_list = PyObject_GetAttrString(atom,"neighbors");
    int nn = PySequence_Length( nb_list );
    for(k=0;k<nn;k++){
      PyObject *nb_atom = PySequence_Fast_GET_ITEM(nb_list,k);
      int atom_id = PyInt_AsLong( PyObject_GetAttrString(nb_atom, "id")) - 1;
      if( i < atom_id ) {
	energy += calc_lj_energy( atom, nb_atom, fudgeLJ);
      }
    }
  }
  return energy;
}

real lj14_energy( PyObject *atomlist )
{
  int i, k;
  real energy = 0.;
  real fudgeLJ = .5;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    PyObject *nb_list = PyObject_GetAttrString(atom,"b14");
    int nn = PySequence_Length( nb_list );
    for(k=0;k<nn;k++){
      PyObject *nb_atom = PySequence_Fast_GET_ITEM(nb_list,k);
      int atom_id = PyInt_AsLong( PyObject_GetAttrString(nb_atom, "id")) - 1;
      if( i < atom_id ) {
	energy += calc_lj_energy( atom, nb_atom, fudgeLJ);
      }
    }
  }
  return energy;
}

real nb_coul_energy( PyObject *atomlist )
{
  int i, k;
  real energy = 0.;
  real fudgeQQ = 1.;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    PyObject *nb_list = PyObject_GetAttrString(atom,"neighbors");
    int nn = PySequence_Length( nb_list );
    for(k=0;k<nn;k++){
      PyObject *nb_atom = PySequence_Fast_GET_ITEM(nb_list,k);
      int atom_id = PyInt_AsLong( PyObject_GetAttrString(nb_atom, "id")) - 1;
      if( i < atom_id ) {
	energy += calc_coulomb_energy( atom, nb_atom, fudgeQQ);
      }
    }
  }
  return energy;
}

real coul14_energy( PyObject *atomlist )
{
  int i, k;
  real energy = 0.;
  real fudgeQQ = 0.8333;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    PyObject *nb_list = PyObject_GetAttrString(atom,"b14");
    int nn = PySequence_Length( nb_list );
    for(k=0;k<nn;k++){
      PyObject *nb_atom = PySequence_Fast_GET_ITEM(nb_list,k);
      int atom_id = PyInt_AsLong( PyObject_GetAttrString(nb_atom, "id")) - 1;
      if( i < atom_id ) {
	energy += calc_coulomb_energy( atom, nb_atom, fudgeQQ);
      }
    }
  }
  return energy;
}


real nb_energy( PyObject *atomlist )
{
  int i, k;
  real energy = 0.;
  real fudgeQQ = 1.;
  real fudgeLJ = 1.;
  int natoms = PySequence_Length(atomlist);
  for(i=0;i<natoms;i++){
    PyObject *atom = PySequence_Fast_GET_ITEM(atomlist, i);
    PyObject *nb_list = PyObject_GetAttrString(atom,"neighbors");
    int nn = PySequence_Length( nb_list );
    for(k=0;k<nn;k++){
      PyObject *nb_atom = PySequence_Fast_GET_ITEM(nb_list,k);
      int atom_id = PyInt_AsLong( PyObject_GetAttrString(nb_atom, "id")) - 1;
      if( i < atom_id ) {
	energy += calc_coulomb_energy( atom, nb_atom, fudgeQQ);
	energy += calc_lj_energy( atom, nb_atom, fudgeLJ);
      }
    }
  }
  return energy;
}


PyObject *wrap_nb_lj_energy(PyObject *self, PyObject *args)
{
  PyObject *atomlist;

  if(!PyArg_ParseTuple(args,"O",&atomlist))
    return NULL;

  real energy = nb_lj_energy( atomlist );
  return Py_BuildValue("d",energy);
}

PyObject *wrap_lj14_energy(PyObject *self, PyObject *args)
{
  PyObject *atomlist;

  if(!PyArg_ParseTuple(args,"O",&atomlist))
    return NULL;

  real energy = lj14_energy( atomlist );
  return Py_BuildValue("d",energy);
}

PyObject *wrap_nb_coul_energy(PyObject *self, PyObject *args)
{
  PyObject *atomlist;

  if(!PyArg_ParseTuple(args,"O",&atomlist))
    return NULL;

  real energy = nb_coul_energy( atomlist );
  return Py_BuildValue("d",energy);
}


PyObject *wrap_coul14_energy(PyObject *self, PyObject *args)
{
  PyObject *atomlist;

  if(!PyArg_ParseTuple(args,"O",&atomlist))
    return NULL;

  real energy = coul14_energy( atomlist );
  return Py_BuildValue("d",energy);
}

PyObject *wrap_nb_energy(PyObject *self, PyObject *args)
{
  PyObject *atomlist;

  if(!PyArg_ParseTuple(args,"O",&atomlist))
    return NULL;

  real energy = nb_energy( atomlist );
  return Py_BuildValue("d",energy);
}
