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

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <iostream>
#include <fstream>

#include "molfile_plugin.h"
#include "periodic_table.h"

#include "gemmi/cif.hpp"
#include "gemmi/smcif.hpp"
#include "gemmi/to_cif.hpp"


typedef struct {
  int step;
  gemmi::cif::Document *doc;
} cifdata;

typedef struct {
  int numatoms;
  int step;
  std::ofstream *stream;
  molfile_atom_t *atomlist;
} cifwritedata;


static void *open_cif_read(const char *filename, const char *filetype, 
                           int *natoms) {
  cifdata *data = (cifdata *) malloc(sizeof(cifdata));
 
  printf("cif) trying to open file '%s'\n", filename);
  try {
    data->doc = new gemmi::cif::Document(gemmi::cif::read_file(filename));
  } catch (std::exception& e) {
    printf("cif) could not read file: %s\n", e.what());
    return NULL;
  }
  printf("cif) found %zd blocks in CIF file\n", data->doc->blocks.size());

  // Iterate to find the first non-empty block
  *natoms = 0;
  for (int i = 0; i < data->doc->blocks.size(); i++) {
    try {
      gemmi::SmallStructure s = gemmi::make_small_structure_from_block(data->doc->blocks[i]);

      if (s.get_all_unit_cell_sites().size() > 0) {
	*natoms = s.get_all_unit_cell_sites().size();
	data->step = i;
	printf("cif) first non-empty block: %d\n", i+1);
	printf("cif) number of atoms: %d\n", *natoms);
	return data;
      }
    } catch (std::exception& e) {
      /* Simply skip to the next block. */
      printf("cif) could not interpret block %d\n", i+1);
    }
  }

  printf("cif) Error: all blocks in CIF file had zero atoms\n");
  return NULL;
}


static int read_cif_structure(void *mydata, int *optflags, 
                              molfile_atom_t *atoms) {
  int i;
  cifdata *data = (cifdata *) mydata;

  gemmi::SmallStructure s = gemmi::make_small_structure_from_block(data->doc->blocks[data->step]);
  auto sites = s.get_all_unit_cell_sites();

  for (i=0; i<sites.size(); i++) {
    strncpy(atoms[i].name, sites[i].label.c_str(), 16);
    strncpy(atoms[i].type, sites[i].element.name(), 16);
    atoms[i].atomicnumber = sites[i].element.atomic_number();
    atoms[i].mass = sites[i].element.weight();
  }

  *optflags = MOLFILE_MASS | MOLFILE_ATOMICNUMBER;
  return MOLFILE_SUCCESS;
}


static int read_cif_timestep(void *mydata, int natoms, molfile_timestep_t *ts) {
  cifdata *data = (cifdata *) mydata;
  int step = data->step;
  int i, n;

  if (step >= data->doc->blocks.size())
    return MOLFILE_EOF;

  printf("cif) reading block %d/%ld\n", step+1, data->doc->blocks.size());

  data->step += 1;
  if (!ts)
    return MOLFILE_SUCCESS;

  gemmi::SmallStructure s = gemmi::make_small_structure_from_block(data->doc->blocks[step]);
  auto sites = s.get_all_unit_cell_sites();
  n = sites.size();
  if (n != natoms) {
    printf("cif) Number of atoms in block %d is %d, not %d as expected\n", step, n, natoms);
    printf("cif) Stopping there\n");
    return MOLFILE_ERROR;
  }

  printf("cif) space group: %s\n", s.spacegroup_hm.c_str());
  printf("cif) cell parameters: %g %g %g %g %g %g\n", s.cell.a, s.cell.b, s.cell.c,
         s.cell.alpha, s.cell.beta, s.cell.gamma);

  /* cell parameters */
  ts->A = s.cell.a;
  ts->B = s.cell.b;
  ts->C = s.cell.c;
  ts->alpha = s.cell.alpha;
  ts->beta = s.cell.beta;
  ts->gamma = s.cell.gamma;

  /* read the coordinates */
  for (i=0; i<natoms; i++) {
    gemmi::Fractional fract = sites[i].fract.wrap_to_unit();
    gemmi::Position p = s.cell.orthogonalize(fract);

    ts->coords[3*i  ] = p.x;
    ts->coords[3*i+1] = p.y;
    ts->coords[3*i+2] = p.z;
  }

  return MOLFILE_SUCCESS;
}
    
