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

import sys
from openeye import oechem

def main(argv=[__name__]):
         
    if len(argv) != 3:
        oechem.OEThrow.Usage("%s <infile> <outfile>" % argv[0])

    ifs = oechem.oemolistream()
    if not ifs.open(argv[1]):
        oechem.OEThrow.Fatal("Unable to open %s for reading" % argv[1])

    ofs = oechem.oemolostream()
    if not ofs.open(argv[2]):
        oechem.OEThrow.Fatal("Unable to open %s for writing" % argv[2])

    for mol in ifs.GetOEMols():
        mol.NewConf(mol)
        mol.NewConf(mol)

        molSty = oechem.OE3DMolStyle()
        molSty.SetAtomStyle(oechem.OEAtomStyle_Wireframe)
        molSty.SetAtomColorer(oechem.OEMolStyleColorer(oechem.OEBlue))
        molSty.SetSurfaceColorer(oechem.OEMolStyleColorer(oechem.OEYellow))
        molSty.SetHydrogenVisibility(oechem.OEHydrogenVisibility_Off)
        molSty.SetSurfaceType(oechem.OESurfaceType_Molecular)

        confSty = oechem.OE3DMolStyle()
        confSty.SetAtomStyle(oechem.OEAtomStyle_CPK)
        confSty.SetHydrogenVisibility(oechem.OEHydrogenVisibility_Polar)

        atomSty = oechem.OE3DMolStyle()
        atomSty.SetAtomStyle(oechem.OEAtomStyle_Stick)
        atomSty.SetAtomColorer(oechem.OEMolStyleColorer(oechem.OEWhite))
        atomSty.SetSurfaceType(oechem.OESurfaceType_Off)

        atomSty2 = oechem.OE3DMolStyle()
        atomSty2.SetAtomColorer(oechem.OEMolStyleColorer(oechem.OEPink))

        oechem.OESetStyle(mol, molSty)

        molConf1 = mol.GetConf(oechem.OEHasConfIdx(0))
        oechem.OESetStyle(molConf1, confSty)

        molAtom1 = mol.GetAtom(oechem.OEHasAtomIdx(0))
        oechem.OESetStyle(molAtom1, atomSty)

        molAtom2 = mol.GetAtom(oechem.OEHasAtomIdx(1))
        oechem.OESetStyle(molAtom2, atomSty2)

        if not oechem.OEHasStyle(mol):
            oechem.OEThrow.Fatal("OE3DMolStyle was not applied to the OEMol")

        oechem.OEWriteMolecule(ofs, mol)
        

#    for conf in mol.GetConfs():
#        tempConfSty = oechem.OE3DMolStyle()
#        tempConfSty = oechem.OEGetStyle(conf)
#        print(tempConfSty.GetString())

#    for atom in mol.GetAtoms():
#        tempAtomSty = oechem.OE3DMolStyle()
#        tempAtomSty = oechem.OEGetStyle(atom)
#        print(tempAtomSty.GetString())

#    tempMolSty = oechem.OE3DMolStyle()
#    tempMolSty = oechem.OEGetStyle(mol)
#    print(tempMolSty.GetString())

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

#This is a complicated example of adding multiple OE3DMolStyles to the connfomer, atoms, and molecule
#The mol will have a style that includes: Wireframe, blue, yellow molecular surface, no hydrogens
#The first conf will have a style that indludes: CPK, blue, yellow molecular surface, polar hydrogens
#The second conf will have the same style of the mol
#The first atom on the first conf will have a style that includes: Stick, white, no surface, polar hydrogens
#The first atom on the second conf will have a style that includes: Stick, white, no surface, no hydrogens
#The second atom on the first conf will have a style that includes: CPK, pink, blue molecular surface, polar H
#This shows the hierarchy of how the stle moves from mol to conf to atom
#The top level is always the mol, but the style in a conf will take priority
#The atom will have the style from the mol and conf, but if a OE3DMolStyle is added to an atom then it will take priority