#!/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 BroodComboOptions(oechem.OEOptions):
    def __init__(self):
        oechem.OEOptions.__init__(self, "BroodComboOptions")

        dbParam = oechem.OEStringParameter("-db")
        dbParam.SetRequired(True)
        dbParam.SetVisibility(oechem.OEParamVisibility_Simple)
        dbParam.SetBrief("Database folder")
        self._dbParam = self.AddParameter(dbParam)

        maxHitsParam = oechem.OEUIntParameter("-primaryQueryMaxHits", 15)
        maxHitsParam.SetVisibility(oechem.OEParamVisibility_Simple)
        maxHitsParam.SetBrief("Maximum number of primary hits collected from the first query")
        self._maxHitsParam = self.AddParameter(maxHitsParam)

        secondaryMaxHitsParam = oechem.OEUIntParameter("-secondaryQueryMaxHits", 15)
        secondaryMaxHitsParam.SetVisibility(oechem.OEParamVisibility_Simple)
        secondaryMaxHitsParam.SetBrief("Maximum number of secondary combo hits collected for each primary hit")
        self._secondaryMaxHitsParam = self.AddParameter(secondaryMaxHitsParam)

    def CreateCopy(self):
        return self

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

    def GetPrimaryQueryMaxHits(self):
        if self._maxHitsParam.GetHasValue():
            return int(self._maxHitsParam.GetStringValue())
        return int(self._maxHitsParam.GetStringDefault())

    def GetSecondaryQueryMaxHits(self):
        if self._secondaryMaxHitsParam.GetHasValue():
            return int(self._secondaryMaxHitsParam.GetStringValue())
        return int(self._secondaryMaxHitsParam.GetStringDefault())

