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('CCCC(=O)O')
    g = rdFingerprintGenerator.GetMorganGenerator(3)
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 14)

    invgen = rdFingerprintGenerator.GetMorganAtomInvGen()
    g = rdFingerprintGenerator.GetMorganGenerator(radius=3, atomInvariantsGenerator=invgen)
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 14)

    invgen = rdFingerprintGenerator.GetMorganFeatureAtomInvGen()
    g = rdFingerprintGenerator.GetMorganGenerator(radius=3, atomInvariantsGenerator=invgen)
    fp = g.GetSparseCountFingerprint(m)
    nz = fp.GetNonzeroElements()
    self.assertEqual(len(nz), 13)

    ms = [Chem.MolFromSmiles(x, sanitize=False) for x in ('C1=CC=CN=N1', 'C1C=CC=NN=1')]
    for m in ms:
      m.UpdatePropertyCache()
      Chem.GetSymmSSSR(m)

    g = rdFingerprintGenerator.GetMorganGenerator(radius=2, useBondTypes=True)
    self.assertNotEqual(g.GetSparseCountFingerprint(ms[0]), g.GetSparseCountFingerprint(ms[1]))
    g = rdFingerprintGenerator.GetMorganGenerator(radius=2, useBondTypes=False)
    self.assertEqual(g.GetSparseCountFingerprint(ms[0]), g.GetSparseCountFingerprint(ms[1]))

    binvgen = rdFingerprintGenerator.GetMorganBondInvGen(useBondTypes=False)
    g2 = rdFingerprintGenerator.GetMorganGenerator(radius=2, bondInvariantsGenerator=binvgen)
    self.assertEqual(g.GetSparseCountFingerprint(ms[0]), g2.GetSparseCountFingerprint(ms[0]))
    self.assertEqual(g.GetSparseCountFingerprint(ms[1]), g2.GetSparseCountFingerprint(ms[1]))

  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()