/**********************************************************************

Copyright (C) 2010 by Reinis Danne

Some portions Copyright (C) 2004 by Chris Morley

This program is free software; you can redistribute it and/or modify

it under the terms of the GNU General Public License as published by

the Free Software Foundation version 2 of the License.

This program is distributed in the hope that it will be useful,

but WITHOUT ANY WARRANTY; without even the implied warranty of

MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

GNU General Public License for more details.

***********************************************************************/

#include <openbabel/babelconfig.h>

#include <openbabel/obmolecformat.h>

#include <openbabel/mol.h>

#include <openbabel/atom.h>

#include <openbabel/elements.h>

#include <openbabel/generic.h>

#include <openbabel/obiter.h>

#include <sstream>

#include <iomanip>

using namespace std;

namespace OpenBabel

{

class GROFormat : public OBMoleculeFormat

{

public:

  //Register this format type ID in the constructor

  GROFormat()

  {

    /* GRO is the file extension and is case insensitive. A MIME type can be

       added as an optional third parameter.

       Multiple file extensions can be registered by adding extra statements.*/

    OBConversion::RegisterFormat("gro",this);

    /* If there are any format specific options they should be registered here

       so that the commandline interface works properly.

       The first parameter is the option name. If it is a single letter it can

       be concatinated with other single letter options. For output options it

       can be multicharcter and is then written as --optionname on the command

       line. The third parameter is the number of parameters the option takes.

       Currently this is either 1 or 0 and if it is 0 can be omitted for output

       options. The parameter is always text and needs to be parsed to extract

       a number.

       Options can apply when writing - 4th parameter is

       OBConversion::OUTOPTIONS or can be omitted as shown. A single letter

       output option is preceded by -x on the command line.

       Or options can apply to the input format - the 4th parameter is

       OBConversion::INOPTIONS. They are then preceded by -a on the command

       line.

       Each option letter may be reused in other formats, but within the same

       group, INOPTIONS or OUTOPTIONS, must take the same number of parameters

       (0 or 1). There will be an error message when OpenBabel  runs if there

       are conflicts between formats. A list of formats currently used (which

       may not be comprehensive) is in docs/options.html.

    */

    OBConversion::RegisterOptionParam("f", this, 1);

    OBConversion::RegisterOptionParam("n", this);

    OBConversion::RegisterOptionParam("s", this, 0, OBConversion::INOPTIONS);

  }

  /* The first line of the description should be a brief identifier, <40 chars,

     because it is used in dropdown lists, etc. in some user interfaces.

     The rest is optional.

     Describe any format specific options here. This text is parsed to provide

     checkboxes, etc for the GUI (for details click the control menu),

     so please try to keep to a similar form.

     Write options are the most common, and the "Write" is optional.

     The option f takes a text parameter, so that it is essential that the

     option is registered in the constructor of the class.

     Finish the options with a blank line as shown, if there are more than one

     group of options, or if there are further comments after them.

  */

  const char* Description() override  // required

  {

    return

    "GRO format\n"

    "This is GRO file format as used in Gromacs.\n"

    "Right now there is only limited support for element perception. It works "

    "for \nelements with one letter symbols if the atomtype starts with the "

    "same letter.\n\n"

    "Read Options e.g. -as\n"

    " s  Consider single bonds only\n"

    " b  Disable bonding entierly\n"

    ;

  }

  //Optional URL where the file format is specified

  const char* SpecificationURL() override

  {

    return "https://manual.gromacs.org/documentation/current/reference-manual/file-formats.html#gro";

  }

  /* This optional function is for formats which can contain more than one

     molecule. It is used to quickly position the input stream after the nth

     molecule without requiring to convert and discard all the n molecules.

     See obconversion.cpp for details.*/

  int SkipObjects(int n, OBConversion* pConv) override

  {

    string line = "";

    int natoms = 0;

    int nlines = 0;

    if (n == 0)

      n++;

    istream& ifs = *pConv->GetInStream();

    getline(ifs, line);

    ifs >> natoms;

    if (natoms <= 0 || natoms >= 100000000)

      return -1;

    nlines = (natoms+3)*n;

    while (ifs && --nlines) {

      getline(ifs, line);

    }

    return ifs.good() ? 1 : -1;

  }

