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

"""Calculate atom contribution of XLogP of molecules and store them in OEB file."""


import argparse
import os
import pathlib
import sys

from openeye import oechem, oemolprop, oequacpac
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Calculates XLogP of molecules."
__SCRIPT_TOOLKITS__ = ["oechem", "oemolprop", "oequacpac"]
__SCRIPT_CATEGORIES__ = ["depiction"]


def parse_options() -> argparse.Namespace:
    """Set up command line options."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
    )
    parser.add_argument("--help-image", action=HelpPreviewAction)

    io_group = parser.add_argument_group("Input/Output options")
    io_group.add_argument(
        "--in-mol",
        type=str,
        required=True,
        metavar="MOL-FILE",
        help="input molecule file (oeb, sdf)",
    )
    io_group.add_argument(
        "--out-mol",
        type=str,
        required=True,
        metavar="OEB-FILE",
        help="outout molecule file (oeb)",
    )
    io_group.add_argument(
        "--tag-name",
        type=str,
        required=True,
        metavar="STR",
        help="generic data tag for atom property",
    )

    return parser.parse_args()


def main() -> int:
    """Calculate XLogP and output to OEB."""
    args = parse_options()

    # check input/output files

    ifs = oechem.oemolistream()
    if not ifs.open(args.in_mol):
        oechem.OEThrow.Fatal("Cannot open input file!")

    ofs = oechem.oemolostream()
    if not ofs.open(args.out_mol):
        oechem.OEThrow.Fatal("Cannot open output file!")

    if ofs.GetFormat() != oechem.OEFormat_OEB:
        oechem.OEThrow.Fatal("Only works for oeb output file!")

    # read molecules and calculate XLogP

    for mol in ifs.GetOEGraphMols():
        set_xlogp(mol, args.tag_name)
        oechem.OEWriteMolecule(ofs, mol)

    return os.EX_OK


def set_xlogp(mol: oechem.OEMolBase, tag_name: str) -> None:
    """Attache the XLogP atom contribution to each atom with the given tag."""
    oequacpac.OERemoveFormalCharge(mol)

    tag = oechem.OEGetTag(tag_name)
    atom_values = oechem.OEFloatArray(mol.GetMaxAtomIdx())
    oemolprop.OEGetXLogP(mol, atom_values)

    for atom in mol.GetAtoms():
        value = atom_values[atom.GetIdx()]
        atom.SetData(tag, value)


setattr(main, "__SCRIPT_NAME__", __SCRIPT_NAME__)
setattr(main, "__SCRIPT_DESC__", __SCRIPT_DESC__)
setattr(main, "__SCRIPT_TOOLKITS__", __SCRIPT_TOOLKITS__)
setattr(main, "__SCRIPT_CATEGORIES__", __SCRIPT_CATEGORIES__)

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