/*
 * fbond.c
 *
 * Routines to evaluate force for bonded atom interactions.
 */

#include <math.h>
#include "fbond.h"


/*
 * prototypes for internal functions that compute a single interaction
 */
static double eval_bond(Force *, const MD_Bond *, const MD_Bond_Param *);
static double eval_angle(Force *, const MD_Angle *, const MD_Angle_Param *);
static double eval_dihed(Force *, const MD_Dihed *, const MD_Dihed_Param *);
  /* note that impropers are computed identically to dihedrals */


double deven_compute_bonds(Force *force)
{
  const MD_Bond_Param *bondprm = force->data->bondprm;
  const MD_Bond *bond = force->data->bond;
  const MD_Int nbonds  = force->data->nbonds;
  double pe_bond = 0.0;
  MD_Int i;

  for (i = 0;  i < nbonds;  i++, bond++) {
    pe_bond += eval_bond(force, bond, bondprm + bond->param);
  }
  return pe_bond;
}



double deven_compute_angles(Force *force)
{
  const MD_Angle_Param *angleprm = force->data->angleprm;
  const MD_Angle *angle = force->data->angle;
  const MD_Int nangles = force->data->nangles;
  double pe_angle = 0.0;
  MD_Int i;

  for (i = 0;  i < nangles;  i++, angle++) {
    pe_angle += eval_angle(force, angle, angleprm + angle->param);
  }
  return pe_angle;
}



double deven_compute_dihedrals(Force *force)
{
  const MD_Dihed_Param *dihedprm = force->data->dihedprm;
  const MD_Dihed *dihed = force->data->dihed;
  const MD_Int ndiheds = force->data->ndiheds;
  double pe_dihed = 0.0;
  MD_Int i;

  for (i = 0;  i < ndiheds;  i++, dihed++) {
    pe_dihed += eval_dihed(force, dihed, dihedprm + dihed->param);
  }
  return pe_dihed;
}



double deven_compute_impropers(Force *force)
{
  const MD_Impr_Param *imprprm = force->data->imprprm;
  const MD_Impr *impr = force->data->impr;
  const MD_Int nimprs = force->data->nimprs;
  double pe_impr = 0.0;
  MD_Int i;

  for (i = 0;  i < nimprs;  i++, impr++) {
    pe_impr += eval_dihed(force, (MD_Dihed *) impr,
                          (MD_Dihed_Param *) imprprm + impr->param);
  }
  return pe_impr;
}



double eval_bond(Force *force, const MD_Bond *bond, const MD_Bond_Param *prm)
{
  MD_Dvec *pos = force->data->pos;
  MD_Dvec *f = force->data->f;
  MD_Dvec r12, f12;
  double r12len, dist, coef;
  double energy;

  r12.x = pos[bond->atom[1]].x - pos[bond->atom[0]].x;
  r12.y = pos[bond->atom[1]].y - pos[bond->atom[0]].y;
  r12.z = pos[bond->atom[1]].z - pos[bond->atom[0]].z;
  r12len = sqrt(r12.x * r12.x + r12.y * r12.y + r12.z * r12.z);
  dist = r12len - prm->r0;
  coef = -2.0 * INV_KCAL_PER_MOL * prm->k * dist / r12len;
  energy = INV_KCAL_PER_MOL * prm->k * dist * dist;
  f12.x = coef * r12.x;
  f12.y = coef * r12.y;
  f12.z = coef * r12.z;
  f[bond->atom[0]].x -= f12.x;
  f[bond->atom[0]].y -= f12.y;
  f[bond->atom[0]].z -= f12.z;
  f[bond->atom[1]].x += f12.x;
  f[bond->atom[1]].y += f12.y;
  f[bond->atom[1]].z += f12.z;

  return energy;
}



