/*
 *   term.c - Energy/force terms, and calculations for them
 *   Copyright (C) 1997 Will Ware <wware@world.std.com>
 *   
 *   This program is free software; you can redistribute it and/or
 *   modify it under the terms of the GNU General Public License
 *   as published by the Free Software Foundation; either version 2
 *   of the License, or (at your option) any later version.
 *   
 *   This program is distributed in the hope that it will be useful,
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *   GNU General Public License for more details.
 *   
 *   You should have received a copy of the GNU General Public License
 *   along with this program; if not, write to the Free Software
 *   Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <stdio.h>
#include <math.h>

#include "debug.h"
#include "atom.h"
#include "term.h"

/********************************************************************/

/*
 * Physical Units (these are all consistent with one another):
 * Distances are in nanometers (10^-9 m)
 * Times are in picoseconds (10^-12 sec)
 * Masses are in units of 10^-27 kg
 * Forces are in piconewtons (10^-12 N)
 * Energies are in milli-attojoules (10^-21 J)
 * Spring constants are in millinewtons-per-meter or pN/nm
 * Charges are in single-proton charges (0.160206*10^-18 coulombs)
 */

/*
 * Might make sense to switch to these in the future:
 * Distances are in nanometers (10^-9 m)
 * Times are in femtoseconds (10^-15 sec)
 * Masses are in units of 10^-27 kg
 * Forces are in micronewtons (10^-6 N)
 * Energies are in femtojoules (10^-15 J)
 * Spring constants are in kilonewtons-per-meter or pN/nm
 * Indications that this is a good idea: spring constants are now
 * huge.
 */

/********************************************************************/
/* Handy little macros for terse vector hacking */

#define COPY(dest,x)    dest[0] = x[0]; \
dest[1] = x[1]; \
dest[2] = x[2];

#define ADD(sum,x,y)    sum[0] = x[0] + y [0]; \
sum[1] = x[1] + y [1]; \
sum[2] = x[2] + y [2];

#define DIFF(diff,x,y)  diff[0] = x[0] - y [0]; \
diff[1] = x[1] - y [1]; \
diff[2] = x[2] - y [2];

#define DOT(x,y)  ((x[0] * y [0]) + (x[1] * y [1]) + (x[2] * y [2]))

#define CROSS(z,x,y) \
z[0] = x[1] * y[2] - x[2] * y[1]; \
z[1] = x[2] * y[0] - x[0] * y[2]; \
z[2] = x[0] * y[1] - x[1] * y[0];

#define LEN(x)    (sqrt (DOT (x, x)))

#define LEN2(x,y)   (sqrt (DOT(x,x) * DOT(y,y)))	/* efficiency hack */

#define SCALE(v,m)   v[0] *= m; v[1] *= m; v[2] *= m;

#define MUL(v,x,m)   v[0] = m * x[0]; v[1] = m * x[1]; v[2] = m * x[2];

#define LIN_COMB(z,a,x,b,y) \
z[0] = a * x[0] + b * y[0]; \
z[1] = a * x[1] + b * y[1]; \
z[2] = a * x[2] + b * y[2];

#define ADD_LIN_COMB(z,a,x,b,y) \
z[0] += a * x[0] + b * y[0]; \
z[1] += a * x[1] + b * y[1]; \
z[2] += a * x[2] + b * y[2];

/********************************************************************/

static double
radial_energy (double (*efunc) (double, double, double, double, double),
	       double c1, double c2, double c3, double c4,
	       atom * a1, atom * a2)
{
  double diff[3], len;

  DIFF (diff, a1->x, a2->x);
  len = LEN (diff);
  return (*efunc) (len, c1, c2, c3, c4);
}

static void
radial_force (double (*ederiv) (double, double, double, double, double),
	      double c1, double c2, double c3, double c4,
	      atom * a1, atom * a2)
{
  double diff[3], len, m;

  DIFF (diff, a1->x, a2->x);
  len = LEN (diff);
  m = (*ederiv) (len, c1, c2, c3, c4) / len;
  SCALE (diff, m);
  DIFF (a1->f, a1->f, diff);
  ADD (a2->f, a2->f, diff);
}

