/***************************************************************************
 *cr
 *cr            (C) Copyright 1995-2006 The Board of Trustees of the
 *cr                        University of Illinois
 *cr                         All Rights Reserved
 *cr
 ***************************************************************************/

/***************************************************************************
 * RCS INFORMATION:
 *
 *      $RCSfile: lammpsplugin.c,v $
 *      $Author: johns $       $Locker:  $             $State: Exp $
 *      $Revision: 1.23 $       $Date: 2009/02/24 05:10:14 $
 *
 ***************************************************************************/

/*
 *  LAMMPS atom style dump file format:
 *    ITEM: TIMESTEP
 *      %d (timestep number)
 *    ITEM: NUMBER OF ATOMS
 *      %d (number of atoms)
 *    ITEM: BOX BOUNDS
 *      %f %f (boxxlo, boxxhi)
 *      %f %f (boxylo, boxyhi)
 *      %f %f (boxzlo, boxzhi)
 *    ITEM: ATOMS
 *      %d %d %f %f %f  (atomid, atomtype, x, y, z)
 *      ...
 *
 */

#include "largefiles.h"   /* platform dependent 64-bit file I/O defines */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <math.h>
#include "molfile_plugin.h"

#define THISPLUGIN plugin
#include "vmdconio.h"

#define VMDPLUGIN_STATIC
#include "hash.h"

#ifndef M_PI_2
#define M_PI_2 1.57079632679489661922
#endif

/* for transparent reading of .gz files */
#ifdef _USE_ZLIB
#include <zlib.h>
#define FileDesc gzFile
#define myFgets(buf,size,fd) gzgets(fd,buf,size)
#define myFprintf gzprintf
#define myFopen gzopen
#define myFclose gzclose
#define myRewind gzrewind
#else
#define FileDesc FILE*
#define myFprintf fprintf
#define myFopen fopen
#define myFclose fclose
#define myFgets(buf,size,fd) fgets(buf,size,fd)
#define myRewind rewind
#endif

typedef struct {
  FileDesc file;
  char *file_name;
  int *atomtypes;
  int numatoms;
  int nstep;
  int scaled_data;
} lammpsdata;

static char* find_next_item(FileDesc fd, char* szItem) {
  char szLine[1024];
  char* pcolon;
  int nReturn;

  while(myFgets(szLine, 1024, fd)) {
    pcolon = strrchr(szLine, ':');
    if(pcolon) {
      *pcolon = 0;
      if(0 == strcmp(szLine, "ITEM")) {
        *pcolon = ':';
        pcolon+=2;
        strcpy(szItem, pcolon);
        nReturn = (int)strcspn(szItem, "\n\r"); /* remove the carriage return */
        if(nReturn)
          szItem[nReturn] = 0;
        return szItem;
      }
    }
  }

  return NULL;
}


static int find_item(FileDesc fd, char* szItem) {
  char szBuffer[1024];
  while(find_next_item(fd, szBuffer)) {
    if(0 == strcmp(szBuffer, szItem))
      return 1;
  }
  return 0;
}

 
static void *open_lammps_read(const char *filename, const char *filetype, 
                           int *natoms) {
  FileDesc fd;
  lammpsdata *data;
  char szBuffer[1024];

  fd = myFopen(filename, "rb");
  if (!fd) return NULL;
 
  data = (lammpsdata *)malloc(sizeof(lammpsdata));
  data->file = fd;
  data->file_name = strdup(filename);
  *natoms = 0;

  if(!find_item(data->file, "NUMBER OF ATOMS")) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) Unable to find 'NUMBER OF ATOMS' item\n");
    return NULL;
  }

  if(!myFgets(szBuffer, 1024, data->file)) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) dump file '%s' should have the number of atoms after line ITEM: NUMBER OF ATOMS.\n", filename);
    return NULL;
  }
  *natoms = atoi(szBuffer);
 
  data->numatoms=*natoms;
  data->scaled_data=1;
  myRewind(data->file); /* prepare for first read_timestep call */
 
  return data;
}


