/*
 * integ.c
 *
 * Integration routines (leapfrog / velocity Verlet integrator).
 */

#include <stdlib.h>
#include <string.h>
#include "data.h"
#include "force.h"
#include "integ.h"


/*
 * prototypes for internal functions
 */
static void set_reduc(Data *data);
static double compute_ke(Data *data);
static void compute_shadow(Data *data);
static MD_Errcode run_shadow(MD_Sim *sim, MD_Int numsteps);



/*
 * called by front end through MD_init
 *
 * init array of inverse masses
 * re-init the force computation object
 * compute initial forces and reductions for system
 */
MD_Errcode deven_init(MD_Sim *sim)
{
  Data *data = (Data *) MD_engine(sim);
  MD_Atom *atom = data->atom;
  MD_Int natoms = data->natoms;
  MD_Int k;

  /* initialize trajectories if shadow energy results desired */
  data->posarr[0] = data->pos;
  data->velarr[0] = data->vel;
  if (data->is_shadow) {
    for (k = 1;  k < NUMARR;  k++) {
      void *tmp;
      tmp = realloc(data->posarr[k], natoms * sizeof(MD_Dvec));
      if (tmp == NULL) {
        MD_ERR(data->sim, MD_engine_name(data->sim), MD_ERR_MEMALLOC,
               "deven_init call to realloc");
        return MD_FAIL;
      }
      data->posarr[k] = (MD_Dvec *) tmp;
      tmp = realloc(data->velarr[k], natoms * sizeof(MD_Dvec));
      if (tmp == NULL) {
        MD_ERR(data->sim, MD_engine_name(data->sim), MD_ERR_MEMALLOC,
               "deven_init call to realloc");
        return MD_FAIL;
      }
      data->velarr[k] = (MD_Dvec *) tmp;
    }
  }

  /* initialize step number counter */
  data->stepnum = data->firststepnum;
  if (MD_set_stepnum(sim, data->stepnum))  return MD_FAIL;

  /* check to make sure simulation parameters are valid */
  if (deven_data_validate_simparam(data))  return MD_FAIL;

  /* in case front end calls init twice, destroy before re-init */
  deven_force_destroy(data->force);
  if (deven_force_init(data->force, data))  return MD_FAIL;

  /* find initial force */
  if (MD_fcallback(sim, data->stepnum, 0, 1) || MD_waitcb(sim)) {
    return MD_FAIL;
  }
  if (deven_force_compute(data->force))  return MD_FAIL;

  /* find reductions for system */
  data->energy[KE] = compute_ke(data);
  data->energy[TE] = data->energy[PE] + data->energy[KE];
  set_reduc(data);

  /* process callbacks for initial step */
  if (MD_callback(sim, data->stepnum) || MD_waitcb(sim))  return MD_FAIL;

  return 0;
}



/*
 * called by front end through MD_run
 *
 * perform leapfrog (velocity Verlet) integration of system for numsteps
 */
MD_Errcode deven_run(MD_Sim *sim, MD_Int numsteps)
{
  Data *data = (Data *) MD_engine(sim);
  const double dt = data->timestep;  /* femptoseconds */
  const double half_dt = 0.5 * dt;
  const MD_Int natoms = data->natoms;
  MD_Dvec *pos = data->pos;
  MD_Dvec *vel = data->vel;
  MD_Dvec *f = data->f;
  const MD_Atom *atom = data->atom;
  double inv_mass;
  MD_Int n, k;

  /* if computing shadow energies, run a different integration routine */
  if (data->is_shadow)  return run_shadow(sim, numsteps);

  for (n = 0;  n < numsteps;  n++) {

    /* half kick */
    for (k = 0;  k < natoms;  k++) {
      inv_mass = 1.0 / atom[k].m;
      vel[k].x += half_dt * f[k].x * inv_mass;
      vel[k].y += half_dt * f[k].y * inv_mass;
      vel[k].z += half_dt * f[k].z * inv_mass;
    }

    /* drift */
    for (k = 0;  k < natoms;  k++) {
      pos[k].x += dt * vel[k].x;
      pos[k].y += dt * vel[k].y;
      pos[k].z += dt * vel[k].z;
    }

    /* compute the force */
    data->stepnum++;
    if (MD_fcallback(sim, data->stepnum, 0, 1) || MD_waitcb(sim)) {
      return MD_FAIL;
    }
    if (deven_force_compute(data->force))  return MD_FAIL;

    /* half kick */
    for (k = 0;  k < natoms;  k++) {
      inv_mass = 1.0 / atom[k].m;
      vel[k].x += half_dt * f[k].x * inv_mass;
      vel[k].y += half_dt * f[k].y * inv_mass;
      vel[k].z += half_dt * f[k].z * inv_mass;
    }

    /* find reductions for system */
    data->energy[KE] = compute_ke(data);
    data->energy[TE] = data->energy[PE] + data->energy[KE];
    set_reduc(data);

    if (MD_callback(sim, data->stepnum) || MD_waitcb(sim))  return MD_FAIL;
  }
  return 0;
}