static void close_cif_read(void *mydata) {
  cifdata *data = (cifdata *) mydata;
  delete data->doc;
  free(mydata);
}


static void *open_cif_write(const char *filename, const char *filetype, 
                           int natoms) {
  cifwritedata *data = (cifwritedata *) malloc(sizeof(cifwritedata));

  data->stream = new std::ofstream();
  data->stream->open(filename);
  if (!data->stream->is_open()) {
    printf("cif) Unable to open CIF file %s for writing\n", filename);
    return NULL;
  }

  data->numatoms = natoms;
  data->step = 0;
  return data;
}


static int write_cif_structure(void *mydata, int optflags, 
                               const molfile_atom_t *atoms) {
  cifwritedata *data = (cifwritedata *) mydata;
  size_t sz = data->numatoms * sizeof(molfile_atom_t);

  data->atomlist = (molfile_atom_t *) malloc(sz);
  memcpy(data->atomlist, atoms, sz);

  return MOLFILE_SUCCESS;
}


static int write_cif_timestep(void *mydata, const molfile_timestep_t *ts) {
  cifwritedata *data = (cifwritedata *) mydata;
  const molfile_atom_t *atom = data->atomlist;
  const float *pos = ts->coords;
  const char *label, *symbol;
  char buf[1024];
  gemmi::UnitCell cell;
  bool periodic = (ts->A * ts->B * ts->C > 1.e-7);
  float a, b, c, alpha, beta, gamma;

#define LINE(X) *data->stream << X << std::endl

  LINE("data_struct_" << data->step);
  LINE("_audit_creation_method         'generated by VMD (cifplugin)'");

  LINE("_symmetry_cell_setting         'triclinic'");
  LINE("_symmetry_space_group_name_H-M 'P 1'");
  LINE("loop_");
  LINE("  _symmetry_equiv_pos_as_xyz");
  LINE("  '+x,+y,+z'");

  if (periodic) {
    a = ts->A;
    b = ts->B;
    c = ts->C;
    alpha = ts->alpha;
    beta = ts->beta;
    gamma = ts->gamma;
  } else {
    /* arbitrary cell parameters for nonperiodic systems */
    a = b = c = 100;
    alpha = beta = gamma = 90;
  }

  LINE("_cell_length_a      " << a);
  LINE("_cell_length_b      " << b);
  LINE("_cell_length_c      " << c);
  LINE("_cell_angle_alpha   " << alpha);
  LINE("_cell_angle_beta    " << beta);
  LINE("_cell_angle_gamma   " << gamma);
  cell.set(a, b, c, alpha, beta, gamma);

  LINE("loop_");
  LINE("  _atom_site_label");
  LINE("  _atom_site_type_symbol");
  LINE("  _atom_site_occupancy");
  LINE("  _atom_site_fract_x");
  LINE("  _atom_site_fract_y");
  LINE("  _atom_site_fract_z");

  for (int i = 0; i < data->numatoms; ++i) {
    label = atom->name;
    if (atom->atomicnumber > 0) {
       symbol = pte_label[atom->atomicnumber];
    } else {
       symbol = atom->type;
    }

    gemmi::Position cart(pos[0], pos[1], pos[2]);
    gemmi::Fractional frac = cell.fractionalize(cart);

    snprintf(buf, sizeof(buf), "%s  %s  %8.6f %15.9f %15.9f %15.9f",
	    label, symbol, atom->occupancy, frac.x, frac.y, frac.z);

    LINE(buf);
    ++atom; 
    pos += 3;
  }

#undef LINE

  return MOLFILE_SUCCESS;
}


static void close_cif_write(void *mydata) {
  cifwritedata *data = (cifwritedata *) mydata;

  data->stream->close();
  delete data->stream;
  free(mydata);
}



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 = "cif";
  plugin.prettyname = "CIF";
  plugin.author = "Francois-Xavier Coudert";
  plugin.majorv = 0;
  plugin.minorv = 1;
  plugin.is_reentrant = VMDPLUGIN_THREADSAFE;
  plugin.filename_extension = "cif";
  plugin.open_file_read = open_cif_read;
  plugin.read_structure = read_cif_structure;
  plugin.read_next_timestep = read_cif_timestep;
  plugin.close_file_read = close_cif_read;
  plugin.open_file_write = open_cif_write;
  plugin.write_structure = write_cif_structure;
  plugin.write_timestep = write_cif_timestep;
  plugin.close_file_write = close_cif_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;
}