#!/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 molecule with fragment highlights."""

import argparse
import enum
import io
import os
import sys
from collections.abc import Callable, Iterator
from pathlib import Path

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

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict molecule with fragment highlights."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oemedchem"]
__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",
    )

    frag_group = parser.add_argument_group("Fragmentation options")
    frag_group.add_argument(
        "--frag-type",
        "--fragmentation-type",
        type=FragmentationType,
        default=FragmentationType.FunctionalGroup,
        choices=list(FragmentationType),
    )

    image_group = parser.add_argument_group("Image options")
    image_group.add_argument(
        "--image",
        type=str,
        required=False,
        metavar="IMAGE-FILE",
        help="output image file (PNG, SVG) (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=400,
        help="height of output image (default: %(default)s)",
    )

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


def main() -> int:
    """Depict molecule with fragment highlights."""
    args = parse_args()
    _check_image_file(args)

    # initialize molecule
    mol: oechem.OEMolBase
    if args.mol:
        mol = _get_molecule(args)
    elif args.smiles:
        mol = oechem.OEGraphMol()
        if not oechem.OESmilesToMol(mol, args.smiles):
            oechem.OEThrow.Fatal("Cannot parse SMILES!")

    # initialize fragmentation function
    frag_func = _get_fragmentation_function(args.frag_type)

    # 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 fragment highlights
    oedepict.OEPrepareDepiction(mol)
    depict_molecule_with_fragment_highlights(image, mol, frag_func, 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 depict_molecule_with_fragment_highlights(
    image: oedepict.OEImageBase,
    mol: oechem.OEMolBase,
    frag_func: Callable[[oechem.OEMolBase], Iterator[oechem.OEAtomBondSet]],
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """
    Depict a molecule with its fragments highlighted.

    Each fragment returned by the fragmentation function is highlighted
    using a color gradient from yellow to brown.
    """
    frag_list = list(frag_func(mol))
    num_frags = len(frag_list)

    color_gradient = oechem.OELinearColorGradient()
    color_gradient.AddStop(oechem.OEColorStop(0, oechem.OEMediumYellow))
    color_gradient.AddStop(oechem.OEColorStop(num_frags, oechem.OEDarkBrown))

    disp = oedepict.OE2DMolDisplay(mol, opts)

    highlight = oedepict.OEHighlightByLasso(oechem.OEWhite)
    highlight.SetConsiderAtomLabelBoundingBox(True)

    for frag_idx, frag in enumerate(frag_list):
        highlight.SetColor(color_gradient.GetColorAt(frag_idx))
        oedepict.OEAddHighlighting(disp, highlight, frag)

    oedepict.OERenderMolecule(image, disp)


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!")


def _get_molecule(args: argparse.Namespace) -> oechem.OEMolBase:
    ifs = oechem.oemolistream()
    if not ifs.open(args.mol):
        oechem.OEThrow.Fatal(f"Cannot open {args.mol} input file!")

    mol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ifs, mol):
        oechem.OEThrow.Fatal(f"Cannot read molecule from {args.mol} input file!")

    return mol


class FragmentationType(enum.Enum):
    """Molecule fragmentation type."""

    FunctionalGroup = "func-group"
    RingChain = "ring-chain"
    RingLinkerSideChain = "ring-linker-sidechain"

    def __str__(self) -> str:
        """Convert to string representation."""
        return self.value


def _get_fragmentation_function(
    frag_type: FragmentationType,
) -> Callable[[oechem.OEMolBase], Iterator[oechem.OEAtomBondSet]]:
    match frag_type:
        case FragmentationType.RingChain:
            return oemedchem.OEGetRingChainFragments
        case FragmentationType.RingLinkerSideChain:
            return oemedchem.OEGetRingLinkerSideChainFragments
    return oemedchem.OEGetFuncGroupFragments


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