/******************************************************************************
  This source file is part of the Avogadro project.
  This source code is released under the 3-Clause BSD License, (see "LICENSE").
******************************************************************************/

#include "dcdformat.h"
#include "struct.h"

#include <avogadro/core/elements.h>
#include <avogadro/core/molecule.h>
#include <avogadro/core/unitcell.h>
#include <avogadro/core/utilities.h>
#include <avogadro/core/vector.h>

#include <cmath>
#include <istream>
#include <ostream>
#include <string>

using std::map;
using std::string;
using std::to_string;
using std::vector;

namespace Avogadro::Io {

using Core::Array;
using Core::Atom;
using Core::Molecule;
using Core::UnitCell;

#ifndef _WIN32
#endif

#define DCD_EOF (-1)
constexpr int DCD_MAGIC = 84;
constexpr int DCD_IS_CHARMM = 0x01;
constexpr int DCD_HAS_4DIMS = 0x02;
constexpr int DCD_HAS_EXTRA_BLOCK = 0x04;

int swap_integer(int inp)
{
  return (((inp << 24) & 0xff000000) | ((inp << 8) & 0x00ff0000) |
          ((inp >> 8) & 0x0000ff00) | ((inp >> 24) & 0x000000ff));
}

char swap_endian(char endian)
{
  if (endian == '>')
    return '<';
  else
    return '>';
}

bool DcdFormat::read(std::istream& inStream, Core::Molecule& mol)
{
  /** Endian type, Buffer and Format char containers for unpacking and storing
   * data using struct library */
  char endian = '>';
  char buff[BUFSIZ];
  char fmt[BUFSIZ];

  /** Variables to store various components from the binary data unpacked using
   * the struct library */
  char raw[84];
  char* remarks;
  double DELTA;
  int magic;
  int charmm;
  int NAMNF;
  int NTITLE;
  int lenRemarks;
  int NATOMS;
  int blockSize;

  // Determining size of file
  inStream.seekg(0, inStream.end);
  int fileLen = inStream.tellg();
  inStream.seekg(0, inStream.beg);

  // Reading magic number
  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &magic);
  if (magic != DCD_MAGIC) {
    magic = swap_integer(magic);
    endian = swap_endian(endian);
    if (magic != DCD_MAGIC) {
      appendError("File does not start with magic number 84.");
      return false;
    }
  }

