#!/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 oedepict
from openeye import oegrapheme


# @ <SNIPPET-2DPROPMAP-PCHARGE-MULTI>
def SetPartialCharge(mol, tagname, minvalue, maxvalue):
    oechem.OEMMFFAtomTypes(mol)
    oechem.OEMMFF94PartialCharges(mol)

    tag = oechem.OEGetTag(tagname)
    for atom in mol.GetAtoms():
        charge = atom.GetPartialCharge()
        atom.SetData(tag, charge)
        minvalue = min(minvalue, charge)
        maxvalue = max(maxvalue, charge)

    return minvalue, maxvalue


smiles = ["c1ccnc(c1)[N+](=O)[O-]",
          "c1ccccc1F",
          "c1ccc(cc1)S(=O)(=O)[O-]"]

minvalue = float("inf")
maxvalue = float("-inf")
tagname = "PartialCharge"

molecules = []
for smi in smiles:
    mol = oechem.OEGraphMol()
    oechem.OESmilesToMol(mol, smi)
    oedepict.OEPrepareDepiction(mol)
    minvalue, maxvalue = SetPartialCharge(mol, tagname, minvalue, maxvalue)
    molecules.append(oechem.OEGraphMol(mol))

width, height = 750, 250
image = oedepict.OEImage(width, height)

rows, cols = 1, 3
grid = oedepict.OEImageGrid(image, rows, cols)

opts = oedepict.OE2DMolDisplayOptions(grid.GetCellWidth(), grid.GetCellHeight(),
                                      oedepict.OEScale_Default)
opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)
opts.SetTitleLocation(oedepict.OETitleLocation_Hidden)

propmap = oegrapheme.OE2DPropMap(opts.GetBackgroundColor())
propmap.SetNegativeColor(oechem.OEDarkRed)
propmap.SetPositiveColor(oechem.OEDarkBlue)
propmap.SetLegendLocation(oegrapheme.OELegendLocation_Left)

propmap.SetMinValue(minvalue)
propmap.SetMaxValue(maxvalue)

for cell, mol in zip(grid.GetCells(), molecules):
    disp = oedepict.OE2DMolDisplay(mol, opts)
    propmap.Render(disp, tagname)

    oedepict.OERenderMolecule(cell, disp)

oedepict.OEWriteImage("Draw2DPropMapPartialChargeMulti.png", image)
# @ </SNIPPET-2DPROPMAP-PCHARGE-MULTI>
oedepict.OEWriteImage("Draw2DPropMapPartialChargeMulti.pdf", image)
