import os
import unittest
from tempfile import NamedTemporaryFile
from glob import glob
from moleculekit.home import home
from moleculekit.smallmol.smallmol import SmallMol
from moleculekit.smallmol.smallmollib import SmallMolLib
import rdkit
from pandas import core


SDF_N_MOLS = 100

SDF_IDS_DELETE = [1,10,15,16]
SDF_MOLNAME_DELETE = ['ZINC02583363', 'ZINC86860147', 'ZINC04342657', 'ZINC02023420']
SDF_FIELDS = ['ligname', '_mol']
SDF_LOC_0_99 = 'ZINC02141008'

class _TestSmallMol(unittest.TestCase):

    def setUp(self):
        self.dataDir= home('test-smallmol')

    def test_loaSdffile(self):
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')
        lib = SmallMolLib(sdffile)
        n_mols = lib.numMols
        self.assertEqual(n_mols, SDF_N_MOLS, 'Molecules not correctly loaded. '
                                             'Expected: {}; Now: {}'.format(SDF_N_MOLS, n_mols))

        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf.gz')
        lib = SmallMolLib(sdffile)
        n_mols = lib.numMols
        self.assertEqual(n_mols, SDF_N_MOLS, 'Molecules not correctly loaded. '
                                             'Expected: {}; Now: {}'.format(SDF_N_MOLS, n_mols))

    def test_writeSdf(self):
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')
        lib = SmallMolLib(sdffile)

        sdfname = NamedTemporaryFile().name + '.sdf'
        lib.writeSdf(sdfname)

        sdf_exists = os.path.isfile(sdfname)

        self.assertTrue(sdf_exists, msg="The sdf written was not found")

        sdf = SmallMolLib(sdfname)

        self.assertIsInstance(sdf, SmallMolLib, msg="The sdf written was not correctly loaded. Probably the previous"
                                                    "writing went wrong")


    def test_appendSmallMolLib(self):
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')
        lib = SmallMolLib(sdffile)
        lib2 = SmallMolLib(sdffile)

        lib.appendSmallLib(lib2)

        n_mol2_merged = lib.numMols

        self.assertEqual(n_mol2_merged, SDF_N_MOLS*2, msg="The number of molecules in the SmallMolLib is not as expected."
                                                          "The two sdf were not correctly merged. ")


    def test_appendSmallMol(self):
        mol2file = os.path.join(self.dataDir, 'benzamidine.mol2')
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')

        lib = SmallMolLib(sdffile)
        sm = SmallMol(mol2file)
        lib.appendSmallMol(sm)

        n_mol2_append = lib.numMols

        self.assertEqual(n_mol2_append, SDF_N_MOLS+1, msg="The number of molecules in the SmallMolLib is not as expected."
                                                          "The mol2 were not correctly append. ")

    def test_removeMols(self):
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')

        lib = SmallMolLib(sdffile)

        mols_ids = SDF_IDS_DELETE
        ref_mols_name = SDF_MOLNAME_DELETE

        mols_name = [ s.ligname for s in lib.getMols(mols_ids) ]

        self.assertListEqual(mols_name, ref_mols_name, msg="The molecules at the given indexes do not match with the"
                                                           "expected")
        lib.removeMols(mols_ids)

        mols_name_now = [ s.ligname for s in lib.getMols(mols_ids) ]

        self.assertFalse(mols_name_now == mols_name, msg="The molecules seem to not be deleted correctly")

    def test_convertToDataFrame(self):
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')

        lib = SmallMolLib(sdffile)

        df = lib.toDataFrame()

        self.assertIsInstance(df, core.frame.DataFrame, msg="The SmallMolLib object was not correctly converted into pandas"
                                                      "DataFrame")

        cols = df.columns.tolist()
        ref_cols =SDF_FIELDS

        self.assertEqual(cols, ref_cols, msg="The fields in the SmallMolLib object was not the expected one")

        ligname_99 = df.iloc[99][0]
        ref_ligname = SDF_LOC_0_99

        self.assertEqual(ligname_99, ref_ligname, msg="The ligand name found is not the expected one")

    @unittest.skipIf(rdkit.__version__ == '2019.03.4', "new rdkit versions break 2D depiction")
    def test_depict(self):
        import IPython
        refimg = os.path.join(self.dataDir, 'sdf.svg')
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')

        lib = SmallMolLib(sdffile)

        img_name = NamedTemporaryFile().name + '.svg'
        lib.depict(sketch=True, filename=img_name)
        _img = lib.depict(sketch=True, ipython=True)

        refimg_size = os.path.getsize(refimg)
        sm_img_size = os.path.getsize(img_name)

        self.assertIsInstance(_img, IPython.core.display.SVG, msg="The object is not an IPython image as expected")
        self.assertEqual(sm_img_size, refimg_size, msg="The svg image does not have the same size of the reference")


    def test_writeSmiles(self):
        sdffile = os.path.join(self.dataDir, 'fda_drugs_light.sdf')

        lib = SmallMolLib(sdffile)
        tmpfile = NamedTemporaryFile().name + '.smi'
        lib.writeSmiles(tmpfile)

        with open(tmpfile, 'r') as f:
            filelines = f.readlines()[1:]

        with open(os.path.join(self.dataDir, 'fda_drugs_light.smi'), 'r') as f:
            reflines = f.readlines()[1:]

        self.assertEqual(filelines, reflines)

    def test_readSmiles(self):
        smifile = os.path.join(self.dataDir, 'fda_drugs_light.smi')
        lib = SmallMolLib(smifile)
        assert len(lib) == 100

        smifile = os.path.join(self.dataDir, 'fda_drugs_light.smi.gz')
        lib = SmallMolLib(smifile)
        assert len(lib) == 100



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