#!/usr/bin/env python
# (C) 2021 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.

#############################################################################################################
# This program demonstrates how to find pockets and some of its properties from a protein or DesignUnit file.
#############################################################################################################
import sys
from openeye import oechem
from openeye import oespruce


def readProteinMol(ifilename):
    if oechem.OEGetFileExtension(ifilename) == 'oedu':
        du = oechem.OEDesignUnit()
        if not oechem.OEReadDesignUnit(ifilename, du):
            oechem.OEThrow.Fatal("Unable to open %s for reading OEDesignUnit" % ifilename)
        return du
    else:
        ifs = oechem.oemolistream()
        if not ifs.open(ifilename):
            oechem.OEThrow.Fatal("Unable to open %s for reading" % ifilename)
        mol = oechem.OEGraphMol()
        oechem.OEReadMolecule(ifs, mol)
        return mol


def main(argv=None):
    if argv is None:
        argv = [__name__]
    if len(sys.argv) != 2:
        oechem.OEThrow.Usage("%s <protein or DesignUnit input file>" % argv[0])

    mol = readProteinMol(sys.argv[1])

    pockets = oespruce.OEFindPockets(mol)
    print("pockets count: %s" % len(list(pockets)))
    pockets.ToFirst()
    pocket_cntr = 0
    for pocket in pockets:
        pocket_cntr += 1
        pocket_residues = pocket.GetResidues()
        print("pocket_%s Residues count: " % pocket_cntr, len(list(pocket_residues)))
        pocket_residues.ToFirst()
        print("pocket_%s Residues: " % pocket_cntr)
        for res in pocket_residues:
            print(res)
        print("pocket_%s Surface Area: " % pocket_cntr,  pocket.GetSurfaceArea())


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