from __future__ import print_function
import unittest
import os
import io

from rdkit.six.moves import cPickle as pickle

from rdkit import Chem
from rdkit.Chem import rdPartialCharges
from rdkit import RDConfig


def feq(v1, v2, tol2=1e-4):
  return abs(v1 - v2) <= tol2


class TestCase(unittest.TestCase):

  def setUp(self):
    pass

  def test0HalgrenSet(self):
    smiSup = Chem.SmilesMolSupplier(
      os.path.join(RDConfig.RDBaseDir, 'Code', 'GraphMol', 'PartialCharges', 'Wrap', 'test_data',
                   'halgren.smi'), delimiter='\t')

    #parse the original file
    with open(
        os.path.join(RDConfig.RDBaseDir, 'Code', 'GraphMol', 'PartialCharges', 'Wrap', 'test_data',
                     'halgren_out.txt'), 'r') as infil:
      lines = infil.readlines()

    tab = Chem.GetPeriodicTable()

    olst = []
    for mol in smiSup:
      rdPartialCharges.ComputeGasteigerCharges(mol)
      tstr = "Molecule: "
      tstr += mol.GetProp("_Name")
      olst.append(tstr)
      for i in range(mol.GetNumAtoms()):
        at = mol.GetAtomWithIdx(i)
        en = tab.GetElementSymbol(at.GetAtomicNum())
        chg = float(at.GetProp("_GasteigerCharge"))
        tstr = "%i %s %6.4f" % (i, en, chg)
        olst.append(tstr)

    i = 0
    for line in lines:
      self.assertTrue(line.strip() == olst[i])
      i += 1

  def test1PPDataset(self):
    fileN = os.path.join(RDConfig.RDBaseDir, 'Code', 'GraphMol', 'PartialCharges', 'Wrap',
                         'test_data', 'PP_descrs_regress.2.csv')
    infil = open(fileN, 'r')
    lines = infil.readlines()
    infil.close()

    infile = os.path.join(RDConfig.RDBaseDir, 'Code', 'GraphMol', 'PartialCharges', 'Wrap',
                          'test_data', 'PP_combi_charges.pkl')
    with open(infile, 'r') as cchtFile:
      buf = cchtFile.read().replace('\r\n', '\n').encode('utf-8')
      cchtFile.close()
    with io.BytesIO(buf) as cchFile:
      combiCharges = pickle.load(cchFile)

    for lin in lines:
      if (lin[0] == '#'):
        continue
      tlst = lin.strip().split(',')
      smi = tlst[0]
      rdmol = Chem.MolFromSmiles(smi)
      rdPartialCharges.ComputeGasteigerCharges(rdmol)

      nat = rdmol.GetNumAtoms()
      failed = False
      for ai in range(nat):
        rdch = float(rdmol.GetAtomWithIdx(ai).GetProp('_GasteigerCharge'))
        if not feq(rdch, combiCharges[smi][ai], 1.e-2):
          failed = True
          print(smi, ai, rdch, combiCharges[smi][ai])
      if failed:
        rdmol.Debug()
      self.assertFalse(failed)

  def test2Params(self):
    """ tests handling of Issue187 """
    m1 = Chem.MolFromSmiles('C(=O)[O-]')
    rdPartialCharges.ComputeGasteigerCharges(m1)

    m2 = Chem.MolFromSmiles('C(=O)[O-].[Na+]')
    rdPartialCharges.ComputeGasteigerCharges(m2)

    for i in range(m1.GetNumAtoms()):
      c1 = float(m1.GetAtomWithIdx(i).GetProp('_GasteigerCharge'))
      c2 = float(m2.GetAtomWithIdx(i).GetProp('_GasteigerCharge'))
      self.assertTrue(feq(c1, c2, 1e-4))

  def test3Params(self):
    """ tests handling of Issue187 """
    m2 = Chem.MolFromSmiles('C(=O)[O-].[Na+]')
    with self.assertRaisesRegexp(Exception, ""):
      rdPartialCharges.ComputeGasteigerCharges(m2, 12, 1)

  def testGithubIssue20(self):
    """ tests handling of Github issue 20 """
    m1 = Chem.MolFromSmiles('CB(O)O')
    rdPartialCharges.ComputeGasteigerCharges(m1)
    chgs = [-0.030, 0.448, -0.427, -0.427]
    for i in range(m1.GetNumAtoms()):
      c1 = float(m1.GetAtomWithIdx(i).GetProp('_GasteigerCharge'))
      self.assertAlmostEqual(c1, chgs[i], 3)

  def testGithubIssue577(self):
    """ tests handling of Github issue 577 """
    m1 = Chem.MolFromSmiles('CCO')
    from locale import setlocale, LC_NUMERIC
    try:
      setlocale(LC_NUMERIC, "de_DE")
    except Exception:
      # can't set the required locale, might as well just return
      return
    try:
      rdPartialCharges.ComputeGasteigerCharges(m1)
      for at in m1.GetAtoms():
        float(at.GetProp('_GasteigerCharge'))
    finally:
      setlocale(LC_NUMERIC, "C")
    rdPartialCharges.ComputeGasteigerCharges(m1)
    for at in m1.GetAtoms():
      float(at.GetProp('_GasteigerCharge'))


if __name__ == '__main__':
  unittest.main()