#!/opt/anaconda1anaconda2anaconda3/bin/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.

#############################################################################
# generate an exact/SSS SMARTS query from a structure exemplar
#############################################################################

import sys
from openeye import oechem


def ApplySMARTSConstraints(tag, mol, atomFLAGS, bondFLAGS):
    for ats in mol.GetAtoms():
        smarts = oechem.OECreateAtomSmartsString(ats, atomFLAGS)
        if smarts:
            ats.SetStringData(tag, smarts)
    for bds in mol.GetBonds():
        smarts = oechem.OECreateBondSmartsString(bds, bondFLAGS)
        if smarts:
            bds.SetStringData(tag, smarts)


def main(argv=[__name__]):
    itf = oechem.OEInterface(InterfaceData, argv)

    if itf.GetBool("-exact") and itf.GetBool("-sss"):
        oechem.OEThrow.Warning("option -exact in use, overrides -sss")

    atomFLAGS = oechem.OESMARTSAtomFlag_AtomExact
    bondFLAGS = oechem.OESMARTSBondFlag_BondExact
    if not itf.GetBool("-exact") and itf.GetBool("-sss"):
        # generalize query for SSS mappings
        atomFLAGS = (
            oechem.OESMARTSAtomFlag_AtomType
            | oechem.OESMARTSAtomFlag_IgnoreArom
        )
        bondFLAGS = oechem.OESMARTSBondFlag_AnyBondType

    ifs = oechem.oemolistream()
    if not ifs.open(itf.GetString("-in")):
        oechem.OEThrow.Fatal("Unable to open {} for input".format(itf.GetString("-in")))

    ofs = None
    if itf.HasString("-out"):
        ofs = oechem.oeofstream()
        if not ofs.open(itf.GetString("-out")):
            oechem.OEThrow.Fatal("Unable to open {} for output".format(itf.GetString("-out")))

    # option flags
    smartsTAG = "makeSMARTS"

    recno = -1
    for mol in ifs.GetOEMols():
        ++recno
        ApplySMARTSConstraints(smartsTAG, mol, atomFLAGS, bondFLAGS)
        smarts = oechem.OECreateSmartsString(mol, smartsTAG)
        if not smarts:
            oechem.OEThrow("{}: error generating SMARTS pattern", recno)
            continue

        if ofs is None:
            print(smarts)
        else:
            # output file specified
            ofs.write("{}\n".format(smarts))

    if ofs is not None:
        ofs.close()

InterfaceData = """
!BRIEF -in mol_exemplar_file [-out <output>]
    !CATEGORY I/O
        !PARAMETER -in 1
          !ALIAS -i
          !TYPE string
          !REQUIRED true
          !BRIEF Input file name
        !END
        !PARAMETER -out 2
          !ALIAS -o
          !TYPE string
          !REQUIRED false
          !BRIEF Output file name to capture the SMARTS pattern (if successful)
        !END
    !END
    !CATEGORY Options
        !PARAMETER -exact 1
          !TYPE bool
          !REQUIRED false
          !DEFAULT false
          !BRIEF generate an exact SMARTS matching pattern from the input exemplar (overrides -sss)
        !END
        !PARAMETER -sss 2
          !TYPE bool
          !REQUIRED false
          !DEFAULT false
          !BRIEF generate a substructure generalized SMARTS matching pattern from the input exemplar
        !END
    !END
!END
"""

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