/*
 * topoxp.c
 *
 * Summary:  Reader for XPLOR topology files (PSF).
 *
 * Author:  David Hardy
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "topo.h"
#include "topoxp.h"
#include "pdata.h"
#include "error.h"
#include "mdutil.h"


#define  ATOMNAME_BUFSIZE  5
#define  ATOMNAME_FORMAT   "%4s"  /* format string for reading atom type */
#define  WILDCARD          "X"    /* used to match any atom name */


static int read_atom(Topoxp *t)
{
  Tbuf *tbuf = t->tbuf;
  FILE *f = t->f;
  Pdata *pdata = t->pdata;
  char *atomstring = t->atomstring;
  int *atomid = t->atomid;
  MD_Atom *atom = tbuf->atom;
  int n, k, num, index;
  float mass, charge;  /* mass and charge */
  char atomtype[ATOMNAME_BUFSIZE];
  char atomname[ATOMNAME_BUFSIZE];

  /* determine number of atoms to be read on this call */
  n = tbuf->atom_max - t->natom;
  if (tbuf->atom_len < n)  n = tbuf->atom_len;

  /* read n atoms and set topology info */
  /* also need to update atom param mass field */
  for (k = 0;  k < n;  k++, t->natom++) {
    fscanf(f, "%d %*s %*s %*s " ATOMNAME_FORMAT " " ATOMNAME_FORMAT \
              " %f %f %*f", &num, atomname, atomtype, &charge, &mass);
    MDIO_ASSERT(num - 1 == t->natom);
    index = pdata_find_atom(pdata, atomtype);
    MDIO_ASSERT(index >= 0 && index < pdata->pbuf->atomp_num);
    if (atomstring[index * ATOMNAME_BUFSIZE] == '\0') {
      strcpy(&atomstring[index * ATOMNAME_BUFSIZE], atomtype);
    }
    atom[k].m = mass;
    atom[k].q = charge;
    atom[k].param = index;
    strcpy(atom[k].name, atomname);
    atomid[t->natom] = index;
  }
  tbuf->atom_num = n;
  return MD_ATOM;
}


static int read_bond(Topoxp *t)
{
  Tbuf *tbuf = t->tbuf;
  FILE *f = t->f;
  Pdata *pdata = t->pdata;
  char *atomstring = t->atomstring;
  int *atomid = t->atomid;
  MD_Bond *bond = tbuf->bond;
  int n, k, index, atom1, atom2;
  char *name1, *name2;

  /* determine number of bonds to be read on this call */
  n = tbuf->bond_max - t->nbond;
  if (tbuf->bond_len < n)  n = tbuf->bond_len;

  /* read n bonds and set topology info */
  for (k = 0;  k < n;  k++, t->nbond++) {
    fscanf(f, "%d %d", &atom1, &atom2);
    MDIO_ASSERT(atom1 > 0 && atom1 <= t->natom
             && atom2 > 0 && atom2 <= t->natom);
    name1 = &atomstring[ atomid[atom1 - 1] * ATOMNAME_BUFSIZE ];
    name2 = &atomstring[ atomid[atom2 - 1] * ATOMNAME_BUFSIZE ];
    index = pdata_find_bond(pdata, name1, name2);
    MDIO_ASSERT(index >= 0 && index < pdata->pbuf->bondp_num);
    bond[k].atom[0] = atom1 - 1;
    bond[k].atom[1] = atom2 - 1;
    bond[k].param = index;
  }
  tbuf->bond_num = n;
  return MD_BOND;
}


