#!/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 take curated Brood scores and inspect connection-table
# and molecule-building outcomes with OEMolCTBuilder, OEBroodMolBuilder, and
# OEBroodBuildResult.

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

        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 curated scores used for CT and molecule building")
        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 = BroodConnectionMolBuilderOptions()
    opts = oechem.OESimpleAppOptions(
        broodOpts,
        "BroodConnectionMolBuilder",
        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))

    # Gather scored fragment matches for the query from the Brood database.
    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)

    # Curate the raw scores before running CT and molecule building.
    scoreHitlist = oebioisostere.OEScoreHitlist(broodOpts.GetMaxScoreHits())
    packet = oebioisostere.OEBroodDBPacket()
    while reader.GetNextPacket(packet):
        scoreHitlist.AddScores(overlay.Overlay(packet))
    scoreHitlist.Build()

    if scoreHitlist.GetHitCount() == 0:
        oechem.OEThrow.Warning("No fragments survived overlay scoring; nothing to build.")
        return 0

    # Compare connection building and molecule building on the curated scores.
    ctBuilder = oebioisostere.OEMolCTBuilder(query)
    molBuilder = oebioisostere.OEBroodMolBuilder(query)

    written = 0
    failedBuild = 0
    failedCT = 0
    ctAttempts = 0

    for idx, score in enumerate(scoreHitlist.GetHits(), start=1):
        # Rebuild the candidate connection table from the scored fragment.
        ctMol = oechem.OEGraphMol(score.GetFrag())
        ctOk = ctBuilder.Build(ctMol)
        ctAttempts += 1
        if not ctOk:
            failedCT += 1

        # Build the full Brood hit directly from the same scored fragment.
        hit = oebioisostere.OEBroodHit()
        status = molBuilder.Build(score, hit)
        if status != oebioisostere.OEBroodStatusCode_Success:
            failedBuild += 1
            print("[%4d] Build(score) failed: status=%s, ct=%s"
                  % (idx, oebioisostere.OEGetBroodStatus(status), "pass" if ctOk else "fail"))
            continue

        # Collect build diagnostics for reporting with OEBroodBuildResult.
        diagMol = oechem.OEGraphMol(score.GetFrag())
        result = oebioisostere.OEBroodBuildResult()
        diagStatus = molBuilder.Build(diagMol, result)

        oechem.OEWriteMolecule(ofs, hit.GetMol())
        written += 1

        print("[%4d] ct=%s diagStatus=%-14s buildStatus=%-14s molTC=%.3f strain=%6.2f dStrain=%6.2f"
              % (idx,
                 "pass" if ctOk else "fail",
                 oebioisostere.OEGetBroodStatus(diagStatus),
                 oebioisostere.OEGetBroodStatus(result.GetBuildStatus()),
                 result.GetMolTanimotoCombo(),
                 result.GetLocalStrain(),
                 result.GetDeltaLocalStrain()))

    print("---- Summary ----")
    print("Curated scores:          %d" % scoreHitlist.GetHitCount())
    print("Hits written:            %d" % written)
    print("Mol build failures:      %d" % failedBuild)
    print("CT checks attempted:     %d" % ctAttempts)
    print("CT build failures:       %d" % failedCT)
    return 0


if __name__ == "__main__":
    import sys

    sys.exit(main(sys.argv))