#!/usr/bin/env python
# (C) 2017 OpenEye Scientific Software Inc. All rights reserved.
#
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of OpenEye products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. OpenEye 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 OpenEye 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 OpenEye be
# liable for any damages or liability in connection with the Sample Code
# or its use.

# @ <SNIPPET>
import sys
from openeye import oechem


def SplitMolComplexIter(oms, mol):
    # limiting lig, prot and wat to just near the first binding site
    # (a very common pattern)

    # @ <SNIPPET-SPLIT-MOL-COMPLEX-ITER-SIMPLE>
    # for input molecule mol ...
    for frag in oechem.OEGetMolComplexComponents(mol):
        # ...
        # @ </SNIPPET-SPLIT-MOL-COMPLEX-ITER-SIMPLE>
        oechem.OEWriteMolecule(oms, frag)

    # other api points illustrated below

    # @ <SNIPPET-SPLIT-MOL-COMPLEX-SPLIT-COVALENT>
    opt = oechem.OESplitMolComplexOptions()
    opt.SetSplitCovalent()
    # @ </SNIPPET-SPLIT-MOL-COMPLEX-SPLIT-COVALENT>

    # @ <SNIPPET-SPLIT-MOL-FILTER-SURFACE-WATERS>
    site = 1
    surfaceWaters = True
    opt.ResetFilters(site, surfaceWaters)
    opt.SetMaxSurfaceWaterDist(8.0)
    # @ </SNIPPET-SPLIT-MOL-FILTER-SURFACE-WATERS>


def SplitMolComplexIterLigName(oms, mol, ligName):
    # the ligand has been identified by name

    # @ <SNIPPET-SPLIT-MOL-COMPLEX-ITER-LIGNAME>
    # for input molecule mol and string ligName ...
    opt = oechem.OESplitMolComplexOptions(ligName)

    for frag in oechem.OEGetMolComplexComponents(mol, opt):
        # ...
        # @ </SNIPPET-SPLIT-MOL-COMPLEX-ITER-LIGNAME>
        oechem.OEWriteMolecule(oms, frag)


def SplitMolComplexIterAllSites(oms, mol):
    # @ <SNIPPET-SPLIT-MOL-COMPLEX-ITER-ALL-SITES>
    opt = oechem.OESplitMolComplexOptions()
    allSites = 0
    opt.ResetFilters(allSites)

    for frag in oechem.OEGetMolComplexComponents(mol, opt):
        # ...
        # @ </SNIPPET-SPLIT-MOL-COMPLEX-ITER-ALL-SITES>
        oechem.OEWriteMolecule(oms, frag)


def SplitMolComplexIterAllSitesLigName(oms, mol, ligName):
    # @ <SNIPPET-SPLIT-MOL-COMPLEX-ITER-LIGNAME-AND-ALL-SITES>
    opt = oechem.OESplitMolComplexOptions(ligName)
    allSites = 0
    opt.ResetFilters(allSites)

    for frag in oechem.OEGetMolComplexComponents(mol, opt):
        # ...
        # @ </SNIPPET-SPLIT-MOL-COMPLEX-ITER-LIGNAME-AND-ALL-SITES>
        oechem.OEWriteMolecule(oms, frag)


def SplitMolComplexIterFilter(oms, mol):
    # @ <SNIPPET-SPLIT-MOL-COMPLEX-ADD-COFACTOR>
    db = oechem.OEResidueCategoryData()
    db.AddToDB(oechem.OEResidueDatabaseCategory_Cofactor, "MTQ")

    cat = oechem.OEMolComplexCategorizer()
    cat.SetResidueCategoryData(db)

    opt = oechem.OESplitMolComplexOptions()
    opt.SetCategorizer(cat)
    # @ </SNIPPET-SPLIT-MOL-COMPLEX-ADD-COFACTOR>

    # @ <SNIPPET-SPLIT-MOL-COMPLEX-ITER-FILTER>
    # for input molecule mol and options opt ...
    for l in oechem.OEGetMolComplexComponents(mol, opt, opt.GetLigandFilter()):
        # ...
        # @ </SNIPPET-SPLIT-MOL-COMPLEX-ITER-FILTER>
        oechem.OEWriteMolecule(oms, l)

    # @ <SNIPPET-SPLIT-MOL-COMPLEX-EXAMPLE-FILTER>
    ofilter = opt.GetOtherFilter()
    # @ </SNIPPET-SPLIT-MOL-COMPLEX-EXAMPLE-FILTER>
    rs = oechem.OERoles()
    ofilter(rs)


def main(argv=[__name__]):

    itf = oechem.OEInterface(InterfaceData)

    if not oechem.OEParseCommandLine(itf, argv):
        oechem.OEThrow.Fatal("Unable to interpret command line!")

    iname = itf.GetString("-i")
    oname = itf.GetString("-o")

    ims = oechem.oemolistream()
    if not ims.open(iname):
        oechem.OEThrow.Fatal("Cannot open input file!")

    oms = oechem.oemolostream()
    if not oms.open(oname):
        oechem.OEThrow.Fatal("Cannot open output file!")

    inmol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ims, inmol):
        oechem.OEThrow.Fatal("Cannot read input file!")

    allSites = itf.GetBool("-a")
    filterLigand = itf.GetBool("-f")

    # break out each scheme to better illustrate them
    if filterLigand:
        SplitMolComplexIterFilter(oms, inmol)
    elif itf.HasString("-l"):
        ligName = itf.GetString("-l")
        if allSites:
            SplitMolComplexIterAllSitesLigName(oms, inmol, ligName)
        else:
            SplitMolComplexIterLigName(oms, inmol, ligName)
    else:
        if allSites:
            SplitMolComplexIterAllSites(oms, inmol)
        else:
            SplitMolComplexIter(oms, inmol)

    oms.close()


#############################################################################
# INTERFACE
#############################################################################

InterfaceData = '''
!BRIEF SplitMolComplexIterDoc.py <molfile> [<outfile> [ligname]]

!PARAMETER -i
  !TYPE string
  !BRIEF Input molecule (usually a pdb file)
  !REQUIRED true
  !KEYLESS 1
!END
!PARAMETER -o
  !TYPE string
  !DEFAULT splitmolcomplexiter.oeb.gz
  !BRIEF Output molecule (usually an oeb)
  !REQUIRED false
  !KEYLESS 2
!END
!PARAMETER -l
  !TYPE string
  !BRIEF Explicit ligand name
  !REQUIRED false
  !KEYLESS 3
!END
!PARAMETER -a
  !TYPE bool
  !DEFAULT false
  !BRIEF If true, show all binding sites, not just the first
  !REQUIRED false
!END
!PARAMETER -f
  !TYPE bool
  !DEFAULT false
  !BRIEF If true, use just the ligand filter
  !REQUIRED false
!END
'''


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

# @ </SNIPPET>