static int read_angle(Topoxp *t)
{
  Tbuf *tbuf = t->tbuf;
  FILE *f = t->f;
  Pdata *pdata = t->pdata;
  char *atomstring = t->atomstring;
  int *atomid = t->atomid;
  MD_Angle *angle = tbuf->angle;
  int n, k, index, atom1, atom2, atom3;
  char *name1, *name2, *name3;

  /* determine number of angles to be read on this call */
  n = tbuf->angle_max - t->nangle;
  if (tbuf->angle_len < n)  n = tbuf->angle_len;

  /* read n angles and set topology info */
  for (k = 0;  k < n;  k++, t->nangle++) {
    fscanf(f, "%d %d %d", &atom1, &atom2, &atom3);
    MDIO_ASSERT(atom1 > 0 && atom1 <= t->natom
             && atom2 > 0 && atom2 <= t->natom
             && atom3 > 0 && atom3 <= t->natom);
    name1 = &atomstring[ atomid[atom1 - 1] * ATOMNAME_BUFSIZE ];
    name2 = &atomstring[ atomid[atom2 - 1] * ATOMNAME_BUFSIZE ];
    name3 = &atomstring[ atomid[atom3 - 1] * ATOMNAME_BUFSIZE ];
    index = pdata_find_angle(pdata, name1, name2, name3);
    MDIO_ASSERT(index >= 0 && index < pdata->pbuf->anglep_num);
    angle[k].atom[0] = atom1 - 1;
    angle[k].atom[1] = atom2 - 1;
    angle[k].atom[2] = atom3 - 1;
    angle[k].param = index;
  }
  tbuf->angle_num = n;
  return MD_ANGLE;
}


static int read_dihed(Topoxp *t)
{
  Tbuf *tbuf = t->tbuf;
  FILE *f = t->f;
  Pdata *pdata = t->pdata;
  char *atomstring = t->atomstring;
  int *atomid = t->atomid;
  MD_Dihed *dihed = tbuf->dihed;
  int n, k, index, atom1, atom2, atom3, atom4;
  char *name1, *name2, *name3, *name4;

  /* determine number of impropers to be read on this call */
  n = tbuf->dihed_max - t->ndihed;
  if (tbuf->dihed_len < n)  n = tbuf->dihed_len;

  /* read n dihedrals and set topology info */
  for (k = 0;  k < n;  k++, t->ndihed++) {
    fscanf(f, "%d %d %d %d", &atom1, &atom2, &atom3, &atom4);
    MDIO_ASSERT(atom1 > 0 && atom1 <= t->natom
             && atom2 > 0 && atom2 <= t->natom
             && atom3 > 0 && atom3 <= t->natom
             && atom4 > 0 && atom4 <= t->natom);
    name1 = &atomstring[ atomid[atom1 - 1] * ATOMNAME_BUFSIZE ];
    name2 = &atomstring[ atomid[atom2 - 1] * ATOMNAME_BUFSIZE ];
    name3 = &atomstring[ atomid[atom3 - 1] * ATOMNAME_BUFSIZE ];
    name4 = &atomstring[ atomid[atom4 - 1] * ATOMNAME_BUFSIZE ];
    index = pdata_find_dihed(pdata, name1, name2, name3, name4);
    if (index < 0) {
      index = pdata_find_dihed(pdata, WILDCARD, name2, name3, WILDCARD);
    }
    MDIO_ASSERT(index >= 0 && index < pdata->pbuf->dihedp_num);
    dihed[k].atom[0] = atom1 - 1;
    dihed[k].atom[1] = atom2 - 1;
    dihed[k].atom[2] = atom3 - 1;
    dihed[k].atom[3] = atom4 - 1;
    dihed[k].param = index;
  }
  tbuf->dihed_num = n;
  return MD_DIHED;
}