static int read_lammps_structure(void *mydata, int *optflags, molfile_atom_t *atoms) {
  int i, j, idx;
  char szBuffer[1024];
  molfile_atom_t *atom;
  lammpsdata *data = (lammpsdata *)mydata;
  int atomid;
  char atom_type[1025];
  float x, y, z;
  char *k;

  *optflags = MOLFILE_NOOPTIONS; /* no optional data */
  memset(atoms, 0, data->numatoms * sizeof(molfile_atom_t)); /* clear atoms structure */
 
  /* go to the beginning of the file */
  myRewind(data->file); /* prepare for first read_timestep call */
 
  /* find the sections with atoms */
  if(!find_item(data->file, "ATOMS")) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) couldn't find atoms to read structure from file '%s'.\n", data->file_name);
    return MOLFILE_ERROR;
  }

  /* now read the atoms */
  for(i=0;i<data->numatoms;i++) {
    k = myFgets(szBuffer, 1024, data->file);

    /* Read in atom type, X, Y, Z, skipping any remaining data fields */
    j = 0;
    if (k != NULL)
        j = sscanf(szBuffer, "%d %s %f %f %f", &atomid, atom_type, &x, &y, &z);
  
    if (k == NULL) {
      vmdcon_printf(VMDCON_ERROR, "lammpsplugin) Error while reading structure from lammps dump file '%s': atom missing in the first timestep\n",data->file_name);
      vmdcon_printf(VMDCON_ERROR, "lammpsplugin) expecting '%d' atoms, found only '%d'\n",data->numatoms,i+1);
      return MOLFILE_ERROR;
    } else if (j < 5) {
      vmdcon_printf(VMDCON_ERROR, "lammpsplugin) Error while reading structure from lammps dump file '%s': missing type or coordinate(s) for atom '%d'\n", data->file_name, i+1);
      return MOLFILE_ERROR;
    }

    if ((atomid <= 0) || (atomid > data->numatoms)) {
      vmdcon_printf(VMDCON_ERROR, "lammpsplugin) Error while reading structure from lammps dump file '%s': atomid %d is not in valid range [1, %d]\n",data->file_name, atomid, data->numatoms);
      return MOLFILE_ERROR;
    }

    idx = atomid - 1;
    strncpy(atoms[idx].type, atom_type, 16); /* type is char[16]. */
    atoms[idx].type[15] = '\0';              /* force null termination */
    /* WARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNING *
     * Don't use the atomid as name. This will only waste a _lot_ of memory.  *
     * The _identical_ information is available via the 'serial' or 'index'   *
     * atom properties (numbers starting from 1 or 0, respectively).          *
     * WARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNINGWARNING */
    strncpy(atoms[idx].name, atom_type, 16); /* name is char[16]. */
    atoms[idx].name[15] = '\0';              /* force null termination */
    strncpy(atoms[idx].resname, "UNK", 4);
    atoms[idx].resid = 1;
    strncpy(atoms[idx].chain, "",1);
    strncpy(atoms[idx].segid, "",1);
    /* skip to the end of line */

    /* coordinates can be either scaled (fractions of the box)
       or unscaled (absolute coordinates) */
    if ((x<-0.1) || (x>1.1) || (y<-0.1) || (y>1.1) || (z<-0.1) || (x>1.1)) {
        data->scaled_data=0;
    }
  }

  myRewind(data->file);
  return MOLFILE_SUCCESS;
}



