##
# File:    OeMoleculeProviderTests.py
# Author:  J. Westbrook
# Date:    1-Oct-2019
# Version: 0.001
#
# Update:
#
#
##
"""
Tests for utilities to read and process OE molecule instances and related databases created
from of PDB chemical component definitions.

"""
import logging
import os
import time
import unittest

from importlib.metadata import version as get_package_version
from rcsb.utils.chem.OeMoleculeProvider import OeMoleculeProvider

__version__ = get_package_version("rcsb.utils.chem")

HERE = os.path.abspath(os.path.dirname(__file__))
TOPDIR = os.path.dirname(os.path.dirname(os.path.dirname(HERE)))

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s]-%(module)s.%(funcName)s: %(message)s")
logger = logging.getLogger()


class OeMoleculeProviderTests(unittest.TestCase):
    skipFlag = True

    def setUp(self):
        self.__startTime = time.time()
        #
        self.__dataPath = os.path.join(HERE, "test-data")
        self.__cachePath = os.path.join(HERE, "test-output")
        self.__ccUrlTarget = os.path.join(self.__dataPath, "components-abbrev.cif")
        self.__birdUrlTarget = os.path.join(self.__dataPath, "prdcc-abbrev.cif")
        #
        logger.debug("Running tests on version %s", __version__)
        logger.info("Starting %s at %s", self.id(), time.strftime("%Y %m %d %H:%M:%S", time.localtime()))

    def tearDown(self):
        endTime = time.time()
        logger.info("Completed %s at %s (%.4f seconds)", self.id(), time.strftime("%Y %m %d %H:%M:%S", time.localtime()), endTime - self.__startTime)

    def testBuildMoleculeCacheFilesAbbrev(self):
        for molBuildType in ["oe-iso-smiles", "model-xyz", "ideal-xyz"]:
            self.__testBuildMoleculeCacheFiles(
                ccUrlTarget=self.__ccUrlTarget,
                birdUrlTarget=self.__birdUrlTarget,
                ccFileNamePrefix="cc-abbrev",
                oeFileNamePrefix="oe-abbrev",
                molLimit=24,
                fpTypeList=["TREE", "PATH", "MACCS", "CIRCULAR", "LINGO"],
                screenTypeList=["SMARTS"],
                molBuildType=molBuildType,
            )

    @unittest.skipIf(skipFlag, "Long test")
    def testBuildMoleculeCacheFilesSubset(self):
        """Test construction of OE resource files for a subset (1K) of full public chemical reference dictionary"""
        for molBuildType in ["oe-iso-smiles", "model-xyz", "ideal-xyz"]:
            self.__testBuildMoleculeCacheFiles(
                molBuildType=molBuildType,
                molLimit=1000,
                ccFileNamePrefix="cc-1k",
                oeFileNamePrefix="oe-1k",
                screenTypeList=["SMARTS"],
                fpTypeList=["TREE", "PATH", "MACCS", "CIRCULAR", "LINGO"],
            )
        #

    @unittest.skipIf(skipFlag, "Long test")
    def testBuildMoleculeCacheFilesFiltered(self):
        """Test construction of OE resource files for a filtered subset of full public chemical reference dictionary"""
        for molBuildType in ["oe-iso-smiles", "model-xyz", "ideal-xyz"]:
            self.__testBuildMoleculeCacheFiles(
                molBuildType=molBuildType,
                molLimit=1000,
                ccFileNamePrefix="cc-filtered",
                oeFileNamePrefix="oe-filtered",
                screenTypeList=["SMARTS"],
                fpTypeList=["TREE", "PATH", "MACCS", "CIRCULAR", "LINGO"],
            )
        #

    @unittest.skipIf(skipFlag, "Long test")
    def testBuildMoleculeCacheFilesFull(self):
        """Test construction of OE resource files for the full public chemical reference dictionary

        For 31K cc's the build time is ~3123 secs (4 procs)

        """
        for molBuildType in ["oe-iso-smiles", "model-xyz", "ideal-xyz"]:
            self.__testBuildMoleculeCacheFiles(
                molBuildType=molBuildType, ccFileNamePrefix="cc-full", oeFileNamePrefix="oe-full", fpTypeList=["TREE", "PATH", "MACCS", "CIRCULAR", "LINGO"]
            )

    @unittest.skipIf(skipFlag, "Long test")
    def testBuildMoleculeCacheFilesFullScreened(self):
        """Test construction of OE resource files with screened substructure database for the full public chemical reference dictionary"""
        for molBuildType in ["oe-iso-smiles"]:
            self.__testBuildMoleculeCacheFiles(
                molBuildType=molBuildType,
                ccFileNamePrefix="cc-full",
                oeFileNamePrefix="oe-full",
                screenTypeList=["SMARTS"],
                fpTypeList=["TREE", "PATH", "MACCS", "CIRCULAR", "LINGO"],
            )

    def __testBuildMoleculeCacheFiles(self, **kwargs):
        """Test build OE cache files from full component dictionary"""
        ccUrlTarget = kwargs.get("ccUrlTarget", None)
        birdUrlTarget = kwargs.get("birdUrlTarget", None)
        molLimit = kwargs.get("molLimit", 0)
        quietFlag = kwargs.get("quietFlag", True)
        molBuildType = kwargs.get("molBuildType", "ideal-xyz")
        fpTypeList = kwargs.get("fpTypeList", ["TREE"])
        screenTypeList = kwargs.get("screenTypeList", [])
        ccFileNamePrefix = kwargs.get("ccFileNamePrefix", "cc")
        oeFileNamePrefix = kwargs.get("oeFileNamePrefix", "oe")
        #
        startTime = time.time()

        oemp = OeMoleculeProvider(
            ccUrlTarget=ccUrlTarget,
            birdUrlTarget=birdUrlTarget,
            cachePath=self.__cachePath,
            ccFileNamePrefix=ccFileNamePrefix,
            oeFileNamePrefix=oeFileNamePrefix,
            molBuildType=molBuildType,
            useCache=False,
            quietFlag=quietFlag,
            fpTypeList=fpTypeList,
            screenTypeList=screenTypeList,
        )
        ok = oemp.testCache()
        self.assertTrue(ok)
        oemp = OeMoleculeProvider(
            ccUrlTarget=ccUrlTarget,
            birdUrlTarget=birdUrlTarget,
            cachePath=self.__cachePath,
            ccFileNamePrefix=ccFileNamePrefix,
            oeFileNamePrefix=oeFileNamePrefix,
            molBuildType=molBuildType,
            useCache=True,
        )
        endTime = time.time()
        logger.info(">> Completed load molBuildType %r molLimit %r (%.4f seconds)", molBuildType, molLimit, endTime - startTime)
        #
        # ---

        deltaMol = 2
        minMol = minNumFp = molLimit - deltaMol if molLimit else 30000
        for fpType in fpTypeList:
            fpDb = oemp.getFingerPrintDb(fpType="TREE")
            logger.debug("fpType %r length %d", fpType, fpDb.NumFingerPrints())
            self.assertGreaterEqual(fpDb.NumFingerPrints(), minNumFp)
        #
        ccId = "004"
        oeMol = oemp.getMol(ccId)
        logger.debug("%s atom count %d", ccId, len(list(oeMol.GetAtoms())))
        #
        if molBuildType in ["oe-iso-smiles"]:
            self.assertGreaterEqual(len(list(oeMol.GetAtoms())), 12)
        else:
            self.assertGreaterEqual(len(list(oeMol.GetAtoms())), 20)
        #
        oeDb, oeDbIdx = oemp.getOeMolDatabase()
        logger.debug("Type db %r length %d type idx %r length %d", type(oeDb), oeDb.NumMols(), type(oeDbIdx), len(oeDbIdx))
        self.assertGreaterEqual(oeDb.NumMols(), minMol)
        self.assertGreaterEqual(len(oeDbIdx), minMol)
        #
        if molBuildType in ["oe-iso-smiles"] and screenTypeList:
            ssDb = oemp.getSubSearchDb()
            self.assertGreaterEqual(ssDb.NumMolecules(), minMol)
        return True


def buildCacheFiles():
    suiteSelect = unittest.TestSuite()
    suiteSelect.addTest(OeMoleculeProviderTests("testBuildMoleculeCacheFilesAbbrev"))
    suiteSelect.addTest(OeMoleculeProviderTests("testBuildMoleculeCacheFilesSubset"))
    suiteSelect.addTest(OeMoleculeProviderTests("testBuildMoleculeCacheFilesFiltered"))
    suiteSelect.addTest(OeMoleculeProviderTests("testBuildMoleculeCacheFilesFull"))
    return suiteSelect


if __name__ == "__main__":

    mySuite = buildCacheFiles()
    unittest.TextTestRunner(verbosity=2).run(mySuite)
