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

# generate szmap points and probe orientations at ligand atoms

import sys
from openeye import oechem
from openeye import oeszmap

InterfaceData = """
!BRIEF SzmapBestOrientations.py [-prob #.#] [-du] <designunitfile> [-o] <molfile>
!PARAMETER -du
  !TYPE string
  !BRIEF Input designunit
  !REQUIRED true
  !KEYLESS 1
!END
!PARAMETER -o
  !TYPE string
  !BRIEF Output file for points and probes molecules
  !REQUIRED true
  !KEYLESS 2
!END
!PARAMETER -prob
  !TYPE double
  !DEFAULT 0.5
  !BRIEF Cutoff for cumulative probability of probes
  !REQUIRED false
!END
"""


def GenerateSzmapProbes(oms, cumulativeProb, lig, prot):
    """
    generate multiconf probes and data-rich points at ligand coords
    @rtype : None
    @param oms: output mol stream for points and probes
    @param cumulativeProb: cumulative probability for cutoff of point set
    @param lig: mol defining coordinates for szmap calcs
    @param prot: context mol for szmap calcs (must have charges and radii)
    """

    sz = oeszmap.OESzmapEngine(prot)

    rslt = oeszmap.OESzmapResults()

    points = oechem.OEGraphMol()
    points.SetTitle("points %s" % lig.GetTitle())
    probes = oechem.OEMol()

    coord = oechem.OEFloatArray(3)

    for i, atom in enumerate(lig.GetAtoms()):
        lig.GetCoords(atom, coord)

        if not oeszmap.OEIsClashing(sz, coord):
            oeszmap.OECalcSzmapResults(rslt, sz, coord)

            rslt.PlaceNewAtom(points)

            clear = False
            rslt.PlaceProbeSet(probes, cumulativeProb, clear)

            name = oeszmap.OEGetEnsembleName(oeszmap.OEEnsemble_NeutralDiffDeltaG)
            nddG = rslt.GetEnsembleValue(oeszmap.OEEnsemble_NeutralDiffDeltaG)
            print("%2d (%7.3f, %7.3f, %7.3f): %s = %.3f"
                  % (i, coord[0], coord[1], coord[2], name, nddG))
        else:
            print("%2d (%7.3f, %7.3f, %7.3f): CLASH"
                  % (i, coord[0], coord[1], coord[2]))

    oechem.OEWriteMolecule(oms, points)
    oechem.OEWriteMolecule(oms, probes)


def main(argv=(__name__)):
    """
    the protein should have charges and radii but the ligand need not
    """
    itf = oechem.OEInterface()

    if not oechem.OEConfigure(itf, InterfaceData):
        oechem.OEThrow.Fatal("Problem configuring OEInterface!")

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

    duFile = itf.GetString("-du")
    du = oechem.OEDesignUnit()
    if not oechem.OEReadDesignUnit(duFile, du):
        oechem.OEThrow.Fatal("Unable to open %s for reading" % duFile)

    if not du.HasLigand():
        oechem.OEThrow.Fatal("Input designunit %s does not have a ligand bound" % du.GetTitle())

    lig = oechem.OEGraphMol()
    du.GetLigand(lig)
    prot = oechem.OEGraphMol()
    du.GetComponents(prot, oechem.OEDesignUnitComponents_Protein)

    outputFile = itf.GetString("-o")
    if not oechem.OEIsWriteable(outputFile):
        oechem.OEThrow.Fatal("Invalid file extension for output %s" % outputFile)

    oms = oechem.oemolostream()
    if not oms.open(outputFile):
        oechem.OEThrow.Fatal("Unable to open %s for writing" % outputFile)

    GenerateSzmapProbes(oms, itf.GetDouble("-prob"), lig, prot)


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