def main(argv=[__name__]):
    broodOpts = BroodComboOptions()
    opts = oechem.OERefInputAppOptions(
        broodOpts,
        "BroodComboBuilder",
        oechem.OEFileStringType_Mol3D,
        oechem.OEFileStringType_Mol3D,
        oechem.OEFileStringType_Mol3D,
        "-in2",
    )
    if oechem.OEConfigureOpts(opts, argv, False) == oechem.OEOptsConfigureStatus_Help:
        return 0
    broodOpts.UpdateValues(opts)

    ifs1 = oechem.oemolistream()
    if not ifs1.open(opts.GetInFile()):
        oechem.OEThrow.Fatal("Unable to open %s for reading" % opts.GetInFile())

    ifs2 = oechem.oemolistream()
    if not ifs2.open(opts.GetRefFile()):
        oechem.OEThrow.Fatal("Unable to open %s for reading" % opts.GetRefFile())

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

    query1 = oebioisostere.OEBroodQuery()
    retCode = oebioisostere.OEReadBroodQuery(ifs1, query1)
    if retCode != oebioisostere.OEBroodStatusCode_Success:
        oechem.OEThrow.Fatal("Failed to read -in1 query: %s" % oebioisostere.OEGetBroodStatus(retCode))

    query2 = oebioisostere.OEBroodQuery()
    retCode = oebioisostere.OEReadBroodQuery(ifs2, query2)
    if retCode != oebioisostere.OEBroodStatusCode_Success:
        oechem.OEThrow.Fatal("Failed to read -in2 query: %s" % oebioisostere.OEGetBroodStatus(retCode))

    # Build primary hits from the first query.
    vecPrimaryHits = []
    primaryReader = oebioisostere.OEDBReader()
    if primaryReader.Init(broodOpts.GetDatabase()) != oebioisostere.OEBroodStatusCode_Success:
        oechem.OEThrow.Fatal("Unable to open database for primary query hits")
    primaryOverlay = oebioisostere.OEBroodOverlay()
    primaryOverlay.SetupRef(query1)
    primaryBuilder = oebioisostere.OEBroodMolBuilder(query1)
    primaryPacket = oebioisostere.OEBroodDBPacket()
    while primaryReader.GetNextPacket(primaryPacket) and len(vecPrimaryHits) < broodOpts.GetPrimaryQueryMaxHits():
        for match in primaryOverlay.Overlay(primaryPacket):
            hit = oebioisostere.OEBroodHit()
            if primaryBuilder.Build(match, hit) == oebioisostere.OEBroodStatusCode_Success:
                vecPrimaryHits.append(hit)
                if len(vecPrimaryHits) >= broodOpts.GetPrimaryQueryMaxHits():
                    break

    if len(vecPrimaryHits) == 0:
        oechem.OEThrow.Fatal("Unable to build primary hits from -in1")

    totalSecondaryHits = 0
    written = 0
    allSecondaryHits = []

    # The attempt budget below bounds Build() calls per primary hit
    # regardless of success rate.
    maxBuildAttempts = broodOpts.GetSecondaryQueryMaxHits() * 20

    # For each primary hit from query1, run query2 as the reference query in combo mode.
    # The per-primary cap keeps each combo expansion bounded by -secondaryQueryMaxHits.
    for primaryHit in vecPrimaryHits:
        secondaryReader = oebioisostere.OEDBReader()
        if secondaryReader.Init(broodOpts.GetDatabase()) != oebioisostere.OEBroodStatusCode_Success:
            oechem.OEThrow.Fatal("Unable to build secondary combo hits from -in2")
        # Build an overlay pipeline for query2, then combine each score with the current
        # primary hit via OEBroodComboBuilder.
        secondaryOverlay = oebioisostere.OEBroodOverlay()
        secondaryOverlay.SetupRef(query2)
        secondaryBuilder = oebioisostere.OEBroodComboBuilder(query2, primaryHit)
        secondaryHits = []
        secondaryPacket = oebioisostere.OEBroodDBPacket()
        buildAttempts = 0
        done = False
        # Stream the database packets; stop once this primary hit reaches its cap
        # or exhausts the attempt budget.
        while secondaryReader.GetNextPacket(secondaryPacket) and not done:
            for match in secondaryOverlay.Overlay(secondaryPacket):
                hit = oebioisostere.OEBroodHit()
                # Combo build succeeds only when the secondary replacement can be merged with
                # the primary context for the current query2 reference.
                if secondaryBuilder.Build(match, hit) == oebioisostere.OEBroodStatusCode_Success:
                    secondaryHits.append(hit)
                    if len(secondaryHits) >= broodOpts.GetSecondaryQueryMaxHits():
                        done = True
                        break
                buildAttempts += 1
                if buildAttempts >= maxBuildAttempts:
                    done = True
                    break
        allSecondaryHits.extend(secondaryHits)

    if len(allSecondaryHits) > 0:
        # Cluster and rank all secondary hits before writing output molecules.
        clBuilder = oebioisostere.OEBroodClusterBuilder(query1)
        vecHits = []
        for hit in allSecondaryHits:
            vecHits.append(hit)

        if not clBuilder.Add(vecHits):
            oechem.OEThrow.Fatal("Unable to cluster secondary combo hits")
        if not clBuilder.Rank():
            oechem.OEThrow.Fatal("Unable to rank secondary combo clusters")

        vecClusters = clBuilder.GetClusters()

        for cluster in vecClusters:
            oechem.OEWriteMolecule(ofs, cluster.GetHead().GetMol())
            written += 1
            totalSecondaryHits += 1

        for cluster in vecClusters:
            for hit in cluster.GetMembers():
                oechem.OEWriteMolecule(ofs, hit.GetMol())
                written += 1
                totalSecondaryHits += 1

    print("Primary hits: %d" % len(vecPrimaryHits))
    print("Secondary combo hits: %d" % totalSecondaryHits)
    print("Wrote %d molecules to %s" % (written, opts.GetOutFile()))
    return 0


if __name__ == "__main__":
    import sys

    sys.exit(main(sys.argv))