/*
 * perform leapfrog integration with shadow energy computation
 */
MD_Errcode run_shadow(MD_Sim *sim, MD_Int numsteps)
{
  Data *data = (Data *) MD_engine(sim);
  const double dt = data->timestep;  /* femptoseconds */
  const double half_dt = 0.5 * dt;
  const MD_Int natoms = data->natoms;
  MD_Dvec *pos = data->pos;
  MD_Dvec *vel = data->vel;
  MD_Dvec *oldpos, *oldvel;
  MD_Double *beta = data->beta_arr[data->arr_index];
  MD_Double *oldbeta;
  MD_Dvec *f = data->f;
  const MD_Atom *atom = data->atom;
  double inv_mass, pos_dot_f;
  MD_Int n, k;

  for (n = 0;  n < numsteps;  n++) {

    /*
     * first halfstep of leapfrog
     */
    data->arr_index = (data->arr_index + 1) & NUMARR_MINUS_1;
    oldpos = pos;  pos = data->posarr[data->arr_index];
    oldvel = vel;  vel = data->velarr[data->arr_index];
    oldbeta = beta;  beta = data->beta_arr[data->arr_index];

    /* half kick */
    pos_dot_f = 0;
    for (k = 0;  k < natoms;  k++) {
      inv_mass = 1.0 / atom[k].m;
      vel[k].x = oldvel[k].x + half_dt * f[k].x * inv_mass;
      vel[k].y = oldvel[k].y + half_dt * f[k].y * inv_mass;
      vel[k].z = oldvel[k].z + half_dt * f[k].z * inv_mass;
      pos_dot_f += oldpos[k].x * f[k].x + oldpos[k].y * f[k].y
                   + oldpos[k].z * f[k].z;
    }
    beta[0] = oldbeta[0] + half_dt * (-pos_dot_f - 2 * data->energy[PE]);

    /* half drift */
    for (k = 0;  k < natoms;  k++) {
      pos[k].x = oldpos[k].x + half_dt * vel[k].x;
      pos[k].y = oldpos[k].y + half_dt * vel[k].y;
      pos[k].z = oldpos[k].z + half_dt * vel[k].z;
    }

    /*
     * second halfstep of leapfrog
     */
    data->arr_index = (data->arr_index + 1) & NUMARR_MINUS_1;
    oldpos = pos;  pos = data->posarr[data->arr_index];
    oldvel = vel;  vel = data->velarr[data->arr_index];
    oldbeta = beta;  beta = data->beta_arr[data->arr_index];

    /* half drift */
    for (k = 0;  k < natoms;  k++) {
      pos[k].x = oldpos[k].x + half_dt * oldvel[k].x;
      pos[k].y = oldpos[k].y + half_dt * oldvel[k].y;
      pos[k].z = oldpos[k].z + half_dt * oldvel[k].z;
    }

    /* compute the force */
    data->pos = pos;
    data->vel = oldvel;
    data->stepnum++;
    if (MD_fcallback(sim, data->stepnum, 0, 1) || MD_waitcb(sim)) {
      return MD_FAIL;
    }
    if (deven_force_compute(data->force))  return MD_FAIL;

    /* half kick */
    pos_dot_f = 0;
    for (k = 0;  k < natoms;  k++) {
      inv_mass = 1.0 / atom[k].m;
      vel[k].x = oldvel[k].x + half_dt * f[k].x * inv_mass;
      vel[k].y = oldvel[k].y + half_dt * f[k].y * inv_mass;
      vel[k].z = oldvel[k].z + half_dt * f[k].z * inv_mass;
      pos_dot_f += pos[k].x * f[k].x + pos[k].y * f[k].y + pos[k].z * f[k].z;
    }
    beta[0] = oldbeta[0] + half_dt * (-pos_dot_f - 2 * data->energy[PE]);
    data->vel = vel;

    /* find reductions for system */
    data->energy[KE] = compute_ke(data);
    data->energy[TE] = data->energy[PE] + data->energy[KE];
    if (data->is_shadow_ready) {
      compute_shadow(data);
    }
    else if (data->arr_index >= (NUMARR / 2)) {
      data->is_shadow_ready = 1;
      compute_shadow(data);
    }
    set_reduc(data);

    if (MD_callback(sim, data->stepnum) || MD_waitcb(sim))  return MD_FAIL;
  }
  return 0;
}



