#!/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

# Demonstrates how to collect, curate, and write scored Brood fragment hits
# with OEScoreHitlist starting from a query and fragment database.

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

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

        maxHitsParam = oechem.OEUIntParameter("-maxScoreHits", 100)
        maxHitsParam.SetVisibility(oechem.OEParamVisibility_Simple)
        maxHitsParam.SetBrief("Maximum number of unique scored fragments curated by OEScoreHitlist")
        self._maxHitsParam = self.AddParameter(maxHitsParam)

    def CreateCopy(self):
        return self

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

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


def main(argv=[__name__]):
    broodOpts = BroodScoreHitlistOptions()
    opts = oechem.OESimpleAppOptions(
        broodOpts,
        "BroodScoreHitlist",
        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 to read query: %s" % oebioisostere.OEGetBroodStatus(retCode))

    # Read database packets and overlay them against the query.
    reader = oebioisostere.OEDBReader()
    if reader.Init(broodOpts.GetDatabase(), query) != oebioisostere.OEBroodStatusCode_Success:
        oechem.OEThrow.Fatal("Unable to open Brood database '%s'" % broodOpts.GetDatabase())

    overlay = oebioisostere.OEBroodOverlay()
    overlay.SetupRef(query)

    # Use OEScoreHitlist to deduplicate and rank scored fragment matches.
    scoreHitlist = oebioisostere.OEScoreHitlist(broodOpts.GetMaxScoreHits())

    packetCount = 0
    packet = oebioisostere.OEBroodDBPacket()
    while reader.GetNextPacket(packet):
        packetCount += 1
        scoreHitlist.AddScores(overlay.Overlay(packet))

    scoreHitlist.Build()

    print("Database packets processed:    %d" % packetCount)
    print("Total scores added:            %d" % scoreHitlist.GetAddCount())
    print("Unique fragment matches:       %d" % scoreHitlist.GetMatchCount())
    print("Duplicate scores filtered:     %d" % scoreHitlist.GetDuplicateCount())
    print("Curated score hits:            %d" % scoreHitlist.GetHitCount())

    # Write the curated fragment hits selected by the score hitlist.
    for score in scoreHitlist.GetHits():
        oechem.OEWriteMolecule(ofs, score.GetFrag())

    print("Wrote %d molecules to %s" % (scoreHitlist.GetHitCount(), opts.GetOutFile()))
    return 0


if __name__ == "__main__":
    import sys

    sys.exit(main(sys.argv))