#!/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("Database folder")
        self._dbParam = self.AddParameter(dbParam)

        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 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, "BroodHitlist", oechem.OEFileStringType_Mol3D,
                                     oechem.OEFileStringType_Mol3D)
    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())

    ofs = oechem.oemolostream()
    if not ofs.open(opts.GetOutFile()):
        oechem.OEThrow.Fatal("Unable to open %s for writing" % opts.GetOutFile())

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

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

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

    packetCount = 0
    packet = oebioisostere.OEBroodDBPacket()
    while reader.GetNextPacket(packet):
        packetCount += 1
        print("Processing packet %d with %d fragments" % (packetCount, packet.GetFragCount()))
        vecScores = overlay.Overlay(packet)
        hlist.AddScores(vecScores);
    print("Generating hitlist...")
    hlist.Build();
    print("Total number of fragments overlayed: %d" % hlist.GetAddCount())
    print("Number of final hits: %d" % hlist.GetHitCount())

    for idx, hit in enumerate(hlist.GetHits()):
        oechem.OEWriteMolecule(ofs, hit.GetMol())
        comboScore = hit.GetComboScore()
        if broodOpts.GetGenOpts().GetScoreType() == oebioisostere.OEBroodScoreType_ET:
            comboScore = hit.GetETComboScore()
        print("Hit: %d %s Combo Score: %2f Belief Score: %.2f Complexity: %2f"
              % (idx+1, hit.GetMol().GetTitle(), comboScore, hit.GetBeliefScore(), hit.GetComplexity()))


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