void set_reduc(Data *data)
{
  MD_Int i;

  for (i = PE_BOND;  i <= SHADOW8;  i++) {
    data->reduc[i] = data->energy[i] * KCAL_PER_MOL;
  }
}



double compute_ke(Data *data)
{
  const MD_Dvec *v = data->vel;
  const MD_Atom *a = data->atom;
  const MD_Int natoms = data->natoms;
  MD_Int k;
  double twice_ke = 0.0;

  for (k = 0;  k < natoms;  k++) {
    twice_ke += (v[k].x*v[k].x + v[k].y*v[k].y + v[k].z*v[k].z) * a[k].m;
  }
  return (0.5 * twice_ke);
}



void compute_shadow(Data *data)
{
  const int natoms = data->natoms;
  const MD_Atom *atom = data->atom;
  const int n = (data->arr_index - 4) & NUMARR_MINUS_1;
                                                 /* n is 2 steps back */
  const int np1 = (n + 1) & NUMARR_MINUS_1;      /* n + 1/2 */
  const int nm1 = (n - 1) & NUMARR_MINUS_1;      /* n - 1/2 */
  const int np2 = (n + 2) & NUMARR_MINUS_1;      /* n + 1 */
  const int nm2 = (n - 2) & NUMARR_MINUS_1;      /* n - 1 */
  const int np3 = (n + 3) & NUMARR_MINUS_1;      /* n + 3/2 */
  const int nm3 = (n - 3) & NUMARR_MINUS_1;      /* n - 3/2 */
  const int np4 = (n + 4) & NUMARR_MINUS_1;      /* n + 2 */
  const int nm4 = (n - 4) & NUMARR_MINUS_1;      /* n - 2 */
  int k;
  const double c = 0.5 / data->timestep;

  const MD_Dvec *pos_np4 = data->posarr[np4];
  const MD_Dvec *pos_np3 = data->posarr[np3];
  const MD_Dvec *pos_np2 = data->posarr[np2];
  const MD_Dvec *pos_np1 = data->posarr[np1];
  const MD_Dvec *pos_n   = data->posarr[n];
  const MD_Dvec *pos_nm1 = data->posarr[nm1];
  const MD_Dvec *pos_nm2 = data->posarr[nm2];
  const MD_Dvec *pos_nm3 = data->posarr[nm3];
  const MD_Dvec *pos_nm4 = data->posarr[nm4];

  const MD_Dvec *vel_np4 = data->velarr[np4];
  const MD_Dvec *vel_np3 = data->velarr[np3];
  const MD_Dvec *vel_np2 = data->velarr[np2];
  const MD_Dvec *vel_np1 = data->velarr[np1];
  const MD_Dvec *vel_n   = data->velarr[n];
  const MD_Dvec *vel_nm1 = data->velarr[nm1];
  const MD_Dvec *vel_nm2 = data->velarr[nm2];
  const MD_Dvec *vel_nm3 = data->velarr[nm3];
  const MD_Dvec *vel_nm4 = data->velarr[nm4];

  const double *beta_np4 = data->beta_arr[np4];
  const double *beta_np3 = data->beta_arr[np3];
  const double *beta_np2 = data->beta_arr[np2];
  const double *beta_np1 = data->beta_arr[np1];
  const double *beta_nm1 = data->beta_arr[nm1];
  const double *beta_nm2 = data->beta_arr[nm2];
  const double *beta_nm3 = data->beta_arr[nm3];
  const double *beta_nm4 = data->beta_arr[nm4];

  MD_Dvec q0, p0, q1, p1, q2, p2, q3, p3, q4, p4;
  double a10 = 0, a12 = 0, a14 = 0, a30 = 0, a32 = 0, a34 = 0,
         b10 = 0, b12 = 0, b30 = 0, b32 = 0;

  for (k = 0;  k < natoms;  k++) {

    /* compute full step formulaes */
    q0.x = pos_n[k].x;
    q0.y = pos_n[k].y;
    q0.z = pos_n[k].z;
    p0.x = vel_n[k].x * atom[k].m;
    p0.y = vel_n[k].y * atom[k].m;
    p0.z = vel_n[k].z * atom[k].m;
    q1.x = 0.5 * (pos_np2[k].x - pos_nm2[k].x);
    q1.y = 0.5 * (pos_np2[k].y - pos_nm2[k].y);
    q1.z = 0.5 * (pos_np2[k].z - pos_nm2[k].z);
    p1.x = 0.5 * (vel_np2[k].x - vel_nm2[k].x) * atom[k].m;
    p1.y = 0.5 * (vel_np2[k].y - vel_nm2[k].y) * atom[k].m;
    p1.z = 0.5 * (vel_np2[k].z - vel_nm2[k].z) * atom[k].m;
    q2.x = pos_np2[k].x - 2 * pos_n[k].x + pos_nm2[k].x;
    q2.y = pos_np2[k].y - 2 * pos_n[k].y + pos_nm2[k].y;
    q2.z = pos_np2[k].z - 2 * pos_n[k].z + pos_nm2[k].z;
    p2.x = (vel_np2[k].x - 2 * vel_n[k].x + vel_nm2[k].x) * atom[k].m;
    p2.y = (vel_np2[k].y - 2 * vel_n[k].y + vel_nm2[k].y) * atom[k].m;
    p2.z = (vel_np2[k].z - 2 * vel_n[k].z + vel_nm2[k].z) * atom[k].m;
    q3.x = 0.5 * pos_np4[k].x - pos_np2[k].x + pos_nm2[k].x
           - 0.5 * pos_nm4[k].x;
    q3.y = 0.5 * pos_np4[k].y - pos_np2[k].y + pos_nm2[k].y
           - 0.5 * pos_nm4[k].y;
    q3.z = 0.5 * pos_np4[k].z - pos_np2[k].z + pos_nm2[k].z
           - 0.5 * pos_nm4[k].z;
    p3.x = (0.5 * vel_np4[k].x - vel_np2[k].x + vel_nm2[k].x
           - 0.5 * vel_nm4[k].x) * atom[k].m;
    p3.y = (0.5 * vel_np4[k].y - vel_np2[k].y + vel_nm2[k].y
           - 0.5 * vel_nm4[k].y) * atom[k].m;
    p3.z = (0.5 * vel_np4[k].z - vel_np2[k].z + vel_nm2[k].z
           - 0.5 * vel_nm4[k].z) * atom[k].m;
    q4.x = pos_np4[k].x - 4 * pos_np2[k].x + 6 * pos_n[k].x
           - 4 * pos_nm2[k].x + pos_nm4[k].x;
    q4.y = pos_np4[k].y - 4 * pos_np2[k].y + 6 * pos_n[k].y
           - 4 * pos_nm2[k].y + pos_nm4[k].y;
    q4.z = pos_np4[k].z - 4 * pos_np2[k].z + 6 * pos_n[k].z
           - 4 * pos_nm2[k].z + pos_nm4[k].z;
    p4.x = (vel_np4[k].x - 4 * vel_np2[k].x + 6 * vel_n[k].x
           - 4 * vel_nm2[k].x + vel_nm4[k].x) * atom[k].m;
    p4.y = (vel_np4[k].y - 4 * vel_np2[k].y + 6 * vel_n[k].y
           - 4 * vel_nm2[k].y + vel_nm4[k].y) * atom[k].m;
    p4.z = (vel_np4[k].z - 4 * vel_np2[k].z + 6 * vel_n[k].z
           - 4 * vel_nm2[k].z + vel_nm4[k].z) * atom[k].m;
    a10 += (q1.x * p0.x + q1.y * p0.y + q1.z * p0.z)
           - (p1.x * q0.x + p1.y * q0.y + p1.z * q0.z);
    a12 += (q1.x * p2.x + q1.y * p2.y + q1.z * p2.z)
           - (p1.x * q2.x + p1.y * q2.y + p1.z * q2.z);
    a14 += (q1.x * p4.x + q1.y * p4.y + q1.z * p4.z)
           - (p1.x * q4.x + p1.y * q4.y + p1.z * q4.z);
    a30 += (q3.x * p0.x + q3.y * p0.y + q3.z * p0.z)
           - (p3.x * q0.x + p3.y * q0.y + p3.z * q0.z);
    a32 += (q3.x * p2.x + q3.y * p2.y + q3.z * p2.z)
           - (p3.x * q2.x + p3.y * q2.y + p3.z * q2.z);
    a34 += (q3.x * p4.x + q3.y * p4.y + q3.z * p4.z)
           - (p3.x * q4.x + p3.y * q4.y + p3.z * q4.z);

    /* compute half step formulaes */
    q0.x = 0.5 * (pos_np1[k].x + pos_nm1[k].x);
    q0.y = 0.5 * (pos_np1[k].y + pos_nm1[k].y);
    q0.z = 0.5 * (pos_np1[k].z + pos_nm1[k].z);
    p0.x = 0.5 * (vel_np1[k].x + vel_nm1[k].x) * atom[k].m;
    p0.y = 0.5 * (vel_np1[k].y + vel_nm1[k].y) * atom[k].m;
    p0.z = 0.5 * (vel_np1[k].z + vel_nm1[k].z) * atom[k].m;
    q1.x = pos_np1[k].x - pos_nm1[k].x;
    q1.y = pos_np1[k].y - pos_nm1[k].y;
    q1.z = pos_np1[k].z - pos_nm1[k].z;
    p1.x = (vel_np1[k].x - vel_nm1[k].x) * atom[k].m;
    p1.y = (vel_np1[k].y - vel_nm1[k].y) * atom[k].m;
    p1.z = (vel_np1[k].z - vel_nm1[k].z) * atom[k].m;
    q2.x = 0.5 * (pos_np3[k].x - pos_np1[k].x - pos_nm1[k].x
           + pos_nm3[k].x);
    q2.y = 0.5 * (pos_np3[k].y - pos_np1[k].y - pos_nm1[k].y
           + pos_nm3[k].y);
    q2.z = 0.5 * (pos_np3[k].z - pos_np1[k].z - pos_nm1[k].z
           + pos_nm3[k].z);
    p2.x = 0.5 * (vel_np3[k].x - vel_np1[k].x - vel_nm1[k].x
           + vel_nm3[k].x) * atom[k].m;
    p2.y = 0.5 * (vel_np3[k].y - vel_np1[k].y - vel_nm1[k].y
           + vel_nm3[k].y) * atom[k].m;
    p2.z = 0.5 * (vel_np3[k].z - vel_np1[k].z - vel_nm1[k].z
           + vel_nm3[k].z) * atom[k].m;
    q3.x = pos_np3[k].x - 3 * pos_np1[k].x + 3 * pos_nm1[k].x
           - pos_nm3[k].x;
    q3.y = pos_np3[k].y - 3 * pos_np1[k].y + 3 * pos_nm1[k].y
           - pos_nm3[k].y;
    q3.z = pos_np3[k].z - 3 * pos_np1[k].z + 3 * pos_nm1[k].z
           - pos_nm3[k].z;
    p3.x = (vel_np3[k].x - 3 * vel_np1[k].x + 3 * vel_nm1[k].x
           - vel_nm3[k].x) * atom[k].m;
    p3.y = (vel_np3[k].y - 3 * vel_np1[k].y + 3 * vel_nm1[k].y
           - vel_nm3[k].y) * atom[k].m;
    p3.z = (vel_np3[k].z - 3 * vel_np1[k].z + 3 * vel_nm1[k].z
           - vel_nm3[k].z) * atom[k].m;
    b10 += (q1.x * p0.x + q1.y * p0.y + q1.z * p0.z)
           - (p1.x * q0.x + p1.y * q0.y + p1.z * q0.z);
    b12 += (q1.x * p2.x + q1.y * p2.y + q1.z * p2.z)
           - (p1.x * q2.x + p1.y * q2.y + p1.z * q2.z);
    b30 += (q3.x * p0.x + q3.y * p0.y + q3.z * p0.z)
           - (p3.x * q0.x + p3.y * q0.y + p3.z * q0.z);
    b32 += (q3.x * p2.x + q3.y * p2.y + q3.z * p2.z)
           - (p3.x * q2.x + p3.y * q2.y + p3.z * q2.z);
  }
  a10 -= 0.5 * (beta_np2[0] - beta_nm2[0]);
  a30 -= 0.5 * beta_np4[0] - beta_np2[0] + beta_nm2[0] - 0.5 * beta_nm4[0];
  b10 -= beta_np1[0] - beta_nm1[0];
  b30 -= beta_np3[0] - 3 * beta_np1[0] + 3 * beta_nm1[0] - beta_nm3[0];
  data->energy[SHADOW2] = c * b10;
  data->energy[SHADOW4] = c * (a10 - (1.0/6) * a12);
  data->energy[SHADOW6] = c * (b10 - (7.0/20) * b12 + (11.0/60) * b30
                        + (1.0/30) * b32);
  data->energy[SHADOW8] = c * (a10 - (2.0/7) * a12 - (19.0/210) * a14
                        + (5.0/42) * a30 + (13.0/105) * a32 - (1.0/140) * a34);
}