static int read_lammps_timestep(void *mydata, int natoms, 
                                molfile_timestep_t *ts) {
  int i, j;
  char atom_name[1025], szBuffer[1025];
  float x, y, z;
  int atomid;
  float lo[3],hi[3], l[3];

  lammpsdata *data = (lammpsdata *)mydata;

  /* check if there is antother time steps */
  if(!find_item(data->file, "TIMESTEP")) return MOLFILE_ERROR;
 
  /* check if we actually have to read something */
  if(!ts) return MOLFILE_SUCCESS;

  /* check the number of atoms in the timestep */
  if(!find_item(data->file, "NUMBER OF ATOMS")) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) lammps dump file '%s', unable to find item: NUMBER OF ATOMS for current timestep.\n", data->file_name);
    return MOLFILE_ERROR;
  }

  if(!myFgets(szBuffer, 1024, data->file)) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) lammps dump file '%s', should have the number of atoms after line ITEM: NUMBER OF ATOMS.\n", data->file_name);
    return MOLFILE_ERROR;
  }

  if(natoms != atoi(szBuffer)) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) lammps dump file '%s', wrong number of atoms in timestep.\n", data->file_name);
    return MOLFILE_ERROR;
  }

  /* read the boundary box of the timestep */
  if(!find_item(data->file, "BOX BOUNDS")) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) lammps dump file '%s', could not find ITEM: BOX BOUNDS.\n", data->file_name);
    return MOLFILE_ERROR;
  }

  if (NULL == myFgets(szBuffer, 1024, data->file))  return MOLFILE_ERROR;
  sscanf(szBuffer,"%f %f", lo, hi);
  l[0] = hi[0] - lo[0];

  if (NULL == myFgets(szBuffer, 1024, data->file))  return MOLFILE_ERROR;
  sscanf(szBuffer,"%f %f", lo+1, hi+1);
  l[1] = hi[1] - lo[1];

  if (NULL == myFgets(szBuffer, 1024, data->file))  return MOLFILE_ERROR;
  sscanf(szBuffer,"%f %f", lo+2, hi+2);
  l[2] = hi[2] - lo[2];
 
  ts->A = l[0];
  ts->B = l[1];
  ts->C = l[2];
  /* XXX: non-orthogonal boxes are now possible, too. */
  ts->alpha = 90.0;
  ts->beta = 90.0;
  ts->gamma = 90.0;

  /* read the coordinates */
  if (!find_item(data->file, "ATOMS")) {
    vmdcon_printf(VMDCON_ERROR, "lammpsplugin) lammps dump file '%s', could not find ITEM: ATOMS.\n", data->file_name);
    return MOLFILE_ERROR;
  }

  for (i=0; i<natoms; i++) {
    if(!myFgets(szBuffer, 1024, data->file)) {
      /* there is an error or end of file, both cases should not occur */
      vmdcon_printf(VMDCON_ERROR, "lammpsplugin) lammps dump file '%s', unexpected end of file or error while reading coordinates for atom %d.\n", data->file_name, i);
      return MOLFILE_ERROR;
    }

    /* Read in atom type, X, Y, Z, skipping any remaining data fields */
    j = sscanf(szBuffer, "%d %s %f %f %f", &atomid, atom_name, &x, &y, &z);
    if (j != 5) {
      vmdcon_printf(VMDCON_ERROR, "lammpsplugin) lammps dump file '%s', error while parsing coordinates for atom %d.\n", data->file_name, i);
      return MOLFILE_ERROR;
    }


    if (atomid > 0 && atomid <= natoms) {
      int addr = 3 * (atomid - 1);
      if(data->scaled_data) {
          /* LAMMPS coordinates are fractional unit cell coords 
           * by default, so they need to be scaled by a/b/c etc. */
          ts->coords[addr    ] = lo[0] + x * l[0];
          ts->coords[addr + 1] = lo[1] + y * l[1];
          ts->coords[addr + 2] = lo[2] + z * l[2];
      } else {
          /* ... but they can also be absolute values */
          ts->coords[addr    ] = x;
          ts->coords[addr + 1] = y;
          ts->coords[addr + 2] = z;
      }
    } else {
      vmdcon_printf(VMDCON_WARN, "lammpsplugin) ignoring illegal atom index %d\n", atomid);
    }
  }

  return MOLFILE_SUCCESS;
}
    
static void close_lammps_read(void *mydata) {
  lammpsdata *data = (lammpsdata *)mydata;
  myFclose(data->file);
  free(data->file_name);
  free(data);
}

static void *open_lammps_write(const char *filename, const char *filetype, 
                           int natoms) {
  FileDesc *fd;
  lammpsdata *data;

  fd = myFopen(filename, "w");
  if (!fd) { 
    vmdcon_printf(VMDCON_ERROR, "Error) Unable to open lammpstrj file %s for writing\n",
            filename);
    return NULL;
  }
  
  data = (lammpsdata *)malloc(sizeof(lammpsdata));
  data->numatoms = natoms;
  data->file = fd;
  data->file_name = strdup(filename);
  data->nstep = 0;
  return data;
}

static int write_lammps_structure(void *mydata, int optflags, 
                               const molfile_atom_t *atoms) {
  lammpsdata *data = (lammpsdata *)mydata;
  int i, j;
  hash_t atomtypehash;

  hash_init(&atomtypehash,127);

  /* generate 1 based lookup table for atom types */
  for (i=0, j=1; i < data->numatoms; i++)
    if (hash_insert(&atomtypehash, atoms[i].type, j) == HASH_FAIL)
      j++;
  
  data->atomtypes = (int *) malloc(data->numatoms * sizeof(int));

  for (i=0; i < data->numatoms ; i++)
    data->atomtypes[i] = hash_lookup(&atomtypehash, atoms[i].type);

  hash_destroy(&atomtypehash);
  
  return MOLFILE_SUCCESS;
}

