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

"""Depict molecular symmetry classes."""

import argparse
import io
import os
import sys
from pathlib import Path

from openeye import oechem, oedepict
from PIL import Image
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict molecular symmetry classes."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict"]
__SCRIPT_CATEGORIES__ = ["depiction"]


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
    )

    input_group = parser.add_argument_group("Input options")
    exclusive_group = input_group.add_mutually_exclusive_group(required=True)
    exclusive_group.add_argument(
        "--mol",
        type=str,
        metavar="MOL-FILE",
        help="input molecule file",
    )
    exclusive_group.add_argument(
        "--smiles",
        type=str,
        metavar="SMILES",
        help="input molecule SMILES",
    )

    image_group = parser.add_argument_group("Image options")
    image_group.add_argument(
        "--image",
        type=str,
        required=False,
        metavar="IMAGE-FILE",
        help="output image file (SVG, PNG) (required: %(required)s) -- if no output is provided the image will be displayed on the screen",
    )
    image_group.add_argument(
        "--width",
        type=int,
        default=800,
        help="width of output image (default: %(default)s)",
    )
    image_group.add_argument(
        "--height",
        type=int,
        default=800,
        help="height of output image (default: %(default)s)",
    )

    parser.add_argument("--help-image", action=HelpPreviewAction)
    return parser.parse_args()


def depict_symmetry(
    image: oedepict.OEImageBase,
    mol: oechem.OEMolBase,
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """
    Depict a molecule with symmetry classes highlighted.

    Perceives symmetry classes on the molecule and highlights each
    equivalence class (with more than one atom) using a distinct color.

    Args:
        image: Image to render into.
        mol: Molecule to depict.
        opts: 2D molecule display options.

    """
    oechem.OEPerceiveSymmetry(mol)

    sym_dict: dict[int, oechem.OEAtomBondSet] = {}
    for atom in mol.GetAtoms():
        sym = atom.GetSymmetryClass()
        if sym not in sym_dict:
            sym_dict[sym] = oechem.OEAtomBondSet()
        sym_dict[sym].AddAtom(atom)

    # Remove unique symmetry classes (only one atom)
    sym_dict = {sym: abset for sym, abset in sym_dict.items() if abset.NumAtoms() > 1}

    oedepict.OEPrepareDepiction(mol)
    scale = oedepict.OEGetMoleculeScale(mol, opts)
    opts.SetScale(scale * 0.9)
    disp = oedepict.OE2DMolDisplay(mol, opts)

    colors = list(oechem.OEGetContrastColors())
    for (_sym, atom_set), color in zip(sym_dict.items(), colors, strict=False):
        oedepict.OEAddHighlighting(
            disp, color, oedepict.OEHighlightStyle_BallAndStick, atom_set
        )

    oedepict.OERenderMolecule(image, disp)


def main() -> int:
    """Depict molecular symmetry classes."""
    args = parse_args()

    _check_image_file(args)

    # initialize molecule
    mol = oechem.OEGraphMol()
    if args.mol:
        ifs = oechem.oemolistream()
        if not ifs.open(args.mol):
            oechem.OEThrow.Fatal("Cannot open input file!")
        if not oechem.OEReadMolecule(ifs, mol):
            oechem.OEThrow.Fatal("Cannot read input file!")
    elif args.smiles:
        if not oechem.OESmilesToMol(mol, args.smiles):
            oechem.OEThrow.Fatal("Cannot parse SMILES!")

    # create image
    width, height = args.width, args.height
    image = oedepict.OEImage(width, height)

    # setup depiction options
    opts = oedepict.OE2DMolDisplayOptions(width, height, oedepict.OEScale_AutoScale)
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)

    # depict molecule with symmetry highlighting
    depict_symmetry(image, mol, opts)

    if args.image:
        oedepict.OEWriteImage(args.image, image)
    else:
        _img = Image.open(io.BytesIO(oedepict.OEWriteImageToBytes("png", image)))
        _img.show()

    return os.EX_OK


def _check_image_file(args: argparse.Namespace) -> None:
    # script will terminate if there is some issues
    if not args.image:
        # image will be displayed on the screen
        return
    ext = Path(args.image).suffix[1:].upper()
    if not oedepict.OEIsRegisteredImageFile(ext):
        oechem.OEThrow.Fatal("Unknown image output type!")

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


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