  ////////////////////////////////////////////////////

  /// Declarations for the "API" interface functions. Definitions are below

  bool ReadMolecule(OBBase* pOb, OBConversion* pConv) override;

  bool WriteMolecule(OBBase* pOb, OBConversion* pConv) override;

};

  ////////////////////////////////////////////////////

//Make an instance of the format class

GROFormat theGROFormat;

/////////////////////////////////////////////////////////////////

bool GROFormat::ReadMolecule(OBBase* pOb, OBConversion* pConv)

{

  OBMol* pmol = pOb->CastAndClear<OBMol>();

  if (pmol == nullptr)

      return false;

  istream& ifs = *pConv->GetInStream();

  pmol->BeginModify();

  /** Parse the input stream and use the OpenBabel API to populate the OBMol **/

  char buffer[BUFF_SIZE];

  string line = "";

  stringstream errorMsg;

  int natoms = 0;

  string title = "";

  long int resid = 0; // 5

  string resname = ""; //5

  string atomtype = ""; //5

  //long int atomid = 0; //5

  double x = 0.0; // 8.3

  double y = 0.0; // 8.3

  double z = 0.0; // 8.3

  double vx = 0.0; // 8.4

  double vy = 0.0; // 8.4

  double vz = 0.0; // 8.4

  string tempstr = "";

  long int residx = 0;

  OBAtom* atom;

  OBResidue* res = nullptr;

  OBVectorData* velocity;

  if (!ifs || ifs.peek() == EOF) {

    return false; // Trying to read past end of the file

  }

  if (!ifs.getline(buffer, BUFF_SIZE)) {

    errorMsg << "Problems reading a GRO file: "

             << "Cannot read the first line!";

    obErrorLog.ThrowError(__FUNCTION__, errorMsg.str(), obWarning);

    return false;

  }

  // Get the title

  title.assign(buffer);

  if (title.size() < 1) {

    title = pConv->GetTitle();

    pmol->SetTitle(title);

  } else {

    pmol->SetTitle(title);

  }

  if (!ifs.getline(buffer, BUFF_SIZE)) {

    errorMsg << "Problems reading a GRO file: "

             << "Cannot read the second line!";

    obErrorLog.ThrowError(__FUNCTION__, errorMsg.str(), obWarning);

    return false;

  }

  // Get the number of atoms

  stringstream(buffer) >> natoms;

  if (natoms < 1 || natoms >= 100000000) {

    errorMsg << "Problems reading a GRO file: "

             << "There are no atoms in the file or the second line is"

             << " incorrectly written.";

    obErrorLog.ThrowError(__FUNCTION__, errorMsg.str(), obWarning);

    return false;

  }

  pmol->ReserveAtoms(natoms);

  // Read all atom records

  for (int i=1; i<=natoms; i++) {

    if (!ifs.getline(buffer,BUFF_SIZE)) {

      errorMsg << "Problems reading a GRO file: "

               << "Could not read line #" << i+2 << ", file error." << endl

               << " According to the second line, there should be " << natoms

               << " atoms, and therefore " << natoms+3 << " lines in the file.";

      obErrorLog.ThrowError(__FUNCTION__, errorMsg.str(), obWarning);

      return false;

    }

    line = buffer;

    // Get atom

    atom  = pmol->NewAtom();

    tempstr.assign(line,0,5);

    stringstream(tempstr) >> resid;

    resname.assign(line,5,5);

    Trim(resname);

    atomtype.assign(line,10,5);

    Trim(atomtype);

    // Not used, OB assigns its own indizes

    //tempstr.assign(line,15,5);

    //stringstream(tempstr) >> atomid;

    tempstr.assign(line,20,8);

    stringstream(tempstr) >> x;

    tempstr.assign(line,28,8);

    stringstream(tempstr) >> y;

    tempstr.assign(line,36,8);

    stringstream(tempstr) >> z;

    if (line.size() > 44) {

      tempstr.assign(line,44,8);

      stringstream(tempstr) >> vx;

      tempstr.assign(line,52,8);

      stringstream(tempstr) >> vy;

      tempstr.assign(line,60,8);

      stringstream(tempstr) >> vz;

      velocity = new OBVectorData();

      velocity->SetData(vx, vy, vz);

      velocity->SetAttribute("Velocity");

      velocity->SetOrigin(fileformatInput);

      atom->SetData(velocity);

    }

    // OB translates this and, e.g., OW1 turns into O3

    // Type conversion should be done explicitly if that needs to be

    // controlled.

    atom->SetType(atomtype);

    // Set coordinates of the atom, multiply by 10 to convert from nm to

    // angstrom

    atom->SetVector(x*10, y*10, z*10);

    if (resid == residx) {

      // Add atom to an existing residue

      if (res == nullptr) {

        res = pmol->NewResidue();

        res->SetName(resname);

        res->SetNum(resid);

      }

      res->AddAtom(atom);

    } else {

      // Create new residue and use that

      res = pmol->NewResidue();

      res->SetName(resname);

      res->SetNum(resid);

      res->AddAtom(atom);

      residx = resid;

    }

    // Atom type has to be set in residues as AtomID

    res->SetAtomID(atom, atomtype);

    // For determining the element this doesn't work:

    // atom->SetAtomicNum(OBElements::GetAtomicNum(atom->GetType()));

    //

    // So the element simbol should be found while reading in the file. It

    // could be possible to provide an option for external atomtype<->element

    // simbol map. Now such functionality is not implemented.

    // TODO: Make element perception optional and improve it

    string element = "";

    if (atomtype[0] == 'C') {

      if (atomtype.find_first_of("aloru",1) == 1) {

        element.assign(atomtype,0,2);

      } else {

        element.assign(atomtype,0,1);

      }

    } else if (atomtype[0] == 'N') {

      if (atomtype.find_first_of("abei",1) == 1) {

        element.assign(atomtype,0,2);

      } else {

        element.assign(atomtype,0,1);

      }

    } else {

      element.assign(atomtype,0,1);

    }

    atom->SetAtomicNum(OBElements::GetAtomicNum(element.data()));

  }

  // Get periodic box

  if (!ifs.getline(buffer,BUFF_SIZE)) {

    errorMsg << "Problems reading a GRO file: "

             << "Could not read box vectors!" << endl;

    obErrorLog.ThrowError(__FUNCTION__, errorMsg.str(), obWarning);

    return false;

  }

  double v1x=0.0, v2y=0.0, v3z=0.0;

  double v1y=0.0, v1z=0.0, v2x=0.0;

  double v2z=0.0, v3x=0.0, v3y=0.0;

  stringstream ss(buffer);

  if (!ss) {

    errorMsg << "Problems reading a GRO file: "

             << "Could not read box vectors!" << endl;

    obErrorLog.ThrowError(__FUNCTION__, errorMsg.str(), obWarning);

    return false;

  }

  ss >> v1x >> v2y >> v3z;

  if (ss) {

    ss >> v1y >> v1z >> v2x;

    ss >> v2z >> v3x >> v3y;

  }

  if (!(v1x == 0.0 && v2y == 0.0 && v3z == 0.0 &&

        v1y == 0.0 && v1z == 0.0 && v2x == 0.0 &&

        v2z == 0.0 && v3x == 0.0 && v3y == 0.0)) {

  // Set box vectors and convert nm to angstroms

  const vector3 v1(v1x*10,v1y*10,v1z*10);

  const vector3 v2(v2x*10,v2y*10,v2z*10);

  const vector3 v3(v3x*10,v3y*10,v3z*10);

  OBUnitCell* cell = new OBUnitCell();

  cell->SetData(v1, v2, v3);

  pmol->SetData(cell);

  }

  pmol->EndModify();

  // Input options -as and -ab

  if (!pConv->IsOption("b", OBConversion::INOPTIONS)) {

    pmol->ConnectTheDots();

    if (!pConv->IsOption("s", OBConversion::INOPTIONS)) {

      pmol->PerceiveBondOrders();

    }

  }

  pmol->SetChainsPerceived();

  /* For multi-molecule formats, leave the input stream at the start of the

     next molecule, ready for this routine to be called again.

   * Return true if ok. Returning false means discard the OBMol and stop

     converting, unless the -e option is set. With a multi-molecule inputstream

     this will skip the current molecule and continue with the next, if

     SkipObjects() has been defined. If it has not, and continuation after

     errors is still required, it is necessary to leave the input stream at the

     beginning of next object when returning false;*/

  return true;

}

////////////////////////////////////////////////////////////////

bool GROFormat::WriteMolecule(OBBase* pOb, OBConversion* pConv)

{

  OBMol* pmol = dynamic_cast<OBMol*>(pOb);

  if (pmol == nullptr)

      return false; // Stop converting

  ostream& ofs = *pConv->GetOutStream();

  /** Write the representation of the OBMol molecule to the output stream **/

  OBVectorData* vector;

  vector3 v;

  OBResidue* res = nullptr;

  string tempstr = "";

  long int atIdx = 0;

  long int resIdx = 0;

  ofs << pmol->GetTitle() << endl;

  ofs << pmol->NumAtoms() << endl;

  ofs.setf(ios::fixed);

  FOR_ATOMS_OF_MOL(atom, pmol) {

    // CVE-2022-42885: atoms read from formats without residue

    // information return a null OBResidue. Synthesize defaults

    // rather than dereferencing the null pointer.

    res = atom->GetResidue();

    resIdx = (res != nullptr) ? res->GetNum() : 1;

    // Check if residue index excedes the field width and should be wrapped

    if (resIdx > 99999) {

      ofs << setw(5) << resIdx - 100000*int(resIdx/100000);

    } else {

      ofs << setw(5) << resIdx;

    }

    ofs << setw(5) << left

        << ((res != nullptr) ? res->GetName() : string("UNK"));

    // Remove whitespace from AtomID left by other formats

    tempstr = (res != nullptr) ? res->GetAtomID(&(*atom))

                               : string(atom->GetType());

    ofs << setw(5) << right << Trim(tempstr);

    atIdx = atom->GetIdx();

    // Check if atom index excedes the field width and should be wrapped

    if (atIdx > 99999) {

      ofs << setw(5) << atIdx - 100000*int(atIdx/100000);

    } else {

      ofs << setw(5) << atIdx;

    }

    ofs.precision(3);

    // OpenBabel uses angstrom, so converting to nm by dividing by 10

    ofs << setw(8) << atom->x()/10

        << setw(8) << atom->y()/10

        << setw(8) << atom->z()/10;

    if (atom->GetData("Velocity")) {

      vector = (OBVectorData*) atom->GetData("Velocity");

      v = vector->GetData();

      ofs.precision(4);

      ofs << setw(8) << v.x()

          << setw(8) << v.y()

          << setw(8) << v.z();

    }

    ofs << endl;

  }

  // On the last line of the file goes periodic box specification

  if (pmol->HasData(OBGenericDataType::UnitCell)) {

    OBUnitCell* cell = (OBUnitCell*) pmol->GetData(OBGenericDataType::UnitCell);

    matrix3x3 m = cell->GetCellMatrix();

    vector3 v1 = m.GetRow(0);

    vector3 v2 = m.GetRow(1);

    vector3 v3 = m.GetRow(2);

    // Gromacs itself uses precision of 5, so it should be fine

    ofs.precision(5);

    ofs << "   " << v1.x()/10

        << "   " << v2.y()/10

        << "   " << v3.z()/10;

    // If there is any non-zero value among others, then write them all

    const double TRESHOLD = 1.0e-8;

    if (fabs(v1.y()) > TRESHOLD || fabs(v1.z()) > TRESHOLD ||

        fabs(v2.x()) > TRESHOLD || fabs(v2.z()) > TRESHOLD ||

        fabs(v3.x()) > TRESHOLD || fabs(v3.y()) > TRESHOLD) {

      ofs << "   " << v1.y()/10

          << "   " << v1.z()/10

          << "   " << v2.x()/10

          << "   " << v2.z()/10

          << "   " << v3.x()/10

          << "   " << v3.y()/10;

    }

  } else {

    // Set to zero if there is no box data in the molecule

    ofs << "   0.00000   0.00000   0.00000";

  }

  ofs << endl;

  return true;

}

} //namespace OpenBabel