static int read_impr(Topoxp *t)
{
  Tbuf *tbuf = t->tbuf;
  FILE *f = t->f;
  Pdata *pdata = t->pdata;
  char *atomstring = t->atomstring;
  int *atomid = t->atomid;
  MD_Impr *impr = tbuf->impr;
  int n, k, index, atom1, atom2, atom3, atom4;
  char *name1, *name2, *name3, *name4;

  /* determine number of impropers to be read on this call */
  n = tbuf->impr_max - t->nimpr;
  if (tbuf->impr_len < n)  n = tbuf->impr_len;

  /* read n impropers and set topology info */
  for (k = 0;  k < n;  k++, t->nimpr++) {
    fscanf(f, "%d %d %d %d", &atom1, &atom2, &atom3, &atom4);
    MDIO_ASSERT(atom1 > 0 && atom1 <= t->natom
             && atom2 > 0 && atom2 <= t->natom
             && atom3 > 0 && atom3 <= t->natom
             && atom4 > 0 && atom4 <= t->natom);
    name1 = &atomstring[ atomid[atom1 - 1] * ATOMNAME_BUFSIZE ];
    name2 = &atomstring[ atomid[atom2 - 1] * ATOMNAME_BUFSIZE ];
    name3 = &atomstring[ atomid[atom3 - 1] * ATOMNAME_BUFSIZE ];
    name4 = &atomstring[ atomid[atom4 - 1] * ATOMNAME_BUFSIZE ];
    index = pdata_find_impr(pdata, name1, name2, name3, name4);
    if (index < 0) {
      index = pdata_find_impr(pdata, name1, WILDCARD, WILDCARD, name4);
    }
    MDIO_ASSERT(index >= 0 && index < pdata->pbuf->imprp_num);
    impr[k].atom[0] = atom1 - 1;
    impr[k].atom[1] = atom2 - 1;
    impr[k].atom[2] = atom3 - 1;
    impr[k].atom[3] = atom4 - 1;
    impr[k].param = index;
  }
  tbuf->impr_num = n;
  return MD_IMPR;
}


static int read_excl(Topoxp *t)
{
  Tbuf *tbuf = t->tbuf;
  FILE *f = t->f;
  int *excl_list = t->excl_list;
  MD_Excl *excl = tbuf->excl;
  int n, k;

  /* determine number of exclusions to be read on this call */
  n = tbuf->excl_max - t->nexcl;
  if (tbuf->excl_len < n)  n = tbuf->excl_len;

  /* read n exclusions and set topology info */
  for (k = 0;  k < n;  k++, t->nexcl++) {
    while (t->excl_start == t->excl_end) {
      t->excl_start = t->excl_end;
      fscanf(f, "%d", &t->excl_end);
      MDIO_ASSERT(t->excl_end >= t->excl_start 
               && t->excl_end < tbuf->excl_max);
      MDIO_ASSERT(t->excl_atom < t->natom);
      t->excl_atom++;
    }
    MDIO_ASSERT(t->excl_start < tbuf->excl_max);
    excl[k].atom[0] = t->excl_atom;
    excl[k].atom[1] = excl_list[ t->excl_start ] - 1;
    t->excl_start++;
  }
  while (t->excl_start == t->excl_end  &&  t->excl_atom < t->natom) {
    t->excl_start = t->excl_end;
    fscanf(f, "%d", &t->excl_end);
    MDIO_ASSERT(t->excl_end >= t->excl_start 
        && (t->excl_end < tbuf->excl_max 
            || (tbuf->excl_max == 0 && t->excl_end == 0)));
    t->excl_atom++;
  }
  tbuf->excl_num = n;
  return MD_EXCL;
}


/* states for reader */
#define  START  0
#define  TITLE  1
#define  ATOM   2
#define  BOND   3
#define  THETA  4
#define  PHI    5
#define  IMPHI  6
#define  DON    7
#define  ACC    8
#define  NB     9

#define  BUFSIZE  1024

