from rdkit import Chem
from rdkit.Chem import rdConformerParser
from rdkit import RDConfig
import unittest
import os
from rdkit.RDLogger import logger
logger = logger()


class TestCase(unittest.TestCase):

  def setUp(self):
    pass

  def testReadAmberTraj(self):
    fileN = os.path.join(RDConfig.RDBaseDir, 'Contrib', 'ConformerParser', 'test_data',
                         'water_coords.trx')
    mol = Chem.MolFromSmiles('O')
    mol = Chem.AddHs(mol)
    ids = rdConformerParser.AddConformersFromAmberTrajectory(mol, fileN)
    self.failUnless(mol.GetNumConformers() == 1)
    self.failUnless(len(ids) == 1)
    self.failUnless(ids[0] == 0)

    fileN = os.path.join(RDConfig.RDBaseDir, 'Contrib', 'ConformerParser', 'test_data',
                         'water_coords2.trx')
    ids = rdConformerParser.AddConformersFromAmberTrajectory(mol, fileN, clearConfs=True)
    self.failUnless(mol.GetNumConformers() == 2)
    ids = rdConformerParser.AddConformersFromAmberTrajectory(mol, fileN, clearConfs=False)
    self.failUnless(mol.GetNumConformers() == 4)
    ids = rdConformerParser.AddConformersFromAmberTrajectory(mol, fileN, numConfs=1,
                                                             clearConfs=True)
    self.failUnless(mol.GetNumConformers() == 1)


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