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

/***************************************************************************
 * RCS INFORMATION:
 *
 *      $RCSfile: jsplugin.c,v $
 *      $Author: johns $       $Locker:  $             $State: Exp $
 *      $Revision: 1.2 $       $Date: 2004/08/04 18:25:18 $
 *
 ***************************************************************************
 * DESCRIPTION:
 *   Code for reading and writing John's molecular dynamics trajectory files
 *   for I/O performance testing and as a simple example plugin.
 *
 *   No programs actually use this format, so it's only a test/example code.
 *
 *   This plugin reads and writes mostly the same data that a DCD file 
 *   contains, but the ordering of I/O operations avoids having to do
 *   scatter/gather passes on the buffers, so it ends up performing 
 *   two to three times faster than the DCD format for the same data set.
 *
 *   Best measured I/O performance in VMD so far is 631 MB/sec on a Sun V880z 
 *   Best standalone I/O performance so far is 688 MB/sec.
 *
 *  Standalone test binary compilation flags:
 *  cc -fast -xarch=v9a -I../../include -DTEST_JSPLUGIN jsplugin.c \
 *    -o ~/bin/readjs -lm
 *
 *  Profiling flags:
 *  cc -xpg -fast -xarch=v9a -g -I../../include -DTEST_JSPLUGIN jsplugin.c \
 *    -o ~/bin/readjs -lm
 *
 ***************************************************************************/

#if defined(_AIX)
/* Define to enable large file extensions on AIX */
#define _LARGE_FILE
#else
/* Defines which enable LFS I/O interfaces for large (>2GB) file support 
 * on 32-bit machines.  These must be defined before inclusion of any 
 * system headers.
 */
#define _LARGEFILE_SOURCE
#define _FILE_OFFSET_BITS 64
#endif

#include <sys/stat.h>
#include <sys/types.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include "fastio.h"
#include "endianswap.h"

#include "molfile_plugin.h"

#ifndef M_PI_2
#define M_PI_2 1.57079632679489661922
#endif

#define JSHEADERSTRING   "JS Binary Trajectory File Format"                
#define JSMAGICNUMBER    0x31337
#define JSENDIANISM      0x12345678
#define JSMAJORVERSION   1
#define JSMINORVERSION   0

#define JSNFRAMESOFFSET  (strlen(JSHEADERSTRING) + 20)

#define JSNOERR          0
#define JSBADFILE        1
#define JSBADFORMAT      2

typedef struct {
  fio_fd fd;
  int natoms;
  int nframes;
  double tsdelta;
  int reverse;
  int with_unitcell;
} jshandle;

static void *open_js_read(const char *path, const char *filetype, int *natoms) {
  jshandle *js;
  int jsmagicnumber, jsendianism, jsmajorversion, jsminorversion;
  struct stat stbuf;
  char strbuf[1024];
  int rc = 0;

  if (!path) return NULL;

  /* See if the file exists, and get its size */
  memset(&stbuf, 0, sizeof(struct stat));
  if (stat(path, &stbuf)) {
    fprintf(stderr, "Could not access file '%s'.\n", path);
    return NULL;
  }

  js = (jshandle *)malloc(sizeof(jshandle));
  memset(js, 0, sizeof(jshandle));
  if (fio_open(path, FIO_READ, &js->fd) < 0) {
    fprintf(stderr, "Could not open file '%s' for reading.\n", path);
    free(js);
    return NULL;
  }

  /* emit header information */
  fio_fread(strbuf, strlen(JSHEADERSTRING), 1, js->fd);
  strbuf[strlen(JSHEADERSTRING)] = '\0';
  if (strcmp(strbuf, JSHEADERSTRING)) {
    fprintf(stderr, "Bad trajectory header!\n");
    fprintf(stderr, "Read string: %s\n", strbuf);
    return NULL;
  }

  fio_read_int32(js->fd, &jsmagicnumber);
  fio_read_int32(js->fd, &jsendianism);
  fio_read_int32(js->fd, &jsmajorversion);
  fio_read_int32(js->fd, &jsminorversion);
  fio_read_int32(js->fd, &js->natoms);
  fio_read_int32(js->fd, &js->nframes);

  if ((jsmagicnumber != JSMAGICNUMBER) || (jsendianism != JSENDIANISM)) {
    fprintf(stderr, "Opposite endianism trajectory file, enabling byte swapping\n");
    js->reverse = 1;
    swap4_aligned(&jsmagicnumber, 1);
    swap4_aligned(&jsendianism, 1);
    swap4_aligned(&jsmajorversion, 1);
    swap4_aligned(&jsminorversion, 1);
    swap4_aligned(&js->natoms, 1);
    swap4_aligned(&js->nframes, 1);
  }

  if ((jsmagicnumber != JSMAGICNUMBER) || (jsendianism != JSENDIANISM)) {
    fprintf(stderr, "read_jsreader returned %d\n", rc);
    fio_fclose(js->fd);
    free(js);
    return NULL;
  }

  *natoms = js->natoms;
  return js;
}