  // CORD
  snprintf(fmt, sizeof(fmt), "%c%ds", endian, magic);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, raw);
  if (raw[0] != 'C' || raw[1] != 'O' || raw[2] != 'R' || raw[3] != 'D') {
    appendError("Keyword CORD not found.");
    return false;
  }

  // Determining whether the trajectory file is from CHARMM or not
  if (*(reinterpret_cast<int*>(raw + 80)) != 0) {
    charmm = DCD_IS_CHARMM;
    if (*(reinterpret_cast<int*>(raw + 44)) != 0)
      charmm |= DCD_HAS_EXTRA_BLOCK;

    if (*(reinterpret_cast<int*>(raw + 48)) == 1)
      charmm |= DCD_HAS_4DIMS;
  } else {
    charmm = 0;
  }

  // number of fixed atoms
  NAMNF = *(reinterpret_cast<int*>(raw + 36));

  // DELTA (timestep) is stored as a double with X-PLOR but as a float with
  // CHARMM
  if (charmm & DCD_IS_CHARMM) {
    float ftmp;
    ftmp = *(reinterpret_cast<float*>(raw + 40));

    DELTA = static_cast<double>(ftmp);
  } else {
    (DELTA) = *(reinterpret_cast<double*>(raw + 40));
  }

  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &magic);

  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &blockSize);

  if (((blockSize - 4) % 80) == 0) {
    // Read NTITLE, the number of 80 character title strings
    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &NTITLE);
    lenRemarks = NTITLE * 80;
    remarks = reinterpret_cast<char*>(malloc(lenRemarks));
    snprintf(fmt, sizeof(fmt), "%c%ds", endian, lenRemarks);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, remarks);

    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    int endSize;
    struct_unpack(buff, fmt, &endSize);
  } else {
    appendError("Block size must be 4 plus a multiple of 80.");
    return false;
  }

  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  int fourInput;
  struct_unpack(buff, fmt, &fourInput);
  if (fourInput != 4) {
    // Error
  }

  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &NATOMS);

  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &fourInput);
  if (fourInput != 4) {
    appendError("Expected token 4. Read token " + to_string(fourInput));
    return false;
  }

  if (NAMNF != 0) {
    int** FREEINDEXES =
      reinterpret_cast<int**>(calloc((NATOMS - NAMNF), sizeof(int)));
    if (*FREEINDEXES == nullptr) {
      appendError("MALLOC failed.");
      return false;
    }

    /* Read in index array size */
    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    int arrSize;
    struct_unpack(buff, fmt, &arrSize);

    if (arrSize != (NATOMS - NAMNF) * 4) {
      appendError("DCD file contains bad format.");
      return false;
    }

    snprintf(fmt, sizeof(fmt), "%c%di", endian, (NATOMS - NAMNF));
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, *FREEINDEXES);

    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &arrSize);

    if (arrSize != (NATOMS - NAMNF) * 4) {
      appendError("DCD file contains bad format.");
      return false;
    }
  }

  // CHARMM trajectories have an extra block to be read, that contains
  // information about the unit cell
  if ((charmm & DCD_IS_CHARMM) && (charmm & DCD_HAS_EXTRA_BLOCK)) {
    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    int leadingNum;
    struct_unpack(buff, fmt, &leadingNum);

    if (leadingNum == 48) {
      double unitcell[6];
      for (double& aa : unitcell) {
        snprintf(fmt, sizeof(fmt), "%c%dd", endian, 1);
        inStream.read(buff, struct_calcsize(fmt));
        struct_unpack(buff, fmt, &aa);
      }
      if (unitcell[1] >= -1.0 && unitcell[1] <= 1.0 && unitcell[3] >= -1.0 &&
          unitcell[3] <= 1.0 && unitcell[4] >= -1.0 && unitcell[4] <= 1.0) {
        // CHARMM and certain NAMD files have the cosines instead of angles
        // This formulation improves rounding behavior for orthogonal cells
        // so that the angles end up at precisely 90 degrees, unlike acos()
        unitcell[4] = M_PI_2 - asin(unitcell[4]); /* cosBC */
        unitcell[3] = M_PI_2 - asin(unitcell[3]); /* cosAC */
        unitcell[1] = M_PI_2 - asin(unitcell[1]); /* cosAB */
      }

      auto* cell = new UnitCell(unitcell[0], unitcell[2], unitcell[5],
                                unitcell[4], unitcell[3], unitcell[1]);
      if (!cell->isRegular()) {
        appendError("cell matrix is singular");
        delete cell;
        return false;
      }
      mol.setUnitCell(cell);
    } else {
      inStream.read(buff, leadingNum);
    }
    inStream.read(buff, sizeof(int));
  }

  // Reading the atom coordinates
  int formatint[6];
  Array<float> cx, cy, cz;
  cx.reserve(NATOMS);
  cy.reserve(NATOMS);
  cz.reserve(NATOMS);

  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &formatint[0]);

  for (int i = 0; i < NATOMS; ++i) {
    // X coordinates
    snprintf(fmt, sizeof(fmt), "%c%df", endian, 1);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &cx[i]);
    /* code */
  }

  snprintf(fmt, sizeof(fmt), "%c2i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &formatint[1], &formatint[2]);

  for (int i = 0; i < NATOMS; ++i) {
    // Y coordinates
    snprintf(fmt, sizeof(fmt), "%c%df", endian, 1);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &cy[i]);
  }

  snprintf(fmt, sizeof(fmt), "%c2i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &formatint[3], &formatint[4]);

  for (int i = 0; i < NATOMS; ++i) {
    // Z coordinates
    snprintf(fmt, sizeof(fmt), "%c%df", endian, 1);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &cz[i]);
  }

  snprintf(fmt, sizeof(fmt), "%c1i", endian);
  inStream.read(buff, struct_calcsize(fmt));
  struct_unpack(buff, fmt, &formatint[5]);

  typedef map<string, unsigned char> AtomTypeMap;
  AtomTypeMap atomTypes;
  unsigned char customElementCounter = CustomElementMin;

  for (int i = 0; i < NATOMS; ++i) {
    Vector3 pos(cx[i], cy[i], cz[i]);

    AtomTypeMap::const_iterator it;
    atomTypes.insert(std::make_pair(to_string(i), customElementCounter++));
    it = atomTypes.find(to_string(i));
    // if (customElementCounter > CustomElementMax) {
    //   appendError("Custom element type limit exceeded.");
    //   return false;
    // }
    Atom newAtom = mol.addAtom(it->second);
    newAtom.setPosition3d(pos);
  }

  mol.setTimeStep(0, 0);

  // Skipping fourth dimension block
  if ((charmm & DCD_IS_CHARMM) && (charmm & DCD_HAS_EXTRA_BLOCK)) {
    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    int sizeToRead;
    struct_unpack(buff, fmt, &sizeToRead);

    inStream.read(buff, sizeToRead);

    inStream.read(buff, sizeof(int));
  }

  // Set the custom element map if needed
  if (!atomTypes.empty()) {
    Molecule::CustomElementMap elementMap;
    for (const auto& atomType : atomTypes) {
      elementMap.insert(
        std::make_pair(atomType.second, "Atom " + atomType.first));
    }
    mol.setCustomElementMap(elementMap);
  }

  mol.setCoordinate3d(mol.atomPositions3d(), 0);

  // Do we have an animation?
  int coordSet = 1;
  while ((static_cast<int>(inStream.tellg()) != fileLen) &&
         (static_cast<int>(inStream.tellg()) != DCD_EOF)) {
    // Reading the atom coordinates
    Array<Vector3> positions;
    positions.reserve(NATOMS);

    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &formatint[0]);

    for (int i = 0; i < NATOMS; ++i) {
      // X coordinates
      snprintf(fmt, sizeof(fmt), "%c%df", endian, 1);
      inStream.read(buff, struct_calcsize(fmt));
      struct_unpack(buff, fmt, &cx[i]);
    }

    snprintf(fmt, sizeof(fmt), "%c2i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &formatint[1], &formatint[2]);

    for (int i = 0; i < NATOMS; ++i) {
      // Y coordinates
      snprintf(fmt, sizeof(fmt), "%c%df", endian, 1);
      inStream.read(buff, struct_calcsize(fmt));
      struct_unpack(buff, fmt, &cy[i]);
    }

    snprintf(fmt, sizeof(fmt), "%c2i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &formatint[3], &formatint[4]);

    for (int i = 0; i < NATOMS; ++i) {
      // Z coordinates
      snprintf(fmt, sizeof(fmt), "%c%df", endian, 1);
      inStream.read(buff, struct_calcsize(fmt));
      struct_unpack(buff, fmt, &cz[i]);
    }

    snprintf(fmt, sizeof(fmt), "%c1i", endian);
    inStream.read(buff, struct_calcsize(fmt));
    struct_unpack(buff, fmt, &formatint[5]);

    for (int i = 0; i < NATOMS; ++i) {
      Vector3 pos(cx[i], cy[i], cz[i]);
      positions.push_back(pos);
    }

    mol.setTimeStep(DELTA * coordSet, coordSet);

    // Skipping fourth dimension block
    if ((charmm & DCD_IS_CHARMM) && (charmm & DCD_HAS_EXTRA_BLOCK)) {
      snprintf(fmt, sizeof(fmt), "%c1i", endian);
      inStream.read(buff, struct_calcsize(fmt));
      int sizeToRead;
      struct_unpack(buff, fmt, &sizeToRead);

      inStream.read(buff, sizeToRead);

      inStream.read(buff, sizeof(int));
    }

    mol.setCoordinate3d(positions, coordSet++);
  }

  return true;
}

bool DcdFormat::write(std::ostream&, const Core::Molecule&)
{
  return false;
}

std::vector<std::string> DcdFormat::fileExtensions() const
{
  std::vector<std::string> ext;
  ext.emplace_back("dcd");
  return ext;
}

std::vector<std::string> DcdFormat::mimeTypes() const
{
  std::vector<std::string> mime;
  mime.emplace_back("application/octet-stream");
  return mime;
}

} // namespace Avogadro::Io
