//
//  Copyright (C) 2007-2023 Greg Landrum
//
//   @@ 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 <GraphMol/RDKitBase.h>
#include <GraphMol/Descriptors/MolDescriptors.h>
#include <GraphMol/PartialCharges/GasteigerCharges.h>
#include <vector>
#include <algorithm>
#include <boost/format.hpp>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

namespace RDKit {
namespace Descriptors {

double getLabuteAtomContribs(const ROMol &mol, std::vector<double> &Vi,
                             double &hContrib, bool includeHs, bool force) {
  TEST_ASSERT(Vi.size() == mol.getNumAtoms());
  if (!force && mol.hasProp(common_properties::_labuteAtomContribs)) {
    mol.getProp(common_properties::_labuteAtomContribs, Vi);
    mol.getProp(common_properties::_labuteAtomHContrib, hContrib);
    double res;
    mol.getProp(common_properties::_labuteASA, res);
    return res;
  }
  unsigned int nAtoms = mol.getNumAtoms();
  std::vector<double> rads(nAtoms);
  for (unsigned int i = 0; i < nAtoms; ++i) {
    rads[i] = PeriodicTable::getTable()->getRb0(
        mol.getAtomWithIdx(i)->getAtomicNum());
    Vi[i] = 0.0;
  }

  for (const auto bond : mol.bonds()) {
    const double bondScaleFacts[4] = {.1, 0, .2, .3};
    double Ri = rads[bond->getBeginAtomIdx()];
    double Rj = rads[bond->getEndAtomIdx()];
    double bij = Ri + Rj;
    if (!bond->getIsAromatic()) {
      if (bond->getBondType() < 4) {
        bij -= bondScaleFacts[bond->getBondType()];
      }
    } else {
      bij -= bondScaleFacts[0];
    }
    double dij = std::min(std::max(fabs(Ri - Rj), bij), Ri + Rj);
    Vi[bond->getBeginAtomIdx()] += Rj * Rj - (Ri - dij) * (Ri - dij) / dij;
    Vi[bond->getEndAtomIdx()] += Ri * Ri - (Rj - dij) * (Rj - dij) / dij;
  }
  hContrib = 0.0;
  if (includeHs) {
    double Rj = PeriodicTable::getTable()->getRb0(1);
    for (unsigned int i = 0; i < nAtoms; ++i) {
      double Ri = rads[i];
      double bij = Ri + Rj;
      double dij = std::min(std::max(fabs(Ri - Rj), bij), Ri + Rj);
      Vi[i] += Rj * Rj - (Ri - dij) * (Ri - dij) / dij;
      hContrib += Ri * Ri - (Rj - dij) * (Rj - dij) / dij;
    }
  }

  double res = 0.0;
  for (unsigned int i = 0; i < nAtoms; ++i) {
    double Ri = rads[i];
    Vi[i] = M_PI * Ri * (4. * Ri - Vi[i]);
    res += Vi[i];
  }
  if (includeHs && fabs(hContrib) > 1e-4) {
    double Rj = PeriodicTable::getTable()->getRb0(1);
    hContrib = M_PI * Rj * (4. * Rj - hContrib);
    res += hContrib;
  }
  mol.setProp(common_properties::_labuteAtomContribs, Vi, true);
  mol.setProp(common_properties::_labuteAtomHContrib, hContrib, true);
  mol.setProp(common_properties::_labuteASA, res, true);

  return res;
}
double calcLabuteASA(const ROMol &mol, bool includeHs, bool force) {
  if (!force && mol.hasProp(common_properties::_labuteASA)) {
    double res;
    mol.getProp(common_properties::_labuteASA, res);
    return res;
  }
  std::vector<double> contribs;
  contribs.resize(mol.getNumAtoms());
  double hContrib;
  double res;
  res = getLabuteAtomContribs(mol, contribs, hContrib, includeHs, force);
  return res;
}

double getTPSAAtomContribs(const ROMol &mol, std::vector<double> &Vi,
                           bool force, bool includeSandP) {
  TEST_ASSERT(Vi.size() >= mol.getNumAtoms());
  double res = 0;
  std::string pname =
      (boost::format("%s-%s") % common_properties::_tpsa % includeSandP).str();
  std::string contribsName =
      (boost::format("%s-%s") % common_properties::_tpsaAtomContribs %
       includeSandP)
          .str();
  if (!force && mol.hasProp(contribsName)) {
    mol.getProp(contribsName, Vi);
    mol.getProp(pname, res);
    return res;
  }
  unsigned int nAtoms = mol.getNumAtoms();
  std::vector<int> nNbrs(nAtoms, 0), nSing(nAtoms, 0), nDoub(nAtoms, 0),
      nTrip(nAtoms, 0), nArom(nAtoms, 0), nHs(nAtoms, 0);
  for (const auto bond : mol.bonds()) {
    if (bond->getBeginAtom()->getAtomicNum() == 1) {
      nNbrs[bond->getEndAtomIdx()] -= 1;
      nHs[bond->getEndAtomIdx()] += 1;
    } else if (bond->getEndAtom()->getAtomicNum() == 1) {
      nNbrs[bond->getBeginAtomIdx()] -= 1;
      nHs[bond->getBeginAtomIdx()] += 1;
    } else if (bond->getIsAromatic()) {
      nArom[bond->getBeginAtomIdx()] += 1;
      nArom[bond->getEndAtomIdx()] += 1;
    } else {
      switch (bond->getBondType()) {
        case Bond::SINGLE:
          nSing[bond->getBeginAtomIdx()] += 1;
          nSing[bond->getEndAtomIdx()] += 1;
          break;
        case Bond::DOUBLE:
          nDoub[bond->getBeginAtomIdx()] += 1;
          nDoub[bond->getEndAtomIdx()] += 1;
          break;
        case Bond::TRIPLE:
          nTrip[bond->getBeginAtomIdx()] += 1;
          nTrip[bond->getEndAtomIdx()] += 1;
          break;
        default:
          break;
      }
    }
  }

  for (unsigned int i = 0; i < nAtoms; ++i) {
    const Atom *atom = mol.getAtomWithIdx(i);
    int atNum = atom->getAtomicNum();

    if (atNum != 7 && atNum != 8 &&
        (!includeSandP || (atNum != 15 && atNum != 16))) {
      continue;
    }

    nHs[i] += atom->getTotalNumHs();
    int chg = atom->getFormalCharge();
    bool in3Ring = mol.getRingInfo()->isAtomInRingOfSize(i, 3);
    nNbrs[i] += atom->getDegree();

    double tmp = -1;
    if (atNum == 7) {
      switch (nNbrs[i]) {
        case 1:
          if (nHs[i] == 0 && chg == 0 && nTrip[i] == 1) {
            tmp = 23.79;
          } else if (nHs[i] == 1 && chg == 0 && nDoub[i] == 1) {
            tmp = 23.85;
          } else if (nHs[i] == 2 && chg == 0 && nSing[i] == 1) {
            tmp = 26.02;
          } else if (nHs[i] == 2 && chg == 1 && nDoub[i] == 1) {
            tmp = 25.59;
          } else if (nHs[i] == 3 && chg == 1 && nSing[i] == 1) {
            tmp = 27.64;
          }
          break;
        case 2:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 1 && nDoub[i] == 1) {
            tmp = 12.36;
          } else if (nHs[i] == 0 && chg == 0 && nTrip[i] == 1 &&
                     nDoub[i] == 1) {
            tmp = 13.60;
          } else if (nHs[i] == 1 && chg == 0 && nSing[i] == 2 && in3Ring) {
            tmp = 21.94;
          } else if (nHs[i] == 1 && chg == 0 && nSing[i] == 2 && !in3Ring) {
            tmp = 12.03;
          } else if (nHs[i] == 0 && chg == 1 && nTrip[i] == 1 &&
                     nSing[i] == 1) {
            tmp = 4.36;
          } else if (nHs[i] == 1 && chg == 1 && nDoub[i] == 1 &&
                     nSing[i] == 1) {
            tmp = 13.97;
          } else if (nHs[i] == 2 && chg == 1 && nSing[i] == 2) {
            tmp = 16.61;
          } else if (nHs[i] == 0 && chg == 0 && nArom[i] == 2) {
            tmp = 12.89;
          } else if (nHs[i] == 1 && chg == 0 && nArom[i] == 2) {
            tmp = 15.79;
          } else if (nHs[i] == 1 && chg == 1 && nArom[i] == 2) {
            tmp = 14.14;
          }
          break;
        case 3:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 3 && in3Ring) {
            tmp = 3.01;
          } else if (nHs[i] == 0 && chg == 0 && nSing[i] == 3 && !in3Ring) {
            tmp = 3.24;

          } else if (nHs[i] == 0 && chg == 0 && nSing[i] == 1 &&
                     nDoub[i] == 2) {
            tmp = 11.68;
          } else if (nHs[i] == 0 && chg == 1 && nSing[i] == 2 &&
                     nDoub[i] == 1) {
            tmp = 3.01;
          } else if (nHs[i] == 1 && chg == 1 && nSing[i] == 3) {
            tmp = 4.44;
          } else if (nHs[i] == 0 && chg == 0 && nArom[i] == 3) {
            tmp = 4.41;
          } else if (nHs[i] == 0 && chg == 0 && nSing[i] == 1 &&
                     nArom[i] == 2) {
            tmp = 4.93;
          } else if (nHs[i] == 0 && chg == 0 && nDoub[i] == 1 &&
                     nArom[i] == 2) {
            tmp = 8.39;
          } else if (nHs[i] == 0 && chg == 1 && nArom[i] == 3) {
            tmp = 4.10;
          } else if (nHs[i] == 0 && chg == 1 && nSing[i] == 1 &&
                     nArom[i] == 2) {
            tmp = 3.88;
          }
          break;
        case 4:
          if (nHs[i] == 0 && nSing[i] == 4 && chg == 1) {
            tmp = 0.0;
          }
          break;
        default:
          break;
      }
      if (tmp < 0.0) {
        tmp = 30.5 - nNbrs[i] * 8.2 + nHs[i] * 1.5;
        if (tmp < 0) {
          tmp = 0.0;
        }
      }
    } else if (atNum == 8) {
      switch (nNbrs[i]) {
        case 1:
          if (nHs[i] == 0 && chg == 0 && nDoub[i] == 1) {
            tmp = 17.07;
          } else if (nHs[i] == 1 && chg == 0 && nSing[i] == 1) {
            tmp = 20.23;
          } else if (nHs[i] == 0 && chg == -1 && nSing[i] == 1) {
            tmp = 23.06;
          }
          break;
        case 2:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 2 && in3Ring) {
            tmp = 12.53;
          } else if (nHs[i] == 0 && chg == 0 && nSing[i] == 2 && !in3Ring) {
            tmp = 9.23;
          } else if (nHs[i] == 0 && chg == 0 && nArom[i] == 2) {
            tmp = 13.14;
          }
          break;
        default:
          break;
      }
      if (tmp < 0.0) {
        tmp = 28.5 - nNbrs[i] * 8.6 + nHs[i] * 1.5;
        if (tmp < 0) {
          tmp = 0.0;
        }
      }
    } else if (includeSandP && atNum == 15) {
      tmp = 0.0;
      switch (nNbrs[i]) {
        case 2:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 1 && nDoub[i] == 1) {
            tmp = 34.14;
          }
          break;
        case 3:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 3) {
            tmp = 13.59;
          } else if (nHs[i] == 1 && chg == 0 && nSing[i] == 2 &&
                     nDoub[i] == 1) {
            tmp = 23.47;
          }
          break;
        case 4:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 3 && nDoub[i] == 1) {
            tmp = 9.81;
          }
          break;
        default:
          break;
      }
    } else if (includeSandP && atNum == 16) {
      tmp = 0.0;
      switch (nNbrs[i]) {
        case 1:
          if (nHs[i] == 0 && chg == 0 && nDoub[i] == 1) {
            tmp = 32.09;
          } else if (nHs[i] == 1 && chg == 0 && nSing[i] == 1) {
            tmp = 38.80;
          }
          break;
        case 2:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 2) {
            tmp = 25.30;
          } else if (nHs[i] == 0 && chg == 0 && nArom[i] == 2) {
            tmp = 28.24;
          }
          break;
        case 3:
          if (nHs[i] == 0 && chg == 0 && nArom[i] == 2 && nDoub[i] == 1) {
            tmp = 21.70;
          } else if (nHs[i] == 0 && chg == 0 && nSing[i] == 2 &&
                     nDoub[i] == 1) {
            tmp = 19.21;
          }
          break;
        case 4:
          if (nHs[i] == 0 && chg == 0 && nSing[i] == 2 && nDoub[i] == 2) {
            tmp = 8.38;
          }
          break;
        default:
          break;
      }
    }
    Vi[i] = tmp;
    res += tmp;
  }

  mol.setProp(contribsName, Vi, true);
  mol.setProp(pname, res, true);
  return res;
}
double calcTPSA(const ROMol &mol, bool force, bool includeSandP) {
  std::string pname =
      (boost::format("%s-%s") % common_properties::_tpsa % includeSandP).str();
  if (!force && mol.hasProp(pname)) {
    double res;
    mol.getProp(pname, res);
    return res;
  }
  std::vector<double> contribs;
  contribs.resize(mol.getNumAtoms());
  double res;
  res = getTPSAAtomContribs(mol, contribs, force, includeSandP);
  return res;
}

