/*
 * Copyright (C) 2004-2005 by David J. Hardy.  All rights reserved.
 *
 * update.c
 */

#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <math.h>
#include "mdapi/mdengine.h"
#include "force/force.h"
#include "mgrid/mgrid.h"
#include "deven/engine.h"
#include "debug/debug.h"


enum { FALSE = 0, TRUE = 1 };


int32 deven_update(MD_Front *front)
{
  Engine *eng = MD_engine_data(front);
  Param *param = &(eng->param);
  Force *force = &(eng->force);
  ForceParam *p = &(eng->f_param);
  ForceResult *r = &(eng->f_result);
  ForceEnergy *e = &(eng->result.fe);
  Mgrid *mg = &(eng->mgrid);
  MgridParam *mgprm = &(eng->mgrid_param);
  MgridSystem *mgsys = &(eng->mgrid_system);
  Step *step = &(eng->step);
  StepParam *step_param = &(eng->step_param);
  StepSystem *step_system = &(eng->step_system);
  int32 do_remove_com_motion = FALSE;

  /* make sure number of atoms is consistent */
  eng->natoms = eng->atom->attrib.len;
  if (eng->natoms <= 0
      || eng->natoms != eng->pos->attrib.len
      || eng->natoms != eng->vel->attrib.len) {
    return MD_error(front, eng->err_param);
  }

  /* set flag if mgrid is "on" */
  eng->ismgrid = FALSE;
  if (strcasecmp(param->mgrid, "on") == 0
      || strcasecmp(param->mgrid, "yes") == 0) {
    eng->ismgrid = TRUE;
  }
  else if (strcasecmp(param->mgrid, "off") != 0
      && strcasecmp(param->mgrid, "no") != 0) {
    return MD_error(front, eng->err_force_init);
  }

  /* make sure force array is correct length */
  /*** lock several arrays to this length, requires resizing all of them ***/
  if (eng->natoms != eng->force_engdata->attrib.len) {

    /* resize force array to correct length */
    if (MD_engdata_setlen(front, eng->force_engdata, eng->natoms)) {
      return MD_FAIL;
    }

    /* make sure force result uses force array as its buffer space */
    memset(r, 0, sizeof(ForceResult));
    r->f = (MD_Dvec *) (eng->force_engdata->buf);

    /* make sure step system uses this force array */
    step_system->force = r->f;
    step_system->is_force_valid = STEP_FALSE;

#if 0
  /*** for debugging electrostatics ***/
    /* resize force elec array to correct length */
    if (MD_engdata_setlen(front, eng->force_elec, eng->natoms)) {
      return MD_FAIL;
    }
    r->f_elec = (MD_Dvec *) (eng->force_elec->buf);
  /*** end for debugging electrostatics ***/
#endif

    if (eng->ismgrid) {
      /* resize mgrid system arrays to correct length */
      if (MD_engdata_setlen(front, eng->wrap, eng->natoms)
	  || MD_engdata_setlen(front, eng->poswrap, eng->natoms)
	  || MD_engdata_setlen(front, eng->f_elec, eng->natoms)
          || MD_engdata_setlen(front, eng->charge, eng->natoms)) {
        return MD_FAIL;
      }

      /* make sure mgrid system uses these arrays */
      memset(mgsys, 0, sizeof(MgridSystem));
      mgsys->f_elec = (MD_Dvec *) (eng->f_elec->buf);
      mgsys->pos = (MD_Dvec *) (eng->poswrap->buf);
      mgsys->charge = (double *) (eng->charge->buf);
    }
  }

  /* deal with velocity update */
  if (eng->vel->attrib.access & MD_MODIFY) {

    /* make sure step system uses this velocity array */
    step_system->vel = (MD_Dvec *) (eng->vel->buf);

    /* acknowledge velocity modification to interface */
    if (MD_engdata_ackmod(front, eng->vel)) {
      return MD_FAIL;
    }
  }

  /* deal with position update */
  if (eng->pos->attrib.access & MD_MODIFY) {

    /* make sure step system uses this position array */
    step_system->pos = (MD_Dvec *) (eng->pos->buf);
    step_system->is_force_valid = STEP_FALSE;

    /* acknowledge position modification to interface */
    if (MD_engdata_ackmod(front, eng->pos)) {
      return MD_FAIL;
    }
  }

  /* see if we need to (re)initialize step library */
  if (eng->atom->attrib.access & MD_MODIFY) {

    /* determine number of degrees of freedom for system */
    eng->ndegfreedom = 3 * eng->natoms;

    /* allow center-of-mass motion? */
    do_remove_com_motion = FALSE;
    if (strcasecmp(eng->com_motion, "no") == 0
        || strcasecmp(eng->com_motion, "off") == 0) {

      /* remove center of mass motion from initial velocities */
      do_remove_com_motion = TRUE;

      /* this reduces number of degrees of freedom (by one particle) */
      eng->ndegfreedom -= 3;
    }
    else if (strcasecmp(eng->com_motion, "re") == 0
        || strcasecmp(eng->com_motion, "res") == 0) {
      /* restarting from previous simulation */
      /* don't remove center of mass motion because that changes vel */
      /* but do adjust number of degrees of freedom (by one particle) */
      eng->ndegfreedom -= 3;
    }
    else if (strcasecmp(eng->com_motion, "yes") != 0
        && strcasecmp(eng->com_motion, "on") != 0) {
      return MD_error(front, eng->err_force_init);
    }

    /* set step params */
    step_param->random_seed = eng->seed;
    step_param->timestep = eng->timestep;
    step_param->force_object = front;
    step_param->force_compute = deven_force;
    step_param->atom = (MD_Atom *) (eng->atom->buf);
    step_param->natoms = eng->natoms;
    step_param->ndegfreedom = eng->ndegfreedom;

    if (eng->is_step_initialized) {
      printf("# WARNING: reinitializing step library\n");
      /* since these are NOT initial velocities, don't remove com motion */
      do_remove_com_motion = FALSE;
    }
    if (step_init(step, step_param)) {
      return MD_error(front, eng->err_force_init);
    }
    eng->is_step_initialized = TRUE;

    /*** MD_MODIFY flag is cleared from atom later ***/

  } /* done with step library initialization */

  /* see if we need to (re)initialize force library */
  if ((eng->atomprm->attrib.access & MD_MODIFY)
      || (eng->bondprm->attrib.access & MD_MODIFY)
      || (eng->angleprm->attrib.access & MD_MODIFY)
      || (eng->dihedprm->attrib.access & MD_MODIFY)
      || (eng->imprprm->attrib.access & MD_MODIFY)
      || (eng->nbfixprm->attrib.access & MD_MODIFY)
      || (eng->atom->attrib.access & MD_MODIFY)
      || (eng->bond->attrib.access & MD_MODIFY)
      || (eng->angle->attrib.access & MD_MODIFY)
      || (eng->dihed->attrib.access & MD_MODIFY)
      || (eng->impr->attrib.access & MD_MODIFY)
      || (eng->excl->attrib.access & MD_MODIFY)
      || (eng->param_engdata->attrib.access & MD_MODIFY)) {

    /* must first destroy previous force */
    force_done(force);
    if (force_init(force)) {
      return MD_error(front, eng->err_force_init);
    }

    /* zero force param */
    memset(p, 0, sizeof(ForceParam));

    if (eng->ismgrid) {
      /* must first destroy previous mgrid */
      mgrid_done(mg);
      if (mgrid_init(mg)) {
        return MD_error(front, eng->err_force_init);
      }
      /* zero mgrid param */
      memset(mgprm, 0, sizeof(MgridParam));
    }

    /* set flags */
    p->flags = FORCE_ALL;

    /* setup force params */
    p->atomprm = (MD_AtomPrm *) (eng->atomprm->buf);
    p->bondprm = (MD_BondPrm *) (eng->bondprm->buf);
    p->angleprm = (MD_AnglePrm *) (eng->angleprm->buf);
    p->dihedprm = (MD_TorsPrm *) (eng->dihedprm->buf);
    p->imprprm = (MD_TorsPrm *) (eng->imprprm->buf);
    p->nbfixprm = (MD_NbfixPrm *) (eng->nbfixprm->buf);
    p->atomprm_len = eng->atomprm->attrib.len;
    p->bondprm_len = eng->bondprm->attrib.len;
    p->angleprm_len = eng->angleprm->attrib.len;
    p->dihedprm_len = eng->dihedprm->attrib.len;
    p->imprprm_len = eng->imprprm->attrib.len;
    p->nbfixprm_len = eng->nbfixprm->attrib.len;

    p->atom = (MD_Atom *) (eng->atom->buf);
    p->bond = (MD_Bond *) (eng->bond->buf);
    p->angle = (MD_Angle *) (eng->angle->buf);
    p->dihed = (MD_Tors *) (eng->dihed->buf);
    p->impr = (MD_Tors *) (eng->impr->buf);
    p->excl = (MD_Excl *) (eng->excl->buf);
    p->atom_len = eng->atom->attrib.len;
    p->bond_len = eng->bond->attrib.len;
    p->angle_len = eng->angle->attrib.len;
    p->dihed_len = eng->dihed->attrib.len;
    p->impr_len = eng->impr->attrib.len;
    p->excl_len = eng->excl->attrib.len;

    /* can't handle non-orthogonal cells */
    if (param->cellBasisVector1.y != 0.0
        || param->cellBasisVector1.z != 0.0
        || param->cellBasisVector2.x != 0.0
        || param->cellBasisVector2.z != 0.0
        || param->cellBasisVector3.x != 0.0
        || param->cellBasisVector3.y != 0.0) {
      return MD_error(front, eng->err_force_init);
    }
    if (param->cellBasisVector1.x != 0.0) {
      p->flags |= FORCE_X_PERIODIC;
      p->xlen = param->cellBasisVector1.x;
      p->center = param->cellOrigin;
    }
    if (param->cellBasisVector2.y != 0.0) {
      p->flags |= FORCE_Y_PERIODIC;
      p->ylen = param->cellBasisVector2.y;
      p->center = param->cellOrigin;
    }
    if (param->cellBasisVector3.z != 0.0) {
      p->flags |= FORCE_Z_PERIODIC;
      p->zlen = param->cellBasisVector3.z;
      p->center = param->cellOrigin;
    }

    /* setup spherical boundary conditions */
    if (strcasecmp(param->sphericalBC, "on") == 0
        || strcasecmp(param->sphericalBC, "yes") == 0) {
      /* turn periodicity off */
      p->flags &= ~FORCE_PERIODIC;
      /* turn spherical boundary conditions on */
      p->flags |= FORCE_SPHERE;
      /*
       * in this case, xlen, ylen, zlen are still used by grid cell
       * algorithm to determine spatial region of molecule
       * (make sure they correspond to extent of sphere)
       */
      p->center = param->sphericalBCCenter;
      p->radius1 = param->sphericalBCr1;
      p->radius2 = param->sphericalBCr2;
      p->konst1 = param->sphericalBCk1;
      p->konst2 = param->sphericalBCk2;
      p->exp1 = param->sphericalBCexp1;
      p->exp2 = param->sphericalBCexp2;
    }
    else if (strcasecmp(param->sphericalBC, "off") != 0
        && strcasecmp(param->sphericalBC, "no") != 0) {
      return MD_error(front, eng->err_force_init);
    }

    /* setup cylindrical boundary conditions */
    if (strcasecmp(param->cylindricalBC, "on") == 0
        || strcasecmp(param->cylindricalBC, "yes") == 0) {
      /* make sure spherical boundaries were not also set */
      if (p->flags & FORCE_SPHERE) {
        return MD_error(front, eng->err_force_init);
      }
      /* turn periodicity off */
      p->flags &= ~FORCE_PERIODIC;
      /* flags depend on axis value */
      if (strcasecmp(param->cylindricalBCAxis, "x") == 0) {
        p->flags |= FORCE_X_CYLINDER;
      }
      else if (strcasecmp(param->cylindricalBCAxis, "y") == 0) {
        p->flags |= FORCE_Y_CYLINDER;
      }
      else if (strcasecmp(param->cylindricalBCAxis, "z") == 0) {
        p->flags |= FORCE_Z_CYLINDER;
      }
      else {
        return MD_error(front, eng->err_force_init);
      }
      /*
       * in this case, xlen, ylen, zlen are still used by grid cell
       * algorithm to determine spatial region of molecule
       * (make sure they correspond to extent of cylinder)
       */
      p->center = param->cylindricalBCCenter;
      p->radius1 = param->cylindricalBCr1;
      p->radius2 = param->cylindricalBCr2;
      p->length1 = param->cylindricalBCl1;
      p->length2 = param->cylindricalBCl2;
      p->konst1 = param->cylindricalBCk1;
      p->konst2 = param->cylindricalBCk2;
      p->exp1 = param->cylindricalBCexp1;
      p->exp2 = param->cylindricalBCexp2;
    }
    else if (strcasecmp(param->cylindricalBC, "off") != 0
        && strcasecmp(param->cylindricalBC, "no") != 0) {
      return MD_error(front, eng->err_force_init);
    }

    /* full direct electrostatics */
    if (strcasecmp(param->fullDirect, "on") == 0
        || strcasecmp(param->fullDirect, "yes") == 0) {
      p->flags |= FORCE_ELEC_DIRECT;
      /* overrides mgrid */
      eng->ismgrid = FALSE;
    }
    else if (strcasecmp(param->fullDirect, "off") != 0
        && strcasecmp(param->fullDirect, "no") != 0) {
      return MD_error(front, eng->err_force_init);
    }

    /* set exclusion policy */
    if (strcasecmp(param->exclude, "none") == 0) {
      p->flags |= FORCE_EXCL_NONE;
    }
    else if (strcasecmp(param->exclude, "1-2") == 0) {
      p->flags |= FORCE_EXCL_12;
    }
    else if (strcasecmp(param->exclude, "1-3") == 0) {
      p->flags |= FORCE_EXCL_13;
    }
    else if (strcasecmp(param->exclude, "1-4") == 0) {
      p->flags |= FORCE_EXCL_14;
    }
    else if (strcasecmp(param->exclude, "scaled1-4") == 0) {
      p->flags |= FORCE_EXCL_SCAL14;
    }
    else {
      return MD_error(front, eng->err_force_init);
    }

    /* switching and smoothing */
    if (strcasecmp(param->switching, "on") == 0
        || strcasecmp(param->switching, "yes") == 0) {
      if (p->flags & FORCE_ELEC_DIRECT) {
        p->flags |= FORCE_SWITCH;
      }
      else {
        p->flags |= FORCE_CONTINUOUS;  /* i.e. SWITCH + SMOOTH */
      }
      p->switchdist = param->switchDist;
    }
    else if (strcasecmp(param->switching, "off") != 0
        && strcasecmp(param->switching, "no") != 0) {
      return MD_error(front, eng->err_force_init);
    }

    /* nonbonded params */
    p->cutoff = param->cutoff;
    p->elec_const = MD_COULOMB;
    p->dielectric = param->dielectric;
    p->scaling14 = param->scaling;  /* value of "1-4scaling" */

    /* see if mgrid is to be setup (override by full direct elec on) */
    if (eng->ismgrid) {
      double c;
      int32 n;

      /* does use of boundary conditions permit mgrid for cubic domain? */
      if ((p->flags & FORCE_MASK_PERIOD) == FORCE_PERIODIC
          && p->xlen == p->ylen && p->ylen == p->zlen
          && p->zlen == param->mgridLength) {
        mgprm->center = p->center;
        mgprm->boundary = MGRID_PERIODIC;
      }
      else if ((p->flags & FORCE_MASK_PERIOD) == FORCE_NONPERIODIC
          && (p->flags & FORCE_MASK_BC) == FORCE_SPHERE) {
        mgprm->center = p->center;
        mgprm->boundary = MGRID_NONPERIODIC;
      }
      else {
        /* choice of boundary conditions does not work */
        return MD_error(front, eng->err_force_init);
      }

      /* set remaining params */
      mgprm->length = param->mgridLength;
      mgprm->cutoff = param->mgridCutoff;
      mgprm->spacing = param->mgridSpacing;
      mgprm->nspacings = param->mgridNspacings;
      mgprm->nlevels = param->mgridNlevels;
      mgprm->natoms = eng->natoms;
      mgprm->approx = mgrid_string_to_approx(param->mgridApprox);
      mgprm->split = mgrid_string_to_split(param->mgridSplit);

      /* setup charge array, scale by sqrt of multiplicative const */
      c = sqrt(p->elec_const / p->dielectric);
      for (n = 0;  n < eng->natoms;  n++) {
        mgsys->charge[n] = c * p->atom[n].q;
      }

      if (mgrid_param_config(mgprm)) {
        printf("# mgrid failed to configure parameters\n");
        return MD_error(front, eng->err_force_init);
      }

      if (mgrid_setup(mg, mgsys, mgprm)) {
        printf("# mgrid setup failed\n");
        return MD_error(front, eng->err_force_init);
      }
      else {
        printf("# mgrid setup successful with parameters:\n");
        printf("#   center = %g %g %g\n", mgprm->center.x,
            mgprm->center.y, mgprm->center.z);
        printf("#   length = %g\n", mgprm->length);
        printf("#   cutoff = %g\n", mgprm->cutoff);
        printf("#   spacing = %g\n", mgprm->spacing);
        printf("#   nspacings = %d\n", mgprm->nspacings);
        printf("#   nlevels = %d\n", mgprm->nlevels);
        printf("#   boundary = %s\n",
            mgrid_boundary_to_string(mgprm->boundary));
        printf("#   natoms = %d\n", mgprm->natoms);
        printf("#   approx = %s\n", mgrid_approx_to_string(mgprm->approx));
        printf("#   approx = %s\n", mgrid_split_to_string(mgprm->split));
      }

#if 0
      /* turn off all force library elec */
      p->flags &= ~(FORCE_SMOOTH | FORCE_ELEC | FORCE_ELEC_EXCL);
#endif
      p->flags &= ~FORCE_SMOOTH;
      p->flags |= FORCE_ELEC_EXCL;
    }

    /* setup force library */
    if (force_setup(force, p, e, r)) {
      printf("# force setup failed\n");
      return MD_error(front, eng->err_force_init);
    }

#if 0
    if (eng->ismgrid) {
      /* have mgrid process exclusions */
      mgsys->excl_list = force->excl_list;
      mgsys->scaled14_list = force->scaled14_list;
      mgsys->scaling14 = p->scaling14;
    }
#endif

    /* acknowledge engdata modification to interface */
    if (((eng->atomprm->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->atomprm))
        || ((eng->bondprm->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->bondprm))
        || ((eng->angleprm->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->angleprm))
        || ((eng->dihedprm->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->dihedprm))
        || ((eng->imprprm->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->imprprm))
        || ((eng->nbfixprm->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->nbfixprm))
        || ((eng->atom->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->atom))
        || ((eng->bond->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->bond))
        || ((eng->angle->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->angle))
        || ((eng->dihed->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->dihed))
        || ((eng->impr->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->impr))
        || ((eng->excl->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->excl))
        || ((eng->param_engdata->attrib.access & MD_MODIFY)
          && MD_engdata_ackmod(front, eng->param_engdata))) {
      printf("# need to ack!!!\n");
      return MD_error(front, eng->err_force_init);
    }
  } /* done with force library initialization */

  /* deal with temperature update */
  if (eng->init_temp_engdata->attrib.access & MD_MODIFY) {

    /* initialize velocities using this initial temperature */
    if (step_set_random_velocities(step, step_system, eng->init_temp)) {
      return MD_error(front, eng->err_force_init);
    }

    /* acknowledge temperature modification to interface */
    if (MD_engdata_ackmod(front, eng->init_temp_engdata)) {
      return MD_FAIL;
    }
  }

  /* if needed, remove center of mass motion from initial velocities */
  if (do_remove_com_motion) {
    if (step_remove_com_motion(step, step_system)) {
      return MD_error(front, eng->err_force_init);
    }
  }

  return 0;
}