#include <stdlib.h>
#include "simfront.h"


/*
 * read and set initial data for position and velocity coordinate trajectories
 */


/* prototypes for internal routines */
static int read_pdb_trajectory(SimFront *, const char *, MD_Int, double);
static int read_binary_trajectory(SimFront *, const char *, MD_Int);


int setup_trajectory(SimFront *sf)
{
  /* initialize positions */
  if (sf->bincoorfile != NULL) {
    if (read_binary_trajectory(sf, sf->bincoorfile, sf->id_pos)) {
      return -1;
    }
  }
  else if (sf->coorfile != NULL) {
    if (read_pdb_trajectory(sf, sf->coorfile, sf->id_pos, 1.0)) {
      return -1;
    }
  }
  else {
    ERRMSG(sf, "initial positions must be provided");
    return -1;
  }

  /* initialize velocities */
  if (sf->binvelfile != NULL) {
    if (read_binary_trajectory(sf, sf->binvelfile, sf->id_vel)) {
      return -1;
    }
  }
  else if (sf->velfile != NULL) {
    if (read_pdb_trajectory(sf, sf->velfile, sf->id_vel, INV_PDBVEL)) {
      return -1;
    }
  }
  else if (sf->temperature >= 0) {
    ERRMSG(sf, "setting random velocities based on initial temperature "
           "is unsupported");
    return -1;
  }
  else {
    ERRMSG(sf, "initial velocities must be provided");
    return -1;
  }

  return 0;
}


int read_pdb_trajectory(SimFront *sf, const char *fname,
                        MD_Int idnum, double conv_factor)
{
  Pdb *pdb = &(sf->pdb);
  MD_Sim *sim = sf->sim;
  MD_Dvec *dvecbuf = sf->dvecbuf;
  int natoms = sf->tbuf.atom_max;
  int i, k, n;

  /* set length for number of atoms expected */
  if (MD_getlen(sim, idnum) != natoms
      && (MD_setlen(sim, idnum, natoms) || MD_wait(sim))) {
    ERRMSG(sf, MD_errmsg(sim));
    return -1;
  }

  /* open PDB trajectory file */
  if (pdb_open(pdb, fname, PDBTYPE_STANDARD)) {
    ERRMSG(sf, "call to pdb_open failed");
    return -1;
  }

  /* read in PDB trajectory file */
  /*
   * Here is the strangeness in the MDIO Pdb interface:
   *
   *   Pass buffer space to be filled directly to pdb_read().
   *   Only PDB ATOM records are parsed from file.  Buffer space is
   *   an array of vector type followed by two arrays of floating
   *   point type.  Precision of type (whether single or double) is
   *   also indicated.  Pass NULL for any array type whose column is
   *   not needed.  Number of array spaces filled is returned as a
   *   parameter.  Return value indicates success or fail.
   */
  for (k = 0;  k < natoms;  k += n) {
    /* read from PDB file into vector buffer */
    if (pdb_read(pdb, dvecbuf, NULL, NULL, DOUBLETYPE, ARRAYLEN, &n)) {
      ERRMSG(sf, "call to pdb_read failed");
      return -1;
    }
    else if (n + k > natoms) {
      ERRMSG(sf, "more atoms in PDB file than expected");
      return -1;
    }

    /* convert to our internal units */
    for (i = 0;  i < n;  i++) {
      dvecbuf[i].x *= conv_factor;
      dvecbuf[i].y *= conv_factor;
      dvecbuf[i].z *= conv_factor;
    }

    /* set engine data */
    if (MD_setdata(sim, idnum, dvecbuf, n, k) || MD_wait(sim)) {
      ERRMSG(sf, MD_errmsg(sim));
      return -1;
    }
  }

  /* close PDB trajectory file */
  if (pdb_close(pdb)) {
    ERRMSG(sf, "call to pdb_close failed");
    return -1;
  }

  return 0;
}


int read_binary_trajectory(SimFront *sf, const char *fname, MD_Int idnum)
{
  Binres *binres = &(sf->binres);
  MD_Sim *sim = sf->sim;
  MD_Dvec *dvecbuf = sf->dvecbuf;
  int natoms = sf->tbuf.atom_max;
  int k, n, max;

  /* set length for number of atoms expected */
  if (MD_getlen(sim, idnum) != natoms
      && (MD_setlen(sim, idnum, natoms) || MD_wait(sim))) {
    ERRMSG(sf, MD_errmsg(sim));
    return -1;
  }

  /* open binary trajectory file */
  if (binres_open(binres, fname, BINRES_READ, MD_DVEC)) {
    ERRMSG(sf, "call to binres_open failed");
    return -1;
  }

  /* read in binary trajectory file */
  /*
   * Here is the strangeness in the MDIO Binres interface:
   *
   *   Pass buffer space to be filled directly to binres_read().
   *   Unlike reading for PDB files, return value is nonzero if
   *   more to be read or zero if read is finished.  Returns through
   *   argument list the number of elements read and the number of
   *   atoms actually contained in file.
   */
  k = 0;
  while (binres_read(binres, dvecbuf, ARRAYLEN, &n, &max)) {
    if (max != natoms) {
      ERRMSG(sf, "unexpected number of atoms in binary restart file");
      return -1;
    }
    /* set engine data */
    if (MD_setdata(sim, idnum, dvecbuf, n, k) || MD_wait(sim)) {
      ERRMSG(sf, MD_errmsg(sim));
      return -1;
    }
    k += n;
  }
  if (max != natoms) {
    ERRMSG(sf, "unexpected number of atoms in binary restart file");
    return -1;
  }
  /* finish setting engine data */
  if (MD_setdata(sim, idnum, dvecbuf, n, k) || MD_wait(sim)) {
    ERRMSG(sf, MD_errmsg(sim));
    return -1;
  }

  /* close binary trajectory file */
  if (binres_close(binres)) {
    ERRMSG(sf, "call to binres_close failed");
    return -1;
  }

  return 0;
}