static int read_js_timestep(void *v, int natoms, molfile_timestep_t *ts) {
  jshandle *js = (jshandle *)v;
  int framelen = (natoms * 3 * sizeof(float)) + (6 * sizeof(double));;

  /* if we have a valid ts pointer, read the timestep, otherwise skip it */ 
  if (ts != NULL) {
    int readlen; 
    fio_iovec iov[2];
    double unitcell[6];
    unitcell[0] = unitcell[2] = unitcell[5] = 1.0f;
    unitcell[1] = unitcell[3] = unitcell[4] = 90.0f;
  
    /* setup the I/O vector */
    iov[0].iov_base = (fio_caddr_t) ts->coords;   /* read coordinates    */
    iov[0].iov_len  = natoms * 3 * sizeof(float);
    iov[1].iov_base = (fio_caddr_t) &unitcell[0]; /* read PBC unit cell  */
    iov[1].iov_len  = 6 * sizeof(double);

    /* Do all of the reads with a single syscall, for peak efficiency. */
    /* On smart kernels, readv() causes only one context switch, and   */
    /* can effeciently scatter the reads to the various buffers.       */
    readlen = fio_readv(js->fd, &iov[0], 2); 
  
    /* check the number of read bytes versus what we expected */
    if (readlen != framelen)
      return MOLFILE_EOF;

    /* perform byte swapping if necessary */
    if (js->reverse) {
      swap4_aligned(ts->coords, natoms * 3);
      swap8_aligned(unitcell, 6);
    }

    /* copy unit cell values into VMD */
    ts->A = unitcell[0];
    ts->B = unitcell[1];
    ts->C = unitcell[2];
    ts->alpha = 90.0 - asin(unitcell[3]) * 90.0 / M_PI_2;
    ts->beta  = 90.0 - asin(unitcell[4]) * 90.0 / M_PI_2;
    ts->gamma = 90.0 - asin(unitcell[5]) * 90.0 / M_PI_2;
  } else {
    /* skip this frame, seek to the next frame */
    if (fio_fseek(js->fd, framelen, FIO_SEEK_CUR)) 
      return MOLFILE_EOF;
  }
 
  return MOLFILE_SUCCESS;
}
 

static void close_js_read(void *v) {
  jshandle *js = (jshandle *)v;
  fio_fclose(js->fd);
  free(js); 
}


static void *open_js_write(const char *path, const char *filetype, int natoms) {
  jshandle *js;

  js = (jshandle *)malloc(sizeof(jshandle));
  memset(js, 0, sizeof(jshandle));
  if (fio_open(path, FIO_WRITE, &js->fd) < 0) {
    fprintf(stderr, "Could not open file %s for writing\n", path);
    free(js);
    return NULL;
  }

  js->natoms = natoms;
  js->with_unitcell = 1;

  /* emit header information */
  fio_write_str(js->fd, JSHEADERSTRING);
  fio_write_int32(js->fd, JSMAGICNUMBER);
  fio_write_int32(js->fd, JSENDIANISM);
  fio_write_int32(js->fd, JSMAJORVERSION);
  fio_write_int32(js->fd, JSMINORVERSION);

  /* write number of atoms */
  fio_write_int32(js->fd, natoms);

  /* write number of frames, to be updated later */
  js->nframes = 0;
  fio_write_int32(js->fd, js->nframes);

  return js;
}