static int write_lammps_timestep(void *mydata, const molfile_timestep_t *ts) {
  lammpsdata *data = (lammpsdata *)mydata; 
  const float *pos;
  const int  *type;
  int i;

  myFprintf(data->file, "ITEM: TIMESTEP\n");
  myFprintf(data->file, "%d\n", data->nstep);
  myFprintf(data->file, "ITEM: NUMBER OF ATOMS\n");
  myFprintf(data->file, "%d\n", data->numatoms);
  myFprintf(data->file, "ITEM: BOX BOUNDS\n");
  myFprintf(data->file, "%g %g\n", 0.0, ts->A);
  myFprintf(data->file, "%g %g\n", 0.0, ts->B);
  myFprintf(data->file, "%g %g\n", 0.0, ts->C);
  myFprintf(data->file, "ITEM: ATOMS\n");

  pos = ts->coords;
  
  for (i = 0; i < data->numatoms; ++i) {
    myFprintf(data->file, " %d %d %g %g %g\n", 
            i+1, data->atomtypes[i], pos[0], pos[1], pos[2]);
    pos += 3;
  }

  data->nstep ++;
  return MOLFILE_SUCCESS;
}


static void close_lammps_write(void *mydata) {
  lammpsdata *data = (lammpsdata *)mydata;

  myFclose(data->file);
  free(data->atomtypes);
  free(data->file_name);
  free(data);
}


/* registration stuff */
static molfile_plugin_t plugin;

VMDPLUGIN_API int VMDPLUGIN_init() {
  memset(&plugin, 0, sizeof(molfile_plugin_t));
  plugin.abiversion = vmdplugin_ABIVERSION;
  plugin.type = MOLFILE_PLUGIN_TYPE;
  plugin.name = "lammpstrj";
  plugin.prettyname = "LAMMPS Trajectory";
  plugin.author = "Marco Kalweit, Axel Kohlmeyer, Lutz Maibaum, John Stone";
  plugin.majorv = 0;
  plugin.minorv = 10;
  plugin.is_reentrant = VMDPLUGIN_THREADSAFE;
#ifdef _USE_ZLIB
  plugin.filename_extension = "lammpstrj,lammpstrj.gz";
#else
  plugin.filename_extension = "lammpstrj";
#endif
  plugin.open_file_read = open_lammps_read;
  plugin.read_structure = read_lammps_structure;
  plugin.read_next_timestep = read_lammps_timestep;
  plugin.close_file_read = close_lammps_read;
  plugin.open_file_write = open_lammps_write;
  plugin.write_structure = write_lammps_structure;
  plugin.write_timestep = write_lammps_timestep;
  plugin.close_file_write = close_lammps_write;
  return VMDPLUGIN_SUCCESS;
}

VMDPLUGIN_API int VMDPLUGIN_register(void *v, vmdplugin_register_cb cb) {
  (*cb)(v, (vmdplugin_t *)&plugin);
  return VMDPLUGIN_SUCCESS;
}

VMDPLUGIN_API int VMDPLUGIN_fini() {
  return VMDPLUGIN_SUCCESS;
}


#ifdef TEST_PLUGIN

int main(int argc, char *argv[]) {
  molfile_timestep_t timestep;
  molfile_atom_t *atoms;
  void *v;
  int natoms;
  int i, nsets, set, opts;

  while (--argc) {
    ++argv;
    v = open_lammps_read(*argv, "lammps", &natoms);
    if (!v) {
      vmdcon_printf(VMDCON_ERROR, "open_lammps_read failed for file %s\n", *argv);
      return 1;
    }
    vmdcon_printf(VMDCON_INFO, "open_lammps_read succeeded for file %s\n", *argv);
    vmdcon_printf(VMDCON_INFO, "number of atoms: %d\n", natoms);

    timestep.coords = (float *)malloc(3*sizeof(float)*natoms);
    atoms = (molfile_atom_t *)malloc(sizeof(molfile_atom_t)*natoms);
    read_lammps_structure(v, &opts, atoms);
    vmdcon_printf(VMDCON_INFO, "read_lammps_structure: options=0x%08x\n", opts);
#if 0
    for (i=0; i<natoms; ++i) {
      vmdcon_printf(VMDCON_INFO, "atom %09d: name=%s, type=%s, resname=%s, resid=%d, segid=%s, chain=%s\n",
                      i, atoms[i].name, atoms[i].type, atoms[i].resname, atoms[i].resid,
                      atoms[i].segid, atoms[i].chain);
    }
#endif
    i = 0;
    while (!read_lammps_timestep(v, natoms, &timestep)) {
      i++;
    }
    vmdcon_printf(VMDCON_INFO, "ended read_next_timestep on frame %d\n", i);

    close_lammps_read(v);
  }
  return 0;
}

#endif