#!/usr/bin/env python
# (C) 2022 Cadence Design Systems, Inc. (Cadence) 
# All rights reserved.
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of Cadence products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. Cadence claims no rights to Customer's
# modifications. Modification of Sample Code is at Customer's sole and
# exclusive risk. Sample Code may require Customer to have a then
# current license or subscription to the applicable Cadence offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  OPENEYE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be
# liable for any damages or liability in connection with the Sample Code
# or its use.


from openeye import oechem
from openeye import oebioisostere

class BroodOptions(oechem.OEOptions):
    def __init__(self):
        oechem.OEOptions.__init__(self, "BroodOptions")

        dbParam = oechem.OEStringParameter("-db")
        dbParam.SetRequired(True)
        dbParam.SetVisibility(oechem.OEParamVisibility_Simple)
        dbParam.SetBrief("Brood Database Directory")
        self._dbParam = self.AddParameter(dbParam)
        
        cpddbParam = oechem.OEStringParameter("-cpddb")
        cpddbParam.SetRequired(True)
        cpddbParam.SetVisibility(oechem.OEParamVisibility_Simple)
        cpddbParam.SetBrief("CPDDatabase File (database of known compounds to identify available compounds similar to hits)")
        self._cpddbParam = self.AddParameter(cpddbParam)

        self._genOpts = oebioisostere.ToGeneralOptions(self.AddOption(oebioisostere.OEBroodGeneralOptions()))
        self._scoreOpts = oebioisostere.ToScoreOptions(self.AddOption(oebioisostere.OEBroodScoreOptions()))
        self._hitOpts = oebioisostere.ToHitlistOptions(self.AddOption(oebioisostere.OEBroodHitlistOptions()))
        pass

    def CreateCopy(self):
        return self

    def GetDatabase(self):
        return self._dbParam.GetStringValue()
        
    def GetCPDDatabase(self):
        return self._cpddbParam.GetStringValue()

    def GetGenOpts(self):
        return self._genOpts

    def GetScoreOpts(self):
        return self._scoreOpts

    def GetHitlistOpts(self):
        return self._hitOpts


def main(argv=[__name__]):
    broodOpts = BroodOptions()
    opts = oechem.OESimpleAppOptions(broodOpts, "BroodCPDDB", oechem.OEFileStringType_Mol3D,
                                     "oeb")
    if oechem.OEConfigureOpts(opts, argv, False) == oechem.OEOptsConfigureStatus_Help:
        return 0
    broodOpts.UpdateValues(opts)

    ifs = oechem.oemolistream()
    if not ifs.open(opts.GetInFile()):
        oechem.OEThrow.Fatal("Unable to open %s for reading" % opts.GetInFile())
        
    query = oebioisostere.OEBroodQuery()
    retCode = oebioisostere.OEReadBroodQuery(ifs, query)
    if retCode != oebioisostere.OEBroodStatusCode_Success:
        oechem.OEThrow.Fatal("Failed: %s" % oebioisostere.OEGetBroodStatus(retCode))

    reader = oebioisostere.OEDBReader()
    retCode = reader.Init(broodOpts.GetDatabase(), query, broodOpts.GetGenOpts())
    if retCode != oebioisostere.OEBroodStatusCode_Success:
        oechem.OEThrow.Fatal("Unable to load Brood database")
        
    ifsCPDDB = oechem.oemolistream()
    if not ifsCPDDB.open(broodOpts.GetCPDDatabase()):
        oechem.OEThrow.Fatal("Unable to load Brood CPDDatabase file")
        
    cpddb = oebioisostere.OECPDDatabase()
    retValue = cpddb.Prep(ifsCPDDB)
    print("Prepration Mode: %s" % oebioisostere.OEGetBroodStatus(retValue))
    
    numOfCPDDBMolecule = cpddb.GetNumMolecules()
    print("Number of molecules in CPDDatabse: %d" % numOfCPDDBMolecule)
    
    if retValue == oebioisostere.OEBroodStatusCode_NewFPGenerated:
      ofsCPDDB = oechem.oemolostream()
      if not ofsCPDDB.open(opts.GetOutFile()):
        oechem.OEThrow.Fatal("Unable to open %s for writing the newly generated finegrprint" , opts.GetOutFile())
      
      if cpddb.Write(ifsCPDDB, ofsCPDDB):
        print("Fingerprints are written in the %s. You can use this file in the future with -cpddb for improved efficiency over current useage." % ofsCPDDB.GetFileName())
      else:
        oechem.OEThrow.Warning("Fingerprint mismatch in writing updated -cpddb file." % ofsCPDDB.GetFileName())

    overlay = oebioisostere.OEBroodOverlay(broodOpts.GetGenOpts(), broodOpts.GetScoreOpts())
    overlay.SetupRef(query);

    hlist = oebioisostere.OEHitlistBuilder(query, broodOpts.GetGenOpts(), broodOpts.GetHitlistOpts())

    packet = oebioisostere.OEBroodDBPacket()
    while reader.GetNextPacket(packet):
        vecScores = overlay.Overlay(packet)
        hlist.AddScores(vecScores);
    print("Generating hitlist...")
    hlist.Build();
    print("Number of final hits: %d \n" % hlist.GetHitCount())

    for idx, hit in enumerate(hlist.GetHits()):
        vecCpddbValues = cpddb.GetSimilarMolecules(hit.GetMol())
        if not len(vecCpddbValues) == 0:
          print("Hit Mol (SMILES): %s \nAnalog Mol (SMILES): %s \nAnalog Mol Label: %s \n \n"
              % (oechem.OEMolToSmiles(hit.GetMol()),vecCpddbValues[0], vecCpddbValues[1]))

if __name__ == "__main__":
    import sys
    sys.exit(main(sys.argv))