//
// Copyright (C) 2003-2016 Sereina Riniker, Paolo Tosco
//
//   @@ All Rights Reserved @@
//  This file is part of the RDKit.
//  The contents are covered by the terms of the BSD license
//  which is included in the file license.txt, found at the root
//  of the RDKit source tree.
//

#include <RDGeneral/Invariant.h>
#include <RDGeneral/BadFileException.h>
#include <GraphMol/ROMol.h>
#include <GraphMol/Conformer.h>
#include <fstream>
#include <sstream>
#include <set>
#include "Trajectory.h"

namespace RDKit {

RDGeom::Point2D Snapshot::getPoint2D(unsigned int pointNum) const {
  PRECONDITION(d_pos, "d_pos must not be NULL");
  PRECONDITION(d_trajectory, "d_trajectory must not be NULL");
  PRECONDITION(d_trajectory->dimension() == 2, "d_dimension must be == 2");
  PRECONDITION(d_trajectory->numPoints(), "d_numPoints must be > 0");
  URANGE_CHECK(pointNum, d_trajectory->numPoints() - 1);
  unsigned int i = pointNum * d_trajectory->dimension();
  return RDGeom::Point2D(d_pos[i], d_pos[i + 1]);
}

RDGeom::Point3D Snapshot::getPoint3D(unsigned int pointNum) const {
  PRECONDITION(d_pos, "d_pos must not be NULL");
  PRECONDITION(d_trajectory, "d_trajectory must not be NULL");
  PRECONDITION(d_trajectory->dimension() >= 2, "d_dimension must be >= 2");
  PRECONDITION(d_trajectory->numPoints(), "d_numPoints must be > 0");
  URANGE_CHECK(pointNum, d_trajectory->numPoints() - 1);
  unsigned int i = pointNum * d_trajectory->dimension();
  return (RDGeom::Point3D(d_pos[i], d_pos[i + 1],
          (d_trajectory->dimension() == 3) ? d_pos[i + 2] : 0.0));
}

Trajectory::Trajectory(unsigned int dimension, unsigned int numPoints, SnapshotVect *snapshotVect) :
  d_dimension(dimension),
  d_numPoints(numPoints) {
  if (!snapshotVect)
    snapshotVect = new SnapshotVect;
  d_snapshotVect.reset(snapshotVect);
  for (SnapshotVect::iterator vectIt = d_snapshotVect->begin();
    vectIt != d_snapshotVect->end(); ++vectIt)
    vectIt->d_trajectory = this;
}

Trajectory::Trajectory(const Trajectory &other) :
  d_dimension(other.d_dimension),
  d_numPoints(other.d_numPoints),
  d_snapshotVect(new SnapshotVect) {
  for (SnapshotVect::const_iterator vectIt = other.d_snapshotVect->begin();
    vectIt != other.d_snapshotVect->end(); ++vectIt)
    addSnapshot(*vectIt);
}

unsigned int Trajectory::addSnapshot(const Snapshot &s) {
  return insertSnapshot(d_snapshotVect->size(), s);
}

const Snapshot &Trajectory::getSnapshot(unsigned int snapshotNum) const {
  URANGE_CHECK(snapshotNum, d_snapshotVect->size() - 1);
  return (*d_snapshotVect)[snapshotNum];
}

unsigned int Trajectory::insertSnapshot(unsigned int snapshotNum, Snapshot s) {
  URANGE_CHECK(snapshotNum, d_snapshotVect->size());
  s.d_trajectory = this;
  return (d_snapshotVect->insert(d_snapshotVect->begin() + snapshotNum,
          s) - d_snapshotVect->begin());
}

unsigned int Trajectory::removeSnapshot(unsigned int snapshotNum) {
  URANGE_CHECK(snapshotNum, d_snapshotVect->size() - 1);
  return (d_snapshotVect->erase(d_snapshotVect->begin() + snapshotNum) - d_snapshotVect->begin());
}

unsigned int Trajectory::addConformersToMol(ROMol &mol, int from, int to) {
  PRECONDITION(d_numPoints == mol.getNumAtoms(),
               "Number of atom mismatch between ROMol and Trajectory");
  PRECONDITION(from < static_cast<int>(size()), "from must be < size()");
  PRECONDITION(to < static_cast<int>(size()), "to must be < size()");
  if (from < 0)
    from = 0;
  if (to < 0)
    to = size() - 1;
  PRECONDITION(!size() || (from <= to), "from must be <= to");
  int n;
  unsigned int nConf;
  for (n = from, nConf = 0; size() && (n <= to); ++n, ++nConf) {
    Conformer *conf = new Conformer(mol.getNumAtoms());
    for (unsigned int i = 0; i < mol.getNumAtoms(); ++i)
      conf->setAtomPos(i, getSnapshot(n).getPoint3D(i));
    mol.addConformer(conf, true);
  }
  return nConf;
}

unsigned int readAmberTrajectory(const std::string &fName, Trajectory &traj) {
  PRECONDITION(traj.dimension() == 3, "The trajectory must have dimension == 3");
  std::ifstream inStream(fName.c_str());
  if (!inStream || inStream.bad()) {
    std::stringstream ss;
    ss << "Bad input file: " << fName;
    throw RDKit::BadFileException(ss.str());
  }
  std::string tempStr;
  // title
  std::getline(inStream, tempStr);
  // read coordinates
  unsigned int nCoords = traj.numPoints() * 3;
  unsigned int nSnapshots = 0;
  while (inStream.good() && !inStream.eof()) {
    boost::shared_array<double> c(new double[nCoords]());
    unsigned int i = 0;
    while (i < nCoords) {
      if (!(inStream >> c[i])) {
        if (!inStream.eof()) {
          std::stringstream ss;
          ss << "Error while reading file: " << fName;
          throw ValueErrorException(ss.str());
        }
        else if (i && (i < (nCoords - 1))) {
          std::stringstream ss;
          ss << "Premature end of file: " << fName;
          throw ValueErrorException(ss.str());
        }
        break;
      }
      ++i;
    }
    if (!inStream.eof()) {
      traj.addSnapshot(Snapshot(c));
      ++nSnapshots;
    }
  }
  return nSnapshots;
}

unsigned int readGromosTrajectory(const std::string &fName, Trajectory &traj) {
  PRECONDITION(traj.dimension() == 3, "The trajectory must have dimension == 3");
  std::ifstream inStream(fName.c_str());
  if (!inStream || inStream.bad()) {
    std::stringstream ss;
    ss << "Bad input file: " << fName;
    throw RDKit::BadFileException(ss.str());
  }
  std::string tempStr;
  unsigned int nCoords = traj.numPoints() * 3;
  unsigned int nSnapshots = 0;
  const static char *ignoredKeywordArray[] = {
    "TITLE",
    "TIMESTEP",
    "VELOCITYRED",
    "VELOCITY",
    "GENBOX",
    "BOX"
  };
  std::set<std::string> ignoredKeywordSet;
  for (unsigned int i = 0; i < (sizeof(ignoredKeywordArray) / sizeof(char *)); ++i)
    ignoredKeywordSet.insert(std::string(ignoredKeywordArray[i]));
  while (inStream.good() && !inStream.eof()) {
    std::getline(inStream, tempStr);
    if (inStream.bad() || inStream.eof())
      continue;
    if (ignoredKeywordSet.find(tempStr) != ignoredKeywordSet.end()) {
      // ignored block
      while (inStream.good() && !inStream.eof() && (tempStr != "END"))
        std::getline(inStream, tempStr);
    }
    else if ((tempStr == "POSITIONRED") || (tempStr == "POSITION")) {
      // these are the positions
      boost::shared_array<double> c(new double[nCoords]());
      unsigned int j = 0;
      for (unsigned int i = 0; i < traj.numPoints();) {
        std::getline(inStream, tempStr);
        if (inStream.bad() || inStream.eof() || (tempStr == "END"))
          throw ValueErrorException("Wrong number of coordinates");
        // ignore comments
        if (tempStr.find("#") != std::string::npos)
          continue;
        std::stringstream ls(tempStr);
        double x, y, z;
        if (!(ls >> x >> y >> z))
          throw ValueErrorException("Error while reading file");
        // store the coordinates (convert to Angstrom!)
        c[j++] = x * 10.0;
        c[j++] = y * 10.0;
        c[j++] = z * 10.0;
        ++i;
      }
      std::getline(inStream, tempStr); // the END line
      if (inStream.bad() || inStream.eof() || (tempStr != "END"))
        throw ValueErrorException("Wrong number of coordinates");
      traj.addSnapshot(Snapshot(c));
      ++nSnapshots;
    }
    else {
      std::string supportedBlocks("POSITIONRED, POSITION");
      for (std::set<std::string>::const_iterator it = ignoredKeywordSet.begin();
        it != ignoredKeywordSet.end(); ++it)
        supportedBlocks += ", " + *it;
      throw ValueErrorException("Unsupported block: "
        + tempStr + ". Supported blocks are " + supportedBlocks);
    }
  } // read file
  if (inStream.bad()) {
    std::stringstream ss;
		ss << "Bad input file: " << fName;
		throw RDKit::BadFileException(ss.str());
  }
  return nSnapshots;
}

}