static int write_js_timestep(void *v, const molfile_timestep_t *ts) { 
  jshandle *js = (jshandle *)v;
  double unitcell[6];

  js->nframes++; /* increment frame count written to the file so far */

  unitcell[0] = ts->A;
  unitcell[1] = ts->B;
  unitcell[2] = ts->C;
  unitcell[3] = sin((M_PI_2 / 90.0) * (90.0 - ts->alpha));
  unitcell[4] = sin((M_PI_2 / 90.0) * (90.0 - ts->beta));
  unitcell[5] = sin((M_PI_2 / 90.0) * (90.0 - ts->gamma));

  /* coordinates for all atoms */
  fio_fwrite(ts->coords, js->natoms * 3 * sizeof(float), 1, js->fd);

  /* PBC unit cell info */ 
  fio_fwrite(&unitcell[0], 6 * sizeof(double), 1, js->fd);

  return MOLFILE_SUCCESS;
}

static void close_js_write(void *v) {
  jshandle *js = (jshandle *)v;

  /* update the trajectory header information */
  fio_fseek(js->fd, JSNFRAMESOFFSET, FIO_SEEK_SET);
  fio_write_int32(js->fd, js->nframes);
  fio_fseek(js->fd, 0, FIO_SEEK_END);

  fio_fclose(js->fd);
  free(js);
}


/*
 * Initialization stuff here
 */
static molfile_plugin_t jsplugin = {
  vmdplugin_ABIVERSION,                         /* ABI version */
  MOLFILE_PLUGIN_TYPE,                          /* type */
  "js",                                         /* name */
  "John Stone",                                 /* author */
  0,                                            /* major version */
  6,                                            /* minor version */
  VMDPLUGIN_THREADSAFE,                         /* is reentrant  */
  "js",                                         /* filename extension */
  open_js_read,
  0,
  0,
  read_js_timestep,
  close_js_read,
  open_js_write,
  0,
  write_js_timestep,
  close_js_write
};

int VMDPLUGIN_init() {
  return VMDPLUGIN_SUCCESS;
}

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

int VMDPLUGIN_fini() {
  return VMDPLUGIN_SUCCESS;
}

  
#ifdef TEST_JSPLUGIN

#include <sys/time.h>

/* get the time of day from the system clock, and store it (in seconds) */
double time_of_day(void) {
#if defined(_MSC_VER)
  double t;

  t = GetTickCount();
  t = t / 1000.0;

  return t;
#else
  struct timeval tm;
  struct timezone tz;

  gettimeofday(&tm, &tz);
  return((double)(tm.tv_sec) + (double)(tm.tv_usec)/1000000.0);
#endif
}

int main(int argc, char *argv[]) {
  molfile_timestep_t timestep;
  void *v;
  jshandle *js;
  int i, natoms;
  float sizeMB =0.0, totalMB = 0.0;
  double starttime, endtime, totaltime = 0.0;

  while (--argc) {
    ++argv; 
    natoms = 0;
    v = open_js_read(*argv, "js", &natoms);
    if (!v) {
      fprintf(stderr, "open_js_read failed for file %s\n", *argv);
      return 1;
    }
    js = (jshandle *)v;
    sizeMB = ((natoms * 3.0) * js->nframes * 4.0) / (1024.0 * 1024.0);
    totalMB += sizeMB; 
    printf("file: %s\n", *argv);
    printf("  %d atoms, %d frames, size: %6.1fMB\n", natoms, js->nframes, sizeMB);

    starttime = time_of_day();
    timestep.coords = (float *)malloc(3*sizeof(float)*natoms);
    for (i=0; i<js->nframes; i++) {
      int rc = read_js_timestep(v, natoms, &timestep);
      if (rc) {
        fprintf(stderr, "error in read_js_timestep on frame %d\n", i);
        return 1;
      }
    }
    endtime = time_of_day();
    close_js_read(v);
    totaltime += endtime - starttime;
    printf("  Time: %5.1f seconds\n", endtime - starttime);
    printf("  Speed: %5.1f MB/sec, %5.1f timesteps/sec\n", sizeMB / (endtime - starttime), (js->nframes / (endtime - starttime)));
  }
  printf("Overall Size: %6.1f MB\n", totalMB);
  printf("Overall Time: %6.1f seconds\n", totaltime);
  printf("Overall Speed: %5.1f MB/sec\n", totalMB / totaltime);
  return 0;
}
      
#endif