##
#  File:           ChEMBLTargetCofactorProvider.py
#  Date:           15-Jun-2021 jdw
#
#  Updated:
#  20-Aug-2024 dwp Add support for loading and accessing data on MongoDB
##
"""
Accessors for ChEMBL target cofactors.
"""

import datetime
import logging
import math
import os.path
import time

from rcsb.utils.io.FileUtil import FileUtil
from rcsb.utils.io.MarshalUtil import MarshalUtil
from rcsb.utils.io.StashableBase import StashableBase
from rcsb.utils.targets.ChEMBLTargetActivityProvider import ChEMBLTargetActivityProvider
from rcsb.utils.targets.ChEMBLTargetProvider import ChEMBLTargetProvider
from rcsb.utils.targets.TargetCofactorDbProvider import TargetCofactorDbProvider

logger = logging.getLogger(__name__)


class ChEMBLTargetCofactorProvider(StashableBase):
    """Accessors for ChEMBL target cofactors."""

    def __init__(self, **kwargs):
        #
        self.__cachePath = kwargs.get("cachePath", ".")
        self.__dirName = "ChEMBL-cofactors"
        self.__cofactorResourceName = "chembl"
        super(ChEMBLTargetCofactorProvider, self).__init__(self.__cachePath, [self.__dirName])
        self.__dirPath = os.path.join(self.__cachePath, self.__dirName)
        #
        self.__mU = MarshalUtil(workPath=self.__dirPath)
        self.__fD = self.__reload(self.__dirPath, **kwargs)
        #

    def testCache(self, minCount=1):
        logger.info("ChEMBL cofactor count %d", len(self.__fD["cofactors"]) if "cofactors" in self.__fD else 0)
        if self.__fD and "cofactors" in self.__fD and len(self.__fD["cofactors"]) > minCount:
            return True
        else:
            return False

    def hasTarget(self, rcsbEntityId):
        return rcsbEntityId.upper() in self.__fD["cofactors"]

    def getTargets(self, rcsbEntityId):
        try:
            return self.__fD["cofactors"][rcsbEntityId.upper()]
        except Exception:
            return []

    def __getCofactorDataPath(self):
        return os.path.join(self.__dirPath, "ChEMBL-cofactor-data.json")

    def reload(self):
        self.__fD = self.__reload(self.__dirPath, useCache=True)
        return True

    def __reload(self, dirPath, **kwargs):
        startTime = time.time()
        fD = {}
        useCache = kwargs.get("useCache", True)
        ok = False
        cofactorPath = self.__getCofactorDataPath()
        #
        logger.info("useCache %r cofactorPath %r", useCache, cofactorPath)
        if useCache and self.__mU.exists(cofactorPath):
            fD = self.__mU.doImport(cofactorPath, fmt="json")
            ok = len(fD) > 0
        else:
            fU = FileUtil()
            fU.mkdir(dirPath)
        # ---
        logger.info("Completed reload with status (%r) at %s (%.4f seconds)", ok, time.strftime("%Y %m %d %H:%M:%S", time.localtime()), time.time() - startTime)
        return fD

    def getCofactorDataDict(self):
        return self.__fD["cofactors"]

    def buildCofactorList(self, sequenceMatchFilePath, crmpObj=None, lnmpObj=None, maxActivity=5):
        """Build target cofactor list for the matching entities in the input sequence match file.

        Args:
            sequenceMatchFilePath (str): sequence match output file path (e.g., "./CACHE/sequence-search-results/chembl-vs-pdbprent-filtered-results.json")
            crmpObj (obj, optional): instance of ChemRefMappingProviderObj()
            lnmpObj (obj, optional): instance of LigandNeighborMappingProviderObj(). Defaults to None.
            maxActivity (int, optional): maximum number of prioritized activity records per target

        Returns:
            bool: True for success or False otherwise

            Example activity record -

                        "CHEMBL3243": [
                    {
                        "assay_chembl_id": "CHEMBL655768",
                        "assay_description": "In vitro inhibitory activity against recombinant human CD45 using fluorescein diphosphate (FDP) as a substrate",
                        "assay_type": "B",
                        "canonical_smiles": "COC(=O)c1ccc(C2=CC(=O)C(=O)c3ccccc32)cc1",
                        "ligand_efficiency": {
                        "bei": "19.78",
                        "le": "0.36",
                        "lle": "3.11",
                        "sei": "9.57"
                        },
                        "molecule_chembl_id": "CHEMBL301254",
                        "parent_molecule_chembl_id": "CHEMBL301254",
                        "pchembl_value": "5.78",
                        "standard_relation": "=",
                        "standard_type": "IC50",
                        "standard_units": "nM",
                        "standard_value": "1650.0",
                        "target_chembl_id": "CHEMBL3243"
                    },
        """
        rDL = []
        mD = self.__mU.doImport(sequenceMatchFilePath, fmt="json")
        #
        chP = ChEMBLTargetProvider(cachePath=self.__cachePath, useCache=False)
        if not chP.testCache():
            logger.warning("Skipping build of target cofactor list because ChEMBL Target data is missing.")
            return False
        # ---
        chaP = ChEMBLTargetActivityProvider(cachePath=self.__cachePath, useCache=True)
        if not chaP.testCache():
            logger.warning("Skipping build of target cofactor list because ChEMBL Target Activity data is missing.")
            return False
        #
        provenanceSource = "ChEMBL"
        refScheme = "PDB entity"
        assignVersion = chP.getAssignmentVersion()
        for queryId, matchDL in mD.items():
            qCmtD = self.__decodeComment(queryId)
            unpId = qCmtD["uniprotId"]
            queryTaxId = qCmtD["taxId"] if "taxId" in qCmtD else None
            chemblIdL = qCmtD["chemblId"].split(",")
            if queryTaxId == "-1":
                logger.info("Skipping target with missing taxonomy %r (%r)", unpId, chemblIdL)
                continue
            queryName = chP.getTargetDescription(unpId)
            for chemblId in chemblIdL:
                if not chaP.hasTargetActivity(chemblId):
                    logger.debug("Skipping target %r (%r)", unpId, chemblId)
                    # continue
                # --
                chemCompNeighborsD = {}
                if lnmpObj:
                    for matchD in matchDL:
                        tCmtD = self.__decodeComment(matchD["target"])
                        entryId = tCmtD["entityId"].split("_")[0]
                        entityId = tCmtD["entityId"].split("_")[1]
                        rcsbEntityId = entryId + "_" + entityId
                        chemCompIdList = lnmpObj.getLigandNeighbors(rcsbEntityId)
                        chemCompNeighborsD.update({k: True for k in chemCompIdList})
                # --
                for matchD in matchDL:
                    tCmtD = self.__decodeComment(matchD["target"])
                    entryId = tCmtD["entityId"].split("_")[0]
                    entityId = tCmtD["entityId"].split("_")[1]
                    #
                    taDL = chaP.getTargetActivity(chemblId)
                    logger.debug("Target %r has (%d) activity records", chemblId, len(taDL))
                    # ---
                    actL = []
                    for taD in taDL:
                        if taD["assay_type"] in ["B", "F"]:
                            try:
                                if taD["standard_units"] == "nM" and taD["standard_value"] and float(taD["standard_value"]) > 0.0:
                                    pV = -math.log10(float(taD["standard_value"]) * 10.0e-9)
                                    actD = {
                                        "cofactor_id": taD["molecule_chembl_id"],
                                        "assay_id": taD["assay_chembl_id"],
                                        "assay_description": taD["assay_description"],
                                        "measurement_type": "p" + taD["standard_type"],
                                        "measurement_value": round(pV, 2),
                                        "smiles": taD["canonical_smiles"],
                                        "molecule_name": taD["molecule_name"],
                                        "inchi_key": taD["inchi_key"],
                                        "action": taD["action"],
                                        "moa": taD["moa"],
                                        "max_phase": taD["max_phase"],
                                    }
                                    actD = self.__addLocalIds(actD, crmpObj=crmpObj)
                                    actL.append(actD)
                            except Exception as e:
                                logger.debug("Failing for tAD %r with %s", taD, str(e))

                    # ---
                    actL = self.__activityListSelect(actL, chemCompNeighborsD, maxActivity=maxActivity)
                    if not actL:
                        logger.debug("No ChEMBL cofactors for %s %s", chemblId, unpId)
                    # ---
                    # aligned_target.entity_beg_seq_id (current target is PDB entity in json)
                    # aligned_target.target_beg_seq_id (current query is target seq in json)
                    # aligned_target.length
                    fpL = []
                    if "alignedRegions" in matchD:
                        fpL = [
                            {
                                "entity_beg_seq_id": arD["targetBegin"],
                                "target_beg_seq_id": arD["queryBegin"],
                                "length": arD["targetEnd"] - arD["targetBegin"],
                            }
                            for arD in matchD["alignedRegions"]
                        ]
                    else:
                        fpL = [
                            {
                                "entity_beg_seq_id": matchD["targetBegin"],
                                "target_beg_seq_id": matchD["queryBegin"],
                                "length": matchD["alignLen"],
                            }
                        ]
                    # ---
                    rD = {
                        "entry_id": entryId,
                        "entity_id": entityId,
                        "query_uniprot_id": unpId,
                        "query_id": chemblId,
                        "query_id_type": "ChEMBL",
                        "query_name": queryName,
                        "provenance_source": provenanceSource,
                        "reference_scheme": refScheme,
                        "assignment_version": assignVersion,
                        "query_taxonomy_id": int(queryTaxId) if queryTaxId else None,
                        "target_taxonomy_id": int(matchD["targetTaxId"]) if "targetTaxId" in matchD else None,
                        #
                        "aligned_target": fpL,
                        #
                        "taxonomy_match_status": matchD["taxonomyMatchStatus"] if "taxonomyMatchStatus" in matchD else None,
                        "lca_taxonomy_id": matchD["lcaTaxId"] if "lcaTaxId" in matchD else None,
                        "lca_taxonomy_name": matchD["lcaTaxName"] if "lcaTaxName" in matchD else None,
                        "lca_taxonomy_rank": matchD["lcaRank"] if "lcaRank" in matchD else None,
                        "cofactors": actL,
                    }
                    rDL.append(rD)
            #
        qD = {}
        for rD in rDL:
            eId = rD["entry_id"] + "_" + rD["entity_id"]
            qD.setdefault(eId, []).append(rD)
        #
        tS = datetime.datetime.now().isoformat()
        # vS = datetime.datetime.now().strftime("%Y-%m-%d")
        vS = assignVersion
        #
        # Write out cofactor data set
        fp = self.__getCofactorDataPath()
        ok = self.__mU.doExport(fp, {"version": vS, "created": tS, "cofactors": qD}, fmt="json", indent=3)
        #
        return ok

    def __addLocalIds(self, cfD, crmpObj=None):
        #
        if crmpObj:
            localIdL = crmpObj.getLocalIds("CHEMBL", cfD["cofactor_id"])
            if localIdL:
                localId = localIdL[0]
                if localId.startswith("PRD_"):
                    cfD["prd_id"] = localId
                else:
                    cfD["chem_comp_id"] = localId
        return cfD

    def __activityListSelect(self, activityDL, chemCompNeighborsD, maxActivity=5):
        retL = []
        mappedNeighborL = []
        unmappedL = activityDL
        #
        if chemCompNeighborsD:
            unmappedL = []
            # Select out the any cases for molecules that map to a neighbor chemical component.
            for activityD in activityDL:
                if "chem_comp_id" in activityD and activityD["chem_comp_id"] in chemCompNeighborsD:
                    activityD["neighbor_in_pdb"] = "Y"
                    mappedNeighborL.append(activityD)
                else:
                    unmappedL.append(activityD)
                    activityD["neighbor_in_pdb"] = "N"
        #
        numLeft = maxActivity - len(mappedNeighborL)
        if numLeft > 0:
            unmappedL = sorted(unmappedL, key=lambda k: k["measurement_value"], reverse=True)
            retL = mappedNeighborL
            retL.extend(unmappedL[:numLeft])
            retL = sorted(retL, key=lambda k: k["measurement_value"], reverse=True)
        else:
            logger.debug("Mapped neighbor cofactors (%d) excluded unmapped (%d)", len(mappedNeighborL), len(unmappedL))
            retL = sorted(mappedNeighborL, key=lambda k: k["measurement_value"], reverse=True)
        return retL

    def __decodeComment(self, comment, separator="|"):
        dD = {}
        try:
            ti = iter(comment.split(separator))
            dD = {tup[1]: tup[0] for tup in zip(ti, ti)}
        except Exception:
            pass
        return dD

    def loadCofactorData(self, cfgOb, **kwargs):
        """Load cofactor data to MongoDB.
        """
        ok = False
        try:
            tcDbP = TargetCofactorDbProvider(
                cachePath=self.__cachePath,
                cfgOb=cfgOb,
                cofactorResourceName=self.__cofactorResourceName,
                **kwargs
            )
            ok = tcDbP.loadCofactorData(self.__cofactorResourceName, self)
        except Exception as e:
            logger.exception("Failing with %s", str(e))
        #
        logger.info("%r cofactor data DB load status (%r)", self.__cofactorResourceName, ok)
        return ok


class ChEMBLTargetCofactorAccessor:
    def __init__(self, cachePath, cfgOb=None, **kwargs):
        """
        Accessor class for fetching cofactor data from MongoDB.
        """
        self.__cofactorResourceName = "chembl"
        self.__cfgOb = cfgOb
        self.__cachePath = cachePath
        #
        self.__tcDbP = TargetCofactorDbProvider(
            cachePath=self.__cachePath,
            cfgOb=self.__cfgOb,
            cofactorResourceName=self.__cofactorResourceName,
            **kwargs
        )

    def testCache(self, minCount=0):
        docCount = self.__tcDbP.cofactorDbCount()
        logger.info("Loaded %s cofactor DB count %d", self.__cofactorResourceName, docCount)
        return docCount >= minCount

    def reload(self):
        return True

    def getTargets(self, rcsbEntityId, dataFieldName="rcsb_cofactors", **kwargs):
        return self.__tcDbP.fetchCofactorData(rcsbEntityId, dataFieldName, **kwargs)