namespace {
void assignContribsToBins(const std::vector<double> &contribs,
                          const std::vector<double> &binProp,
                          const std::vector<double> &bins,
                          std::vector<double> &res) {
  PRECONDITION(contribs.size() == binProp.size(), "mismatched array sizes");
  PRECONDITION(res.size() >= bins.size() + 1, "mismatched array sizes");
  for (unsigned int i = 0; i < contribs.size(); ++i) {
    double cVal = contribs[i];
    double bVal = binProp[i];
    unsigned int idx =
        std::upper_bound(bins.begin(), bins.end(), bVal) - bins.begin();
    res[idx] += cVal;
  }
}
}  // namespace

std::vector<double> calcSlogP_VSA(const ROMol &mol, std::vector<double> *bins,
                                  bool force) {
  // FIX: use force value to include caching
  std::vector<double> lbins;
  if (!bins) {
    double blist[11] = {-0.4, -0.2, 0,   0.1, 0.15, 0.2,
                        0.25, 0.3,  0.4, 0.5, 0.6};
    lbins.resize(11);
    std::copy(blist, blist + 11, lbins.begin());
  } else {
    lbins.resize(bins->size());
    std::copy(bins->begin(), bins->end(), lbins.begin());
  }
  std::vector<double> res(lbins.size() + 1, 0);

  std::vector<double> vsaContribs(mol.getNumAtoms());
  double tmp;
  getLabuteAtomContribs(mol, vsaContribs, tmp, true, force);
  std::vector<double> logpContribs(mol.getNumAtoms());
  std::vector<double> mrContribs(mol.getNumAtoms());
  getCrippenAtomContribs(mol, logpContribs, mrContribs, force);

  assignContribsToBins(vsaContribs, logpContribs, lbins, res);

  return res;
}

