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

"""Calculates XLogP and visualizes fragment contributions of a molecule."""

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, oegrapheme, oemedchem, oemolprop, oequacpac
from PIL import Image
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict XLogP of molecule (fragment base)."
__SCRIPT_TOOLKITS__ = [
    "oechem",
    "oedepict",
    "oegrapheme",
    "oemolprop",
    "oequacpac",
    "oemedchem",
]
__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]",
    )

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

    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:
    """Visualizes XLogP atom contributions."""
    args = parse_options()

    _check_image_file(args)
    mol = _get_molecule(args)

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

    opts = oedepict.OE2DMolDisplayOptions(
        args.width, args.height, oedepict.OEScale_AutoScale
    )
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)

    frag_func = _get_fragmentation_function(args.frag_type)

    oedepict.OEPrepareDepiction(mol)
    depict_molecule_fragment_xlogp(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 set_atom_properties(mol: oechem.OEMolBase, data_tag: int) -> None:
    """Attach the XLogP atom contribution to each atom with the given tag."""
    oequacpac.OERemoveFormalCharge(mol)

    atom_values = oechem.OEFloatArray(mol.GetMaxAtomIdx())
    xlogp = oemolprop.OEGetXLogP(mol, atom_values)

    mol.SetTitle(f"{mol.GetTitle()} -- OEXLogP = {xlogp:.2f}")

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


def fragment_molecule(
    mol: oechem.OEMolBase,
    frag_func: Callable[[oechem.OEMolBase], Iterator[oechem.OEAtomBondSet]],
    group_tag: int,
) -> None:
    """Fragments the molecule and stores each fragment as a group on the molecule."""
    for frag in frag_func(mol):
        atoms = oechem.OEAtomVector()
        for atom in frag.GetAtoms():
            atoms.append(atom)
        bonds = oechem.OEBondVector()
        for bond in frag.GetBonds():
            bonds.append(bond)

        mol.NewGroup(group_tag, atoms, bonds)  # type: ignore[attr-defined]


def set_fragment_properties(
    mol: oechem.OEMolBase, data_tag: int, group_tag: int
) -> tuple[float, float]:
    """Calculate the fragment contribution based on attached atom properties  for pre-generated fragments."""
    min_value, max_value = float("inf"), float("-inf")

    for group in mol.GetGroups(oechem.OEHasGroupType(group_tag)):
        sum_prop = 0.0
        for atom in group.GetAtoms():
            sum_prop += atom.GetData(data_tag)
        group.SetData(data_tag, sum_prop)

        min_value = min(min_value, sum_prop)
        max_value = max(max_value, sum_prop)

    return min_value, max_value


def depict_molecule_fragment_xlogp(
    image: oedepict.OEImage,
    mol: oechem.OEMolBase,
    frag_func: Callable[[oechem.OEMolBase], Iterator[oechem.OEAtomBondSet]],
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """Generate an image of a molecule depicting the fragment contribution of XLogP."""
    # calculate atom contributions of XLogP
    data_tag: int = oechem.OEGetTag("XLogP")
    set_atom_properties(mol, data_tag)

    # fragment molecule
    group_tag: int = oechem.OEGetTag("fragment")
    fragment_molecule(mol, frag_func, group_tag)

    # calculate fragment contributions
    min_value, max_value = set_fragment_properties(mol, data_tag, group_tag)

    # initialize color gradient
    lightgrey = oechem.OEColor(240, 240, 240)
    color_gradient = oechem.OELinearColorGradient(oechem.OEColorStop(0.0, lightgrey))
    color_gradient.AddStop(oechem.OEColorStop(min_value, oechem.OEDarkGreen))
    color_gradient.AddStop(oechem.OEColorStop(max_value, oechem.OEDarkPurple))

    # generate image frames
    image_width, image_height = image.GetWidth(), image.GetHeight()
    mol_frame = oedepict.OEImageFrame(
        image, image_width, image_height * 0.8, oedepict.OE2DPoint(0.0, 0.0)
    )
    color_frame = oedepict.OEImageFrame(
        image,
        image_width,
        image_height * 0.2,
        oedepict.OE2DPoint(0.0, image_height * 0.8),
    )

    # initialize molecule display
    opts.SetDimensions(
        mol_frame.GetWidth(), mol_frame.GetHeight(), oedepict.OEScale_AutoScale
    )
    disp = oedepict.OE2DMolDisplay(mol, opts)

    # initialize highlighting style
    highlight = oedepict.OEHighlightByLasso(oechem.OEWhite)
    highlight.SetConsiderAtomLabelBoundingBox(True)

    color_gradient_opts = oegrapheme.OEColorGradientDisplayOptions()

    for group in mol.GetGroups(oechem.OEHasGroupType(group_tag)):
        group_value = group.GetData(data_tag)
        color_gradient_opts.AddMarkedValue(group_value)

        # depict fragment contribution
        color = color_gradient.GetColorAt(group_value)
        highlight.SetColor(color)

        ab_set = oechem.OEAtomBondSet(group.GetAtoms(), group.GetBonds())
        oedepict.OEAddHighlighting(disp, highlight, ab_set)

    # render molecule and color gradient
    oedepict.OERenderMolecule(mol_frame, disp)
    oegrapheme.OEDrawColorGradient(color_frame, color_gradient, color_gradient_opts)


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