#!/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 partial charges of molecules and store them in OEB file."""

import argparse
import os
import pathlib
import sys

from openeye import oechem
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Calculates partial charge of molecules."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_CATEGORIES__ = ["cheminformatics"]


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 partial charges 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 partial charge

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

    return os.EX_OK


def set_partial_charge(mol: oechem.OEMolBase, tag_name: str) -> None:
    """Attache the partial change to each atom with the given tag."""
    oechem.OEMMFFAtomTypes(mol)
    oechem.OEMMFF94PartialCharges(mol)

    tag = oechem.OEGetTag(tag_name)
    for atom in mol.GetAtoms():
        atom.SetData(tag, atom.GetPartialCharge())


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())