std::vector<double> calcSMR_VSA(const ROMol &mol, std::vector<double> *bins,
                                bool force) {
  std::vector<double> lbins;
  if (!bins) {
    double blist[9] = {1.29, 1.82, 2.24, 2.45, 2.75, 3.05, 3.63, 3.8, 4.0};
    lbins.resize(9);
    std::copy(blist, blist + 9, lbins.begin());
  } else {
    lbins.resize(bins->size());
    std::copy(bins->begin(), bins->end(), lbins.begin());
  }
  std::vector<double> res(lbins.size() + 1, 0);

  std::vector<double> vsaContribs(mol.getNumAtoms());
  double tmp;
  getLabuteAtomContribs(mol, vsaContribs, tmp, true, force);
  std::vector<double> logpContribs(mol.getNumAtoms());
  std::vector<double> mrContribs(mol.getNumAtoms());
  getCrippenAtomContribs(mol, logpContribs, mrContribs, force);

  assignContribsToBins(vsaContribs, mrContribs, lbins, res);

  return res;
}

std::vector<double> calcPEOE_VSA(const ROMol &mol, std::vector<double> *bins,
                                 bool force) {
  std::vector<double> lbins;
  if (!bins) {
    double blist[13] = {-.3, -.25, -.20, -.15, -.10, -.05, 0,
                        .05, .10,  .15,  .20,  .25,  .30};
    lbins.resize(13);
    std::copy(blist, blist + 13, lbins.begin());
  } else {
    lbins.resize(bins->size());
    std::copy(bins->begin(), bins->end(), lbins.begin());
  }
  std::vector<double> res(lbins.size() + 1, 0);

  std::vector<double> vsaContribs(mol.getNumAtoms());
  double tmp;
  getLabuteAtomContribs(mol, vsaContribs, tmp, true, force);

  std::vector<double> chgs(mol.getNumAtoms(), 0.0);
  computeGasteigerCharges(mol, chgs);
  assignContribsToBins(vsaContribs, chgs, lbins, res);

  return res;
}

std::vector<double> calcCustomProp_VSA(const ROMol &mol,
                                       const std::string &customPropName,
                                       const std::vector<double> &bins,
                                       bool force) {
  std::vector<double> res(bins.size() + 1, 0);

  std::vector<double> vsaContribs(mol.getNumAtoms());
  double tmp;
  getLabuteAtomContribs(mol, vsaContribs, tmp, true, force);

  std::vector<double> prop(mol.getNumAtoms(), 0.0);
  for (unsigned int i = 0; i < mol.getNumAtoms(); ++i) {
    if (mol.getAtomWithIdx(i)->hasProp(customPropName)) {
      prop[i] = mol.getAtomWithIdx(i)->getProp<double>(customPropName);
    } else {
      prop[i] = 1;
    }
  }
  assignContribsToBins(vsaContribs, prop, bins, res);

  return res;
}
}  // end of namespace Descriptors
}  // namespace RDKit
