// $Id$
//
//  Copyright (C) 2003-2011 Greg Landrum and Rational Discovery LLC
//
//   @@ 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 "GasteigerCharges.h"
#include <RDGeneral/types.h>
#include <GraphMol/RDKitBase.h>
#include <GraphMol/ROMol.h>
#include <GraphMol/MolOps.h>
#include "GasteigerParams.h"

namespace Gasteiger {
using namespace RDKit;
/*! \brief split the formal charge across atoms of same type if we have a
 *conjugated system
 *
 *  This function is called before the charge equivalization iteration is
 *started for the
 *  Gasteiger charges. If any of the atom involved in conjugated system have
 *formal charges
 *  set on them, this charges is equally distributed across atoms of the same
 *type in that
 *  conjugated system. So for example the two nitrogens in the benzamidine
 *system start the iteration
 * with equal charges of 0.5
 */
void splitChargeConjugated(const ROMol &mol, DOUBLE_VECT &charges) {
  int aix;
  int natms = mol.getNumAtoms();
  INT_VECT marker;
  INT_VECT_CI mci;
  int aax, yax;
  double formal;
  const Atom *at, *aat, *yat;
  for (aix = 0; aix < natms; aix++) {
    at = mol.getAtomWithIdx(aix);
    formal = at->getFormalCharge();
    // std::cout << aix << " formal charges:" << formal << "\n";
    marker.resize(0);
    if ((fabs(formal) > EPS_DOUBLE) && (fabs(charges[aix]) < EPS_DOUBLE)) {
      marker.push_back(aix);
      ROMol::OEDGE_ITER bnd1, end1, bnd2, end2;
      boost::tie(bnd1, end1) = mol.getAtomBonds(at);
      while (bnd1 != end1) {
        if (mol[*bnd1]->getIsConjugated()) {
          aax = mol[*bnd1]->getOtherAtomIdx(aix);
          aat = mol.getAtomWithIdx(aax);
          boost::tie(bnd2, end2) = mol.getAtomBonds(aat);
          while (bnd2 != end2) {
            if ((*bnd1) != (*bnd2)) {
              if (mol[*bnd2]->getIsConjugated()) {
                yax = mol[*bnd2]->getOtherAtomIdx(aax);
                yat = mol.getAtomWithIdx(yax);
                if (at->getAtomicNum() == yat->getAtomicNum()) {
                  formal += yat->getFormalCharge();
                  marker.push_back(yax);
                }
              }
            }
            bnd2++;
          }
        }
        bnd1++;
      }

      for (mci = marker.begin(); mci != marker.end(); mci++) {
        charges[*mci] = (formal / marker.size());
      }
    }
  }
  /*
  for (aix = 0; aix < natms; aix++) {
    std::cout << "In splitter: " << " charges:" << charges[aix] << "\n";
    }*/
}
}  // end of namespace Gasteiger