double eval_angle(Force *force, const MD_Angle *angle,
                  const MD_Angle_Param *prm)
{
  MD_Dvec *pos = force->data->pos;
  MD_Dvec *f = force->data->f;
  MD_Dvec r21, r23, f1, f2, f3, r13, f13;
  double inv_r21len, inv_r23len, cos_theta, sin_theta, delta_theta, coef;
  double energy, r13len, dist;

  r21.x = pos[angle->atom[0]].x - pos[angle->atom[1]].x;
  r21.y = pos[angle->atom[0]].y - pos[angle->atom[1]].y;
  r21.z = pos[angle->atom[0]].z - pos[angle->atom[1]].z;
  r23.x = pos[angle->atom[2]].x - pos[angle->atom[1]].x;
  r23.y = pos[angle->atom[2]].y - pos[angle->atom[1]].y;
  r23.z = pos[angle->atom[2]].z - pos[angle->atom[1]].z;
  inv_r21len = 1.0 / sqrt(r21.x * r21.x + r21.y * r21.y + r21.z * r21.z);
  inv_r23len = 1.0 / sqrt(r23.x * r23.x + r23.y * r23.y + r23.z * r23.z);
  cos_theta = (r21.x * r23.x + r21.y * r23.y + r21.z * r23.z)
              * inv_r21len * inv_r23len;
  /* cos(theta) should be in [-1,1] */
  /* however, we need to correct in case of roundoff error */
  if (cos_theta > 1.0)        cos_theta = 1.0;
  else if (cos_theta < -1.0)  cos_theta = -1.0;
  sin_theta = sqrt(1.0 - cos_theta * cos_theta);
  delta_theta = acos(cos_theta) - prm->theta0;
  coef = -2.0 * INV_KCAL_PER_MOL * prm->k_theta * delta_theta / sin_theta;
  energy = INV_KCAL_PER_MOL * prm->k_theta * delta_theta * delta_theta;
  f1.x = coef * (cos_theta * r21.x * inv_r21len - r23.x * inv_r23len)
         * inv_r21len;
  f1.y = coef * (cos_theta * r21.y * inv_r21len - r23.y * inv_r23len)
         * inv_r21len;
  f1.z = coef * (cos_theta * r21.z * inv_r21len - r23.z * inv_r23len)
         * inv_r21len;
  f3.x = coef * (cos_theta * r23.x * inv_r23len - r21.x * inv_r21len)
         * inv_r23len;
  f3.y = coef * (cos_theta * r23.y * inv_r23len - r21.y * inv_r21len)
         * inv_r23len;
  f3.z = coef * (cos_theta * r23.z * inv_r23len - r21.z * inv_r21len)
         * inv_r23len;
  f2.x = -(f1.x + f3.x);
  f2.y = -(f1.y + f3.y);
  f2.z = -(f1.z + f3.z);

  /* Urey-Bradley term effects only atoms 1 and 3 */
  if (prm->k_ub != 0.0) {
    r13.x = r23.x - r21.x;
    r13.y = r23.y - r21.y;
    r13.z = r23.z - r21.z;
    r13len = sqrt(r13.x * r13.x + r13.y * r13.y + r13.z * r13.z);
    dist = r13len - prm->r_ub;
    coef = -2.0 * INV_KCAL_PER_MOL * prm->k_ub * dist / r13len;
    energy += INV_KCAL_PER_MOL * prm->k_ub * dist * dist;
    f13.x = coef * r13.x;
    f13.y = coef * r13.y;
    f13.z = coef * r13.z;
    f1.x -= f13.x;
    f1.y -= f13.y;
    f1.z -= f13.z;
    f3.x += f13.x;
    f3.y += f13.y;
    f3.z += f13.z;
  }
  f[angle->atom[0]].x += f1.x;
  f[angle->atom[0]].y += f1.y;
  f[angle->atom[0]].z += f1.z;
  f[angle->atom[1]].x += f2.x;
  f[angle->atom[1]].y += f2.y;
  f[angle->atom[1]].z += f2.z;
  f[angle->atom[2]].x += f3.x;
  f[angle->atom[2]].y += f3.y;
  f[angle->atom[2]].z += f3.z;

  return energy;
}



double eval_dihed(Force *force, const MD_Dihed *dihed,
                  const MD_Dihed_Param *prm)
{
/* have to do this because of the MG module */
#undef C

  MD_Dvec *pos = force->data->pos;
  MD_Dvec *f = force->data->f;
  MD_Dvec r12, r23, r34, A, B, C, dcosdA, dcosdB, dsindC, dsindB, f1, f2, f3;
  double rA, rB, rC, cos_phi, sin_phi, phi, delta, diff, k, K, K1;
  double energy = 0.0;
  MD_Int j, mult, n, not_too_small;

  r12.x = pos[dihed->atom[0]].x - pos[dihed->atom[1]].x;
  r12.y = pos[dihed->atom[0]].y - pos[dihed->atom[1]].y;
  r12.z = pos[dihed->atom[0]].z - pos[dihed->atom[1]].z;
  r23.x = pos[dihed->atom[1]].x - pos[dihed->atom[2]].x;
  r23.y = pos[dihed->atom[1]].y - pos[dihed->atom[2]].y;
  r23.z = pos[dihed->atom[1]].z - pos[dihed->atom[2]].z;
  r34.x = pos[dihed->atom[2]].x - pos[dihed->atom[3]].x;
  r34.y = pos[dihed->atom[2]].y - pos[dihed->atom[3]].y;
  r34.z = pos[dihed->atom[2]].z - pos[dihed->atom[3]].z;

  /* A = cross(r12, r23) */
  A.x = r12.y * r23.z - r12.z * r23.y;
  A.y = r12.z * r23.x - r12.x * r23.z;
  A.z = r12.x * r23.y - r12.y * r23.x;

  /* B = cross(r23, r34) */
  B.x = r23.y * r34.z - r23.z * r34.y;
  B.y = r23.z * r34.x - r23.x * r34.z;
  B.z = r23.x * r34.y - r23.y * r34.x;

  /* C = cross(r23, A) */
  C.x = r23.y * A.z - r23.z * A.y;
  C.y = r23.z * A.x - r23.x * A.z;
  C.z = r23.x * A.y - r23.y * A.x;

  rA = 1.0 / sqrt(A.x * A.x + A.y * A.y + A.z * A.z);
  rB = 1.0 / sqrt(B.x * B.x + B.y * B.y + B.z * B.z);
  rC = 1.0 / sqrt(C.x * C.x + C.y * C.y + C.z * C.z);

  /* normalize B */
  B.x *= rB;
  B.y *= rB;
  B.z *= rB;

  cos_phi = (A.x * B.x + A.y * B.y + A.z * B.z) * rA;
  sin_phi = (C.x * B.x + C.y * B.y + C.z * B.z) * rC;

  phi = -atan2(sin_phi, cos_phi);

  not_too_small = (fabs(sin_phi) > 0.1);
  if (not_too_small) {
    /* normalize A */
    A.x *= rA;
    A.y *= rA;
    A.z *= rA;
    dcosdA.x = rA * (cos_phi * A.x - B.x);
    dcosdA.y = rA * (cos_phi * A.y - B.y);
    dcosdA.z = rA * (cos_phi * A.z - B.z);
    dcosdB.x = rB * (cos_phi * B.x - A.x);
    dcosdB.y = rB * (cos_phi * B.y - A.y);
    dcosdB.z = rB * (cos_phi * B.z - A.z);
  }
  else {
    /* normalize C */
    C.x *= rC;
    C.y *= rC;
    C.z *= rC;
    dsindC.x = rC * (sin_phi * C.x - B.x);
    dsindC.y = rC * (sin_phi * C.y - B.y);
    dsindC.z = rC * (sin_phi * C.z - B.z);
    dsindB.x = rB * (sin_phi * B.x - C.x);
    dsindB.y = rB * (sin_phi * B.y - C.y);
    dsindB.z = rB * (sin_phi * B.z - C.z);
  }

  /* clear f1, f2, f3 before accumulating forces */
  f1.x = f1.y = f1.z = 0.0;
  f2.x = f2.y = f2.z = 0.0;
  f3.x = f3.y = f3.z = 0.0;

  mult = prm->mult;
  for (j = 0;  j < mult;  j++) {
    k  =    INV_KCAL_PER_MOL * prm->k[j];
    delta = prm->phi[j];
    n  =    prm->n[j];
    if (n) {
      K = k * (1.0 + cos(n * phi + delta));
      K1 = -n * k * sin(n * phi + delta);
    }
    else {
      diff = phi - delta;
      if      (diff < -M_PI)  diff += 2.0 * M_PI;
      else if (diff >  M_PI)  diff -= 2.0 * M_PI;
      K = k * diff * diff;
      K1 = 2.0 * k * diff;
    }
    energy += K;

    /* forces */
    if (not_too_small) {
      K1 /= sin_phi;
      f1.x += K1 * (r23.y * dcosdA.z - r23.z * dcosdA.y);
      f1.y += K1 * (r23.z * dcosdA.x - r23.x * dcosdA.z);
      f1.z += K1 * (r23.x * dcosdA.y - r23.y * dcosdA.x);

      f3.x += K1 * (r23.z * dcosdB.y - r23.y * dcosdB.z);
      f3.y += K1 * (r23.x * dcosdB.z - r23.z * dcosdB.x);
      f3.z += K1 * (r23.y * dcosdB.x - r23.x * dcosdB.y);

      f2.x += K1 * (r12.z * dcosdA.y - r12.y * dcosdA.z
                  + r34.y * dcosdB.z - r34.z * dcosdB.y);
      f2.y += K1 * (r12.x * dcosdA.z - r12.z * dcosdA.x
                  + r34.z * dcosdB.x - r34.x * dcosdB.z);
      f2.z += K1 * (r12.y * dcosdA.x - r12.x * dcosdA.y
                  + r34.x * dcosdB.y - r34.y * dcosdB.x);
    }
    else {
      /* phi is too close to 0 or pi, use cos version to avoid 1/sin */
      K1 /= -cos_phi;
      f1.x += K1 * ((r23.y * r23.y + r23.z * r23.z) * dsindC.x
                   - r23.x * r23.y * dsindC.y
                   - r23.x * r23.z * dsindC.z);
      f1.y += K1 * ((r23.z * r23.z + r23.x * r23.x) * dsindC.y
                   - r23.y * r23.z * dsindC.z
                   - r23.y * r23.x * dsindC.x);
      f1.z += K1 * ((r23.x * r23.x + r23.y * r23.y) * dsindC.z
                   - r23.z * r23.x * dsindC.x
                   - r23.z * r23.y * dsindC.y);

      /* f3 += K1 * cross(dsindB, r23) */
      f3.x += K1 * (dsindB.y * r23.z - dsindB.z * r23.y);
      f3.y += K1 * (dsindB.z * r23.x - dsindB.x * r23.z);
      f3.z += K1 * (dsindB.x * r23.y - dsindB.y * r23.x);

      f2.x += K1 * (-(r23.y * r12.y + r23.z * r12.z) * dsindC.x
                   + (2.0 * r23.x * r12.y - r12.x * r23.y) * dsindC.y
                   + (2.0 * r23.x * r12.z - r12.x * r23.z) * dsindC.z
                   + dsindB.z * r34.y - dsindB.y * r34.z);
      f2.y += K1 * (-(r23.z * r12.z + r23.x * r12.x) * dsindC.y
                   + (2.0 * r23.y * r12.z - r12.y * r23.z) * dsindC.z
                   + (2.0 * r23.y * r12.x - r12.y * r23.x) * dsindC.x
                   + dsindB.x * r34.z - dsindB.z * r34.x);
      f2.z += K1 * (-(r23.x * r12.x + r23.y * r12.y) * dsindC.z
                   + (2.0 * r23.z * r12.x - r12.z * r23.x) * dsindC.x
                   + (2.0 * r23.z * r12.y - r12.z * r23.y) * dsindC.y
                   + dsindB.y * r34.x - dsindB.x * r34.y);
    }
  }  /* end loop over multiplicity */
  f[dihed->atom[0]].x += f1.x;
  f[dihed->atom[0]].y += f1.y;
  f[dihed->atom[0]].z += f1.z;
  f[dihed->atom[1]].x += f2.x - f1.x;
  f[dihed->atom[1]].y += f2.y - f1.y;
  f[dihed->atom[1]].z += f2.z - f1.z;
  f[dihed->atom[2]].x += f3.x - f2.x;
  f[dihed->atom[2]].y += f3.y - f2.y;
  f[dihed->atom[2]].z += f3.z - f2.z;
  f[dihed->atom[3]].x += -f3.x;
  f[dihed->atom[3]].y += -f3.y;
  f[dihed->atom[3]].z += -f3.z;

  return energy;
}