from rdkit import Chem
from rdkit.Chem import rdFingerprintGenerator
import unittest


class TestCase(unittest.TestCase):

  def setUp(self):
    pass

  def testAtomPairGenerator(self):
    m = Chem.MolFromSmiles('CCC')
    g = rdFingerprintGenerator.GetAtomPairGenerator()
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 2)

    fp = g.GetCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 2)

    fp = g.GetSparseFingerprint(m)
    nzc = fp.GetNumOnBits()
    self.assertEqual(nzc, 3)

    fp = g.GetFingerprint(m)
    nzc = fp.GetNumOnBits()
    self.assertEqual(nzc, 3)

    g = rdFingerprintGenerator.GetAtomPairGenerator(
      atomInvariantsGenerator=rdFingerprintGenerator.GetAtomPairAtomInvGen())
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 2)

    g = rdFingerprintGenerator.GetAtomPairGenerator(minDistance=2)
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 1)

    g = rdFingerprintGenerator.GetAtomPairGenerator(maxDistance=1)
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 1)

    g = rdFingerprintGenerator.GetAtomPairGenerator(useCountSimulation=False)
    fp = g.GetSparseFingerprint(m)
    nzc = fp.GetNumOnBits()
    self.assertEqual(nzc, 2)

    invGen = rdFingerprintGenerator.GetAtomPairAtomInvGen(includeChirality=False)
    invGenChirality = rdFingerprintGenerator.GetAtomPairAtomInvGen(includeChirality=True)
    g = rdFingerprintGenerator.GetAtomPairGenerator(includeChirality=False,
                                                    atomInvariantsGenerator=invGen)
    gChirality = rdFingerprintGenerator.GetAtomPairGenerator(
      includeChirality=True, atomInvariantsGenerator=invGenChirality)
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    fpChirality = gChirality.GetSparseCountFingerprint(m)
    nzChirality = fpChirality.GetNonzeroElements()
    self.assertNotEqual(nz.keys(), nzChirality.keys())

  def testMorganGenerator(self):
    m = Chem.MolFromSmiles('CCCCC')
    g = rdFingerprintGenerator.GetMorganGenerator(3)
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 7)

  def testRDKitFPGenerator(self):
    m = Chem.MolFromSmiles('CCCCC')
    g = rdFingerprintGenerator.GetRDKitFPGenerator()
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 8)

  def testTopologicalTorsionGenerator(self):
    m = Chem.MolFromSmiles('CCCCC')
    g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 1)

  def testBulk(self):
    m1 = Chem.MolFromSmiles('CCC')
    m2 = Chem.MolFromSmiles('OCCCCC')
    m3 = Chem.MolFromSmiles('CCCCC')

    g = rdFingerprintGenerator.GetAtomPairGenerator()
    results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3],
                                                       rdFingerprintGenerator.AtomPairFP)
    self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetMorganGenerator(2)
    results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3],
                                                       rdFingerprintGenerator.MorganFP)
    self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetRDKitFPGenerator()
    results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3], rdFingerprintGenerator.RDKitFP)
    self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
    results = rdFingerprintGenerator.GetSparseCountFPs([m1, m2, m3],
                                                       rdFingerprintGenerator.TopologicalTorsionFP)
    self.assertEqual(results[0], g.GetSparseCountFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseCountFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseCountFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetAtomPairGenerator()
    results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3], rdFingerprintGenerator.AtomPairFP)
    self.assertEqual(results[0], g.GetSparseFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetMorganGenerator(2)
    results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3], rdFingerprintGenerator.MorganFP)
    self.assertEqual(results[0], g.GetSparseFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetRDKitFPGenerator()
    results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3], rdFingerprintGenerator.RDKitFP)
    self.assertEqual(results[0], g.GetSparseFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
    results = rdFingerprintGenerator.GetSparseFPs([m1, m2, m3],
                                                  rdFingerprintGenerator.TopologicalTorsionFP)
    self.assertEqual(results[0], g.GetSparseFingerprint(m1))
    self.assertEqual(results[1], g.GetSparseFingerprint(m2))
    self.assertEqual(results[2], g.GetSparseFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetAtomPairGenerator()
    results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3], rdFingerprintGenerator.AtomPairFP)
    self.assertEqual(results[0], g.GetCountFingerprint(m1))
    self.assertEqual(results[1], g.GetCountFingerprint(m2))
    self.assertEqual(results[2], g.GetCountFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetMorganGenerator(2)
    results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3], rdFingerprintGenerator.MorganFP)
    self.assertEqual(results[0], g.GetCountFingerprint(m1))
    self.assertEqual(results[1], g.GetCountFingerprint(m2))
    self.assertEqual(results[2], g.GetCountFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetRDKitFPGenerator()
    results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3], rdFingerprintGenerator.RDKitFP)
    self.assertEqual(results[0], g.GetCountFingerprint(m1))
    self.assertEqual(results[1], g.GetCountFingerprint(m2))
    self.assertEqual(results[2], g.GetCountFingerprint(m3))
    self.assertEqual(len(results), 3)

    g = rdFingerprintGenerator.GetTopologicalTorsionGenerator()
    results = rdFingerprintGenerator.GetCountFPs([m1, m2, m3],
                                                 rdFingerprintGenerator.TopologicalTorsionFP)
    self.assertEqual(results[0], g.GetCountFingerprint(m1))
    self.assertEqual(results[1], g.GetCountFingerprint(m2))
    self.assertEqual(results[2], g.GetCountFingerprint(m3))
    self.assertEqual(len(results), 3)

  def testNumBitsPerFeature(self):
    m1 = Chem.MolFromSmiles('CCCO')
    g = rdFingerprintGenerator.GetRDKitFPGenerator(minPath=1, maxPath=2)
    fp = g.GetFingerprint(m1)
    self.assertEqual(fp.GetNumOnBits(), 8)

    g = rdFingerprintGenerator.GetRDKitFPGenerator(minPath=1, maxPath=2, numBitsPerFeature=1)
    fp = g.GetFingerprint(m1)
    self.assertEqual(fp.GetNumOnBits(), 4)


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