namespace RDKit {
void computeGasteigerCharges(const ROMol *mol, int nIter,
                             bool throwOnParamFailure) {
  PRECONDITION(mol, "bad molecule");
  computeGasteigerCharges(*mol, nIter, throwOnParamFailure);
}
void computeGasteigerCharges(const ROMol &mol, int nIter,
                             bool throwOnParamFailure) {
  std::vector<double> chgs(mol.getNumAtoms());
  computeGasteigerCharges(mol, chgs, nIter, throwOnParamFailure);
}

/*! \brief compute the Gasteiger partial charges and return a new molecule with
 *the charges set
 *
 * Ref : J.Gasteiger, M. Marsili, "Iterative Equalization of Oribital
 *Electronegatiity
 *  A Rapid Access to Atomic Charges", Tetrahedron Vol 36 p3219 1980
 */
void computeGasteigerCharges(const ROMol &mol, std::vector<double> &charges,
                             int nIter, bool throwOnParamFailure) {
  PRECONDITION(charges.size() >= mol.getNumAtoms(), "bad array size");

  PeriodicTable *table = PeriodicTable::getTable();
  const GasteigerParams *params = GasteigerParams::getParams();

  double damp = DAMP;
  int natms = mol.getNumAtoms();
  // space for parameters for each atom in the molecule
  std::vector<DOUBLE_VECT> atmPs;
  atmPs.reserve(natms);

  std::fill(charges.begin(), charges.end(), 0.0);

  DOUBLE_VECT
  hChrg;  // total charge on the implicit hydrogen on each heavy atom
  hChrg.resize(natms, 0.0);

  DOUBLE_VECT ionX;
  ionX.resize(natms, 0.0);

  DOUBLE_VECT energ;
  energ.resize(natms, 0.0);

  ROMol::ADJ_ITER nbrIdx, endIdx;

  // deal with the conjugated system - distribute the formal charges on atoms of
  // same type in each
  // conjugated system
  Gasteiger::splitChargeConjugated(mol, charges);

  // now read in the parameters
  ROMol::ConstAtomIterator ai;

  for (ai = mol.beginAtoms(); ai != mol.endAtoms(); ai++) {
    std::string elem = table->getElementSymbol((*ai)->getAtomicNum());
    std::string mode;

    switch ((*ai)->getHybridization()) {
      case Atom::SP3:
        mode = "sp3";
        break;
      case Atom::SP2:
        mode = "sp2";
        break;
      case Atom::SP:
        mode = "sp";
        break;
      default:
        if ((*ai)->getAtomicNum() == 1) {
          // if it is hydrogen
          mode = "*";
        } else if ((*ai)->getAtomicNum() == 16) {
          // we have a sulfur atom with no hydribidation information
          // check how many oxygens we have on the sulfer
          boost::tie(nbrIdx, endIdx) = mol.getAtomNeighbors(*ai);
          int no = 0;
          while (nbrIdx != endIdx) {
            if (mol.getAtomWithIdx(*nbrIdx)->getAtomicNum() == 8) {
              no++;
            }
            nbrIdx++;
          }
          if (no == 2) {
            mode = "so2";
          } else if (no == 1) {
            mode = "so";
          } else {
            // some other sulfur state. Default to sp3
            mode = "sp3";
          }
        }
    }

    // if we get a unknown mode or element type the
    // following will will throw an exception
    atmPs.push_back(params->getParams(elem, mode, throwOnParamFailure));

    // set ionX paramters
    // if Hydrogen treat differently
    int idx = (*ai)->getIdx();
    if ((*ai)->getAtomicNum() == 1) {
      ionX[idx] = IONXH;
    } else {
      ionX[idx] = atmPs[idx][0] + atmPs[idx][1] + atmPs[idx][2];
    }
  }

  // do the iteration here
  int itx, aix, sgn, niHs;
  double enr, dq, dx, qHs, dqH;
  // parameters for hydrogen atoms (for case where the hydrogen are not in the
  // graph (implicit hydrogens)
  DOUBLE_VECT hParams;
  hParams = params->getParams("H", "*", throwOnParamFailure);

  /*
  int itmp;
  for (itmp = 0; itmp < 5; itmp++) {
    std::cout << " aq:" << charges[itmp] << "\n";
    }*/

  for (itx = 0; itx < nIter; itx++) {
    for (aix = 0; aix < natms; aix++) {
      // enr = p0 + charge*(p1 + p2*charge)
      enr = atmPs[aix][0] +
            charges[aix] * (atmPs[aix][1] + atmPs[aix][2] * charges[aix]);
      energ[aix] = enr;
    }

    for (aix = 0; aix < natms; aix++) {
      dq = 0.0;
      boost::tie(nbrIdx, endIdx) =
          mol.getAtomNeighbors(mol.getAtomWithIdx(aix));
      while (nbrIdx != endIdx) {
        dx = energ[*nbrIdx] - energ[aix];
        if (dx < 0.0) {
          sgn = 0;
        } else {
          sgn = 1;
        }
        dq += dx / ((sgn * (ionX[aix] - ionX[*nbrIdx])) + ionX[*nbrIdx]);
        nbrIdx++;
      }
      // now loop over the implicit hydrogens and get their contributions
      // since hydrogens don't connect to anything else, update their charges at
      // the same time
      niHs = mol.getAtomWithIdx(aix)->getTotalNumHs();
      if (niHs > 0) {
        qHs = hChrg[aix] / niHs;
        enr = hParams[0] + qHs * (hParams[1] + hParams[2] * qHs);
        dx = enr - energ[aix];
        if (dx < 0.0) {
          sgn = 0;
        } else {
          sgn = 1;
        }

        dqH = dx / ((sgn * (ionX[aix] - IONXH)) + IONXH);

        dq += (niHs * dqH);

        // adjust the charges on the hydrogens simultaneously (possible because
        // each of the
        // hydrogens have no other neighbors)
        hChrg[aix] -= (niHs * dqH * damp);
      }
      charges[aix] += (damp * dq);
    }

    damp *= DAMP_SCALE;
  }

  for (aix = 0; aix < natms; aix++) {
    mol.getAtomWithIdx(aix)->setProp(common_properties::_GasteigerCharge,
                                     charges[aix], true);
    // set the implicit hydrogen charges
    mol.getAtomWithIdx(aix)->setProp(common_properties::_GasteigerHCharge,
                                     hChrg[aix], true);
  }
}
}