/********************************************************************/

static struct
{
  enum SPECIES species1, species2;
  double ks, r0, De, beta;
}
length_term_coeffs[] =
{
  /*              ks       r0       De     beta */
  /*             pN/nm     nm       maJ    nm^-1 */
  {
    C_SP3, C_SP3, 440000.0, 0.1523, 545.0, 19.89
  }
  ,
  {
    C_SP2, C_SP2, 960000.0, 0.1337, 1207.0, 19.94
  }
  ,
  {
    C_SP, C_SP, 1560000.0, 0.1212, 1614.0, 21.98
  }
  ,
  {
    C_SP3, H, 460000.0, 0.1113, 671.0, 18.51
  }
  ,
  {
    C_SP3, O_SP3, 536000.0, 0.1402, 1343.0, 20.05
  }
  /* just a few for now, add more when it's working */
};

static int num_length_coeffs =
sizeof (length_term_coeffs) / sizeof (length_term_coeffs[0]);

#define KCUBIC 20.0		/* (0.05 nanometer)^-1 */

static double
length_energy_helper (double r, double ks,
		      double r0, double De, double beta)
{
  double rdiff = (r - r0), expr;

  if (rdiff < 0.0)
    return 0.5 * ks * rdiff * rdiff * (1 - KCUBIC * rdiff);
  else
    {
      expr = 1.0 - exp (-beta * rdiff);
      return De * expr * expr;
    }
}

static double
length_energy_function (term * t)
{
  double r;

  r = radial_energy (length_energy_helper,
		     t->coeffs[0],
		     t->coeffs[1],
		     t->coeffs[2],
		     t->coeffs[3],
		     t->atoms[0],
		     t->atoms[1]);
  DBGPRINTF1 ("Length energy: %f\n", r);
  return r;
}

static double
length_energy_deriv (double r, double ks,
		     double r0, double De, double beta)
{
  double rdiff = (r - r0), expr;

  if (rdiff < 0.0)
    return ks * rdiff * (1.0 - 1.5 * KCUBIC * rdiff);
  else
    {
      expr = exp (-beta * rdiff);
      return 2.0 * De * beta * expr * (1.0 - expr);
    }
}

static void
length_add_forces (term * t)
{
  radial_force (length_energy_deriv,
		t->coeffs[0],
		t->coeffs[1],
		t->coeffs[2],
		t->coeffs[3],
		t->atoms[0],
		t->atoms[1]);
}

void
length_term (term * t, atom * a1, atom * a2)
{
  int i, found;

  t->type = LENGTH;
  t->atoms[0] = a1;
  t->atoms[1] = a2;

  for (i = found = 0; i < num_length_coeffs && !found; i++)
    if ((length_term_coeffs[i].species1 == a1->species &&
	 length_term_coeffs[i].species2 == a2->species) ||
	(length_term_coeffs[i].species1 == a2->species &&
	 length_term_coeffs[i].species2 == a1->species))
      {
	t->coeffs[0] = length_term_coeffs[i].ks;
	t->coeffs[1] = length_term_coeffs[i].r0;
	t->coeffs[2] = length_term_coeffs[i].De;
	t->coeffs[3] = length_term_coeffs[i].beta;
	found = 1;
      }

  if (!found)
    {
      t->coeffs[0] = 440000.0;
      t->coeffs[1] = 0.1523;
      t->coeffs[2] = 545.0;
      t->coeffs[3] = 19.89;
    }
}

/********************************************************************/