int topoxp_read(Topoxp *t)
{
  FILE *f = t->f;
  Tbuf *tbuf = t->tbuf;
  int n, k;
  char *tok, *rest;
  const char *delim = "!: \n\t\r\f\v";
  char buf[BUFSIZE];

  switch (t->state) {
  case START:
    fscanf(f, "%4s", buf);
    if (strcmp(buf, "PSF") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }

    /* setup for reading title */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NTITLE") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    t->state = TITLE;

    /* fall through */
  case TITLE:
    for (k = 0;  k < n;  k++) {
      fgets(buf, BUFSIZE, f);  /* skip n lines */
    }

    /* setup for reading atoms */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NATOM") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    tbuf->atom_max = n;
    t->atomid = (int *) malloc(n * sizeof(int));
    MDIO_ASSERT(t->atomid != NULL);
    t->state = ATOM;

    /* fall through */
  case ATOM:
    if (t->natom < tbuf->atom_max) {
      return read_atom(t);
    }

    /* setup for reading bonds */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NBOND") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    tbuf->bond_max = n;
    t->state = BOND;

    /* fall through */
  case BOND:
    if (t->nbond < tbuf->bond_max) {
      return read_bond(t);
    }

    /* setup for reading angles */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NTHETA") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    tbuf->angle_max = n;
    t->state = THETA;

    /* fall through */
  case THETA:
    if (t->nangle < tbuf->angle_max) {
      return read_angle(t);
    }

    /* setup for reading dihedrals */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NPHI") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    tbuf->dihed_max = n;
    t->state = PHI;

    /* fall through */
  case PHI:
    if (t->ndihed < tbuf->dihed_max) {
      return read_dihed(t);
    }

    /* setup for reading impropers */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NIMPHI") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    tbuf->impr_max = n;
    t->state = IMPHI;

    /* fall through */
  case IMPHI:
    if (t->nimpr < tbuf->impr_max) {
      return read_impr(t);
    }

    /* setup for reading donors */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NDON") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    t->state = DON;

    /* fall through */
  case DON:
    for (k = 0;  k < n;  k++) {
      fscanf(f, "%*d %*d");  /* skip over the donor pairs */
    }

    /* setup for reading acceptors */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NACC") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    t->state = ACC;

    /* fall through */
  case ACC:
    for (k = 0;  k < n;  k++) {
      fscanf(f, "%*d %*d");  /* skip over the acceptor pairs */
    }
    free(t->atomid);      /* free this memory, we are done with it */
    t->atomid = NULL;

    /* setup for reading exclusions */
    fscanf(f, "%d", &n);
    fgets(buf, BUFSIZE, f);
    tok = mdutil_strtok_r(buf, delim, &rest);
    if (strcmp(tok, "NNB") != 0) {
      mdio_fatal("file %s, line %d\n", __FILE__, __LINE__);
    }
    tbuf->excl_max = n;
    t->excl_list = (int *) malloc(n * sizeof(int));
    MDIO_ASSERT(t->excl_list != NULL);
    for (k = 0;  k < n;  k++) {
      fscanf(f, "%d", &t->excl_list[k]);
    }
    t->state = NB;

    /* fall through */
  case NB:
    if (t->excl_atom < t->natom) {
      return read_excl(t);
    }
    break;  /* all done! */

  default:
    mdio_fatal("incorrect state, file %s, line %d\n", __FILE__, __LINE__);
  }

  free(t->excl_list);
  t->excl_list = NULL;
  return 0;
}


int topoxp_init(Topoxp *t, Tbuf *tbuf, Pdata *pdata, FILE *f)
{
  t->tbuf = tbuf;
  t->pdata = pdata;
  t->f = f;
  t->state = START;
  t->status = 0;

  t->natom = 0;
  t->nbond = 0;
  t->nangle = 0;
  t->ndihed = 0;
  t->nimpr = 0;
  t->nexcl = 0;

  t->excl_start = 0;
  t->excl_end = 0;
  t->excl_atom = 0;

  t->excl_list = NULL;
  t->atomid = NULL;

  t->atomstring = (char *) calloc(pdata->pbuf->atomp_num * ATOMNAME_BUFSIZE, 
                                  sizeof(char));
  MDIO_ASSERT(t->atomstring != NULL);
  return 0;
}


void topoxp_destroy(Topoxp *t)
{
  free(t->atomstring);
  free(t->atomid);
  free(t->excl_list);

  /* reset everything */
  memset(t, 0, sizeof(Topoxp));
}


int topoxp_status(Topoxp *t)
{
  return t->status;
}


void topoxp_clear_status(Topoxp *t)
{
  if (ferror(t->f) != 0) {
    clearerr(t->f);
  }
  t->status = 0;
}