static struct
{
  enum SPECIES species1, species2, species3;
  double kth, th0;
}
angle_term_coeffs[] =
{
  /*                       kth       th0 */
  /*                    maJ/rad^2  radians */
  {
    C_SP3, C_SP3, C_SP3, 450.0, 1.911
  }
  ,
  {
    C_SP3, C_SP3, H, 360.0, 1.909
  }
  ,
  {
    H, C_SP3, H, 320.0, 1.909
  }
  ,
  {
    C_SP3, C_SP2, C_SP3, 450.0, 2.046
  }
  ,
  {
    C_SP2, C_SP3, C_SP2, 450.0, 1.911
  }
  ,
  {
    C_SP3, C_SP2, C_SP2, 550.0, 2.119
  }
  ,
  {
    C_SP2, C_SP2, H, 360.0, 2.094
  }
  ,
  {
    C_SP2, C_SP2, C_SP2, 430.0, 2.094
  }
  ,
  {
    C_SP3, C_SP, C_SP, 200.0, 3.142
  }
  ,
  {
    C_SP3, C_SP2, O_SP2, 460.0, 2.138
  }
  ,
  {
    C_SP3, O_SP3, C_SP3, 770.0, 1.864
  }
  ,
  {
    C_SP3, N_SP3, C_SP3, 630.0, 1.880
  }
  ,
  {
    C_SP3, C_SP3, O_SP3, 700.0, 1.876
  }
  ,
  {
    C_SP3, C_SP3, N_SP3, 570.0, 1.911
  }
};

static int num_angle_coeffs =
sizeof (angle_term_coeffs) / sizeof (angle_term_coeffs[0]);

static double
angle_energy_function (term * t)
{
  double kth, th0, th, r;
  double adiff[3], bdiff[3];

  kth = t->coeffs[0];
  th0 = t->coeffs[1];
  DIFF (adiff, t->atoms[0]->x, t->atoms[1]->x);
  DIFF (bdiff, t->atoms[2]->x, t->atoms[1]->x);
  th = DOT (adiff, bdiff) / sqrt (DOT (adiff, adiff) * DOT (bdiff, bdiff));
  r = 0.5 * kth * (th - th0) * (th - th0);
  DBGPRINTF1 ("Angle energy: %f\n", r);
  return r;
}

static void
angle_add_forces (term * t)
{
  double kth, th0, th, tdif, du_dth;
  double adiff[3], bdiff[3], L1[3], L3[3];
  double f1[3], f2[3], f3[3];

  kth = t->coeffs[0];
  th0 = t->coeffs[1];
  th = DOT (adiff, bdiff) / sqrt (DOT (adiff, adiff) * DOT (bdiff, bdiff));
  tdif = th - th0;
  du_dth = kth * (tdif * (1.0 + 1.508 * tdif * tdif));

  DIFF (adiff, t->atoms[0]->x, t->atoms[1]->x);
  DIFF (bdiff, t->atoms[2]->x, t->atoms[1]->x);

  LIN_COMB (L1,
	    DOT (adiff, adiff), bdiff,
	    -DOT (adiff, bdiff), adiff);
  MUL (f1, L1, (du_dth / LEN2 (L1, adiff)));

  LIN_COMB (L3,
	    DOT (bdiff, bdiff), adiff,
	    -DOT (adiff, bdiff), bdiff);
  MUL (f3, L3, (du_dth / LEN2 (L3, bdiff)));

  LIN_COMB (f2, -1.0, f1, -1.0, f3);

  ADD (t->atoms[0]->f, t->atoms[0]->f, f1);
  ADD (t->atoms[1]->f, t->atoms[1]->f, f2);
  ADD (t->atoms[2]->f, t->atoms[2]->f, f3);
}

void
angle_term (term * t, atom * a1, atom * a2, atom * a3)
{
  int i, found;

  t->type = ANGLE;
  t->atoms[0] = a1;
  t->atoms[1] = a2;
  t->atoms[2] = a3;

  for (i = found = 0; i < num_angle_coeffs && !found; i++)
    if ((angle_term_coeffs[i].species1 == a1->species &&
	 angle_term_coeffs[i].species2 == a2->species &&
	 angle_term_coeffs[i].species3 == a3->species) ||
	(angle_term_coeffs[i].species1 == a3->species &&
	 angle_term_coeffs[i].species2 == a2->species &&
	 angle_term_coeffs[i].species3 == a1->species))
      {
	t->coeffs[0] = angle_term_coeffs[i].kth;
	t->coeffs[1] = angle_term_coeffs[i].th0;
	found = 1;
      }

  if (!found)
    {
      t->coeffs[0] = 440000.0;
      t->coeffs[1] = 0.1523;
    }
}

/********************************************************************/

static struct
{
  enum SPECIES species1, species2, species3, species4;
  double v1, v2, v3;
}
torsion_term_coeffs[] =
{
  /*                            maJ   maJ   maJ */
  {
    C_SP3, C_SP3, C_SP3, C_SP3, 1.39, 1.88, 0.65
  }
  ,
  {
    C_SP3, C_SP3, C_SP3, H, 0.0, 0.0, 1.85
  }
  ,
  {
    H, C_SP3, C_SP3, H, 0.0, 0.0, 1.65
  }
  ,
  {
    C_SP3, C_SP2, C_SP2, C_SP3, -0.69, 69.47, 0.0
  }
  ,
  {
    C_SP2, C_SP2, C_SP2, C_SP2, -6.46, 55.58, 0.0
  }
  ,
  {
    H, C_SP2, C_SP2, H, 0.0, 104.21, 0.0
  }
};

static int num_torsion_coeffs =
sizeof (torsion_term_coeffs) / sizeof (torsion_term_coeffs[0]);

static double
safe_acos (double x)
{
  if (x >= 1.0)
    return acos (1.0);
  if (x <= -1.0)
    return acos (-1.0);
  return acos (x);
}

static double
torsion_energy_function (term * t)
{
  double v1, v2, v3, w, e;
  double p[3], q[3], r[3];

  v1 = t->coeffs[0];
  v2 = t->coeffs[1];
  v3 = t->coeffs[2];

  DIFF (p, t->atoms[0]->x, t->atoms[1]->x);
  DIFF (q, t->atoms[1]->x, t->atoms[2]->x);
  DIFF (r, t->atoms[3]->x, t->atoms[2]->x);

  LIN_COMB (p,
	    1.0, p,
	    -DOT (p, q) / DOT (q, q), q);
  MUL (p, p, 1.0 / LEN (p));

  LIN_COMB (r,
	    1.0, r,
	    -DOT (r, q) / DOT (q, q), q);
  MUL (r, r, 1.0 / LEN (r));

  w = safe_acos (DOT (p, r));

  e = 0.5 * (v1 * (1.0 + cos (w)) +
	     v2 * (1.0 - cos (2 * w)) +
	     v3 * (1.0 + cos (3 * w)));
  DBGPRINTF1 ("Torsion energy: %f\n", e);
  return e;
}

static void
torsion_add_forces (term * t)
{
  double v1, v2, v3, pp, qq, rr, pq, qr, w, du_dw;
  double alpha, beta, gamma;
  double p[3], q[3], r[3], vm1[3], vq1[3], vm2[3], vq2[3], fm[3], fq[3];

  v1 = t->coeffs[0];
  v2 = t->coeffs[1];
  v3 = t->coeffs[2];

  DIFF (p, t->atoms[0]->x, t->atoms[1]->x);
  DIFF (q, t->atoms[2]->x, t->atoms[1]->x);
  DIFF (r, t->atoms[3]->x, t->atoms[2]->x);

  pp = DOT (p, p);
  qq = DOT (q, q);
  rr = DOT (r, r);
  pq = DOT (p, q);
  qr = DOT (q, r);

  alpha = sqrt (pq * pq / qq);
  beta = sqrt (qq * qq);
  gamma = sqrt (qr * qr / qq);

  CROSS (vm1, q, p);
  CROSS (vq1, r, q);
  MUL (vm2, vm1, 1.0 / LEN (vm1));
  MUL (vq2, vq1, 1.0 / LEN (vq1));

  w = safe_acos (DOT (vm2, vq2));
  du_dw = -0.5 * (v1 * sin (w) +
		  -2.0 * v2 * sin (2 * w) +
		  3.0 * v3 * sin (3 * w));

  MUL (fm, vm2, du_dw / sqrt (pp - (pq * pq) / qq));
  MUL (fq, vq2, du_dw / sqrt (rr - (qr * qr) / qq));

  ADD (t->atoms[0]->f, t->atoms[0]->f, fm);
  ADD (t->atoms[3]->f, t->atoms[3]->f, fq);
  ADD_LIN_COMB (t->atoms[2]->f,
		alpha / beta,
		fm,
		(beta + gamma) / beta,
		fq);
  ADD_LIN_COMB (t->atoms[1]->f,
		gamma / beta,
		fq,
		(beta + alpha) / beta,
		fm);
}

void
torsion_term (term * t, atom * a1, atom * a2, atom * a3, atom * a4)
{
  int i, found;

  t->type = ANGLE;
  t->atoms[0] = a1;
  t->atoms[1] = a2;
  t->atoms[2] = a3;
  t->atoms[3] = a4;

  for (i = found = 0; i < num_torsion_coeffs && !found; i++)
    if ((torsion_term_coeffs[i].species1 == a1->species &&
	 torsion_term_coeffs[i].species2 == a2->species &&
	 torsion_term_coeffs[i].species3 == a3->species &&
	 torsion_term_coeffs[i].species4 == a4->species) ||
	(torsion_term_coeffs[i].species1 == a4->species &&
	 torsion_term_coeffs[i].species2 == a3->species &&
	 torsion_term_coeffs[i].species3 == a2->species &&
	 torsion_term_coeffs[i].species4 == a1->species))
      {
	t->coeffs[0] = torsion_term_coeffs[i].v1;
	t->coeffs[1] = torsion_term_coeffs[i].v2;
	t->coeffs[2] = torsion_term_coeffs[i].v3;
	found = 1;
      }

  if (!found)
    {
      t->coeffs[0] = 440.0;
      t->coeffs[1] = 0.1523;
      t->coeffs[2] = 0.1523;
    }
}

/********************************************************************/

/* units here are (maJ * nm) / (e * e), where e is charge on a proton */
#define ELECTRO_CONSTANT  (8.9876 * 0.160206 * 0.160206 * 1000.0)

static double
electro_energy_helper (double r, double q1,
		       double q2, double dummy1, double dummy2)
{
  return (ELECTRO_CONSTANT * q1 * q2) / r;
}

static double
electrostatic_energy_function (term * t)
{
  double e;

  e = radial_energy (electro_energy_helper,
		     t->coeffs[0],
		     t->coeffs[1],
		     0.0,
		     0.0,
		     t->atoms[0],
		     t->atoms[1]);
  DBGPRINTF1 ("Electrostatic energy: %f\n", e);
  return e;
}

static double
electro_energy_deriv (double r, double q1,
		      double q2, double dummy1, double dummy2)
{
  return (-ELECTRO_CONSTANT * q1 * q2) / (r * r);
}

static void
electrostatic_add_forces (term * t)
{
  radial_force (electro_energy_deriv,
		t->coeffs[0],
		t->coeffs[1],
		0.0,
		0.0,
		t->atoms[0],
		t->atoms[1]);
}

void
electrostatic_term (term * t, atom * a1, atom * a2)
{
  t->type = ELECTROSTATIC;
  t->atoms[0] = a1;
  t->atoms[1] = a2;
  t->coeffs[0] = a1->charge;
  t->coeffs[1] = a2->charge;
}

/********************************************************************/

static double
vdw_energy_helper (double r, double evdw0,
		   double rvdw0, double dummy1, double dummy2)
{
  double ratio = rvdw0 / r, ra3, ra6;
  ra3 = ratio * ratio * ratio;
  ra6 = ra3 * ra3;
  return evdw0 * ra6 * (ra6 - 2.0);
}

static double
vdw_energy_function (term * t)
{
  double e;

  e = radial_energy (vdw_energy_helper,
		     0.5 * (t->atoms[0]->evdw + t->atoms[1]->evdw),
		     t->atoms[0]->rvdw + t->atoms[1]->rvdw,
		     0.0,
		     0.0,
		     t->atoms[0],
		     t->atoms[1]);
  DBGPRINTF1 ("Van der Waals energy: %f\n", e);
  return e;
}

static double
vdw_energy_deriv (double r, double evdw0,
		  double rvdw0, double dummy1, double dummy2)
{
  double ratio = rvdw0 / r, ra3, ra6;
  ra3 = ratio * ratio * ratio;
  ra6 = ra3 * ra3;
  return -12.0 * evdw0 * ratio * ra6 * (ra6 - 1.0);
}

static void
vdw_add_forces (term * t)
{
  radial_force (vdw_energy_deriv,
		0.5 * (t->atoms[0]->evdw + t->atoms[1]->evdw),
		t->atoms[0]->rvdw + t->atoms[1]->rvdw,
		0.0,
		0.0,
		t->atoms[0],
		t->atoms[1]);
}

void
vdw_term (term * t, atom * a1, atom * a2)
{
  t->type = VDW;
  t->atoms[0] = a1;
  t->atoms[1] = a2;
}

/********************************************************************/

static double
sprd_energy_helper (double r, double k,
		    double r0, double dummy1, double dummy2)
{
  /* note: the damper makes no contribution to potential energy */
  double rdiff = r - r0;
  return 0.5 * k * rdiff * rdiff;
}

static double
spring_damper_energy_function (term * t)
{
  double e;

  e = radial_energy (sprd_energy_helper,
		     t->coeffs[0],
		     t->coeffs[1],
		     t->coeffs[2],
		     0.0,
		     t->atoms[0],
		     t->atoms[1]);
  DBGPRINTF1 ("Spring damper energy: %f\n", e);
  return e;
}

/* Because the damping force depends on velocities, we can't just use the
 * normal radial-force routine.
 */

static void
spring_damper_add_forces (term * t)
{
  double k, r0, d, diff[3], vdiff[3], r, rel_v, m;

  k = t->coeffs[0];
  r0 = t->coeffs[1];
  d = t->coeffs[2];
  DIFF (diff, t->atoms[0]->x, t->atoms[1]->x);
  DIFF (vdiff, t->atoms[0]->v, t->atoms[1]->v);
  r = LEN (diff);
  rel_v = DOT (diff, vdiff) / r;
  m = (k * (r - r0) + d * rel_v) / r;
  SCALE (diff, m);
  DIFF (t->atoms[0]->f, t->atoms[0]->f, diff);
  ADD (t->atoms[1]->f, t->atoms[1]->f, diff);
}

void
spring_damper_term (term * t, atom * a1, atom * a2,
		    double k, double r0, double d)
{
  t->type = SPRING_DAMPER;
  t->atoms[0] = a1;
  t->atoms[1] = a2;
  t->coeffs[0] = k;
  t->coeffs[1] = r0;
  t->coeffs[2] = d;
}

/********************************************************************/

double
term_energy (term * t)
{
  switch (t->type)
    {
    case INVALID:
      fprintf (stderr, "Bad term!\n");
      return 0.0;
    case LENGTH:
      return length_energy_function (t);
    case ANGLE:
      return angle_energy_function (t);
    case TORSION:
      return torsion_energy_function (t);
    case ELECTROSTATIC:
      return electrostatic_energy_function (t);
    case VDW:
      return vdw_energy_function (t);
    case SPRING_DAMPER:
      return spring_damper_energy_function (t);
    }
  return 0.0;
}

void
term_add_forces (term * t)
{
  switch (t->type)
    {
    case INVALID:
      fprintf (stderr, "Bad term!\n");
      break;
    case LENGTH:
      length_add_forces (t);
      break;
    case ANGLE:
      angle_add_forces (t);
      break;
    case TORSION:
      torsion_add_forces (t);
      break;
    case ELECTROSTATIC:
      electrostatic_add_forces (t);
      break;
    case VDW:
      vdw_add_forces (t);
      break;
    case SPRING_DAMPER:
      spring_damper_add_forces (t);
      break;
    }
  /* if not a known choice, do nothing */
}