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

"""Visualizes the atom properties store in OEB file."""

import argparse
import enum
import io
import os
import pathlib
import re
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Visualizes the atom properties (supported image formats: svg, png)."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme"]
__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)

    # input options
    input_group = parser.add_argument_group("Input ligand-protein complex")
    input_group.add_argument(
        "--mol",
        type=str,
        required=True,
        metavar="OEB-FILE",
        help="input OEB file with atom properties to depict",
    )
    input_group.add_argument(
        "--tag-name",
        type=str,
        required=True,
        metavar="STR",
        help="generic data tag for atom property",
    )
    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=600,
        help="width of output image (default: %(default)s)",
    )
    image_group.add_argument(
        "--height",
        type=int,
        default=600,
        help="height of output image (default: %(default)s)",
    )

    depiction_group = parser.add_argument_group("Depiction options")
    depiction_group.add_argument(
        "--depiction-style",
        "--style",
        type=DepictionStyle,
        default=DepictionStyle.AtomGlyph,
        choices=list(DepictionStyle),
    )
    depiction_group.add_argument(
        "--negative-color",
        "--n-color",
        type=str,
        default="red",
        choices=[ColorParameter()],
        help="color for negative values (default: %(default)s)",
    )
    depiction_group.add_argument(
        "--positive-color",
        "--p-color",
        type=str,
        default="blue",
        choices=[ColorParameter()],
        help="color for positive values (default: %(default)s)",
    )
    return parser.parse_args()


def main() -> int:
    """Depict atom property."""
    args = parse_options()

    _check_image_file(args)
    mol = _get_molecule(args)

    # check atom properties

    tag: int = oechem.OEGetTag(args.tag_name)
    if not any(atom.HasData(tag) for atom in mol.GetAtoms()):
        oechem.OEThrow.Error(
            f"Cannot find tag {args.tag_name} on atoms of input molecule!"
        )

    # prepare depiction

    clear_coords, suppress_hydrogens = True, True
    oedepict.OEPrepareDepiction(mol, clear_coords, suppress_hydrogens)

    # create image / setup depiction options

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

    negative_color = _get_color(args, "negative_color")
    positive_color = _get_color(args, "positive_color")

    depict_atom_property(
        image,
        mol,
        opts,
        args.tag_name,
        (negative_color, positive_color),
        args.depiction_style,
    )

    oedepict.OEDrawCurvedBorder(image, oedepict.OELightGreyPen, 10.0)

    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_atom_property(
    image: oedepict.OEImageBase,
    mol: oechem.OEMolBase,
    opts: oedepict.OE2DMolDisplayOptions,
    tag_name: str,
    colors: tuple[oechem.OEColor, oechem.OEColor],
    style: str,
) -> None:
    """Depicts atom property using various depiction styles."""
    main_width, main_height = image.GetWidth(), image.GetHeight() * 0.9
    color_width, color_height = image.GetWidth(), image.GetHeight() * 0.1

    main_frame = oedepict.OEImageFrame(
        image, main_width, main_height, oedepict.OE2DPoint(0.0, 0.0)
    )
    color_frame = oedepict.OEImageFrame(
        image, color_width, color_height, oedepict.OE2DPoint(0.0, main_height)
    )

    opts.SetDimensions(main_width, main_height, oedepict.OEScale_AutoScale)
    opts.SetScale(oegrapheme.OEGetMoleculeSurfaceScale(mol, opts))

    int_tag = oechem.OEGetTag(tag_name)
    color_gradient = get_color_gradient(mol, int_tag, colors)

    disp = oedepict.OE2DMolDisplay(mol, opts)

    match style:
        case DepictionStyle.AtomGlyph:
            depict_atom_property_atom_glyph(disp, tag_name, color_gradient)
        case DepictionStyle.PropertyMap:
            depict_atom_property_property_map(disp, tag_name, colors)
        case DepictionStyle.MoleculeSurface:
            depict_atom_property_molecule_surface(disp, tag_name, color_gradient)

    oedepict.OERenderMolecule(main_frame, disp)
    oegrapheme.OEDrawColorGradient(color_frame, color_gradient)

    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Default,
        14,
        oedepict.OEAlignment_Left,
        oechem.OEBlack,
    )
    color_frame.DrawText(oedepict.OE2DPoint(10.0, -10.0), tag_name, font)


class ColorParameter:  # noqa: PLW1641
    """Utility class to handle color parameter."""

    def __init__(self) -> None:  # noqa: D107
        self._build_on_colors = ["red", "blue", "green", "'#rrggbb'"]

    def __repr__(self) -> str:  # noqa: D105
        return ",".join(self._build_on_colors)

    def __eq__(self, param: object) -> bool:  # noqa: D105
        if not isinstance(param, str):
            return False
        if param in ["red", "blue", "green"]:
            return True
        return re.match("^#[0-9a-fA-F]{6}", param) is not None


class DepictionStyle(enum.Enum):
    """Property depiction sty;e."""

    AtomGlyph = "atom-glyph"
    PropertyMap = "property-map"
    MoleculeSurface = "molecule-surface"

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


def get_color_gradient(
    mol: oechem.OEMolBase, tag: int, colors: tuple[oechem.OEColor, oechem.OEColor]
) -> oechem.OELinearColorGradient:
    """Generate color gradient."""
    min_value = min(
        (atom.GetData(tag) for atom in mol.GetAtoms() if atom.HasData(tag)),
        default=float("inf"),
    )
    max_value = max(
        (atom.GetData(tag) for atom in mol.GetAtoms() if atom.HasData(tag)),
        default=float("-inf"),
    )

    color_gradient = oechem.OELinearColorGradient(
        oechem.OEColorStop(0.0, oechem.OEWhite)
    )
    if min_value < 0.0:
        color_gradient.AddStop(oechem.OEColorStop(min_value, colors[0]))
    if max_value > 0.0:
        color_gradient.AddStop(oechem.OEColorStop(max_value, colors[1]))

    return color_gradient


def depict_atom_property_property_map(
    disp: oedepict.OE2DMolDisplay,
    tag_name: str,
    colors: tuple[oechem.OEColor, oechem.OEColor],
) -> None:
    """Depicts atom property using property map style."""
    opts = disp.GetOptions()
    prop_map = oegrapheme.OE2DPropMap(opts.GetBackgroundColor())
    prop_map.SetLegendLocation(oegrapheme.OELegendLocation_Hidden)
    prop_map.SetNegativeColor(colors[0])
    prop_map.SetPositiveColor(colors[1])
    prop_map.Render(disp, tag_name)


def depict_atom_property_atom_glyph(
    disp: oedepict.OE2DMolDisplay,
    tag_name: str,
    color_gradient: oechem.OELinearColorGradient,
) -> None:
    """Depicts atom property using atom glyph style."""
    tag = oechem.OEGetTag(tag_name)
    mol = disp.GetMolecule()

    for atom in mol.GetAtoms():
        if atom.HasData(tag):
            value = atom.GetDoubleData(tag)
            color = color_gradient.GetColorAt(value)
            pen = oedepict.OEPen(color, color, oedepict.OEFill_Off, 3.0)
            glyph = oegrapheme.OEAtomGlyphCircle(
                pen, oegrapheme.OECircleStyle_Default, 1.2
            )
            oegrapheme.OEAddGlyph(disp, glyph, oechem.OEHasAtomIdx(atom.GetIdx()))


def depict_atom_property_molecule_surface(
    disp: oedepict.OE2DMolDisplay,
    tag_name: str,
    color_gradient: oechem.OELinearColorGradient,
) -> None:
    """Depicts atom property using molecule surface style."""
    tag = oechem.OEGetTag(tag_name)
    mol = disp.GetMolecule()

    for atom in mol.GetAtoms():
        if atom.HasData(tag):
            value = atom.GetDoubleData(tag)
            color = color_gradient.GetColorAt(value)
            pen = oedepict.OEPen(color, color, oedepict.OEFill_Off, 4.0)
            oegrapheme.OESetSurfaceArcFxn(mol, atom, oegrapheme.OEDefaultArcFxn(pen))

    oegrapheme.OEDraw2DSurface(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 = pathlib.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!")

    if ifs.GetFormat() != oechem.OEFormat_OEB:
        oechem.OEThrow.Fatal("Expected OEB 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


def _get_color(args: argparse.Namespace, param_name: str) -> oechem.OEColor:
    color_name = getattr(args, param_name, "xffffff")
    color_dict = {"red": oechem.OERed, "blue": oechem.OEBlue, "green": oechem.OEGreen}
    if color_name in color_dict:
        return color_dict[color_name]

    # color_name has format
    color = oechem.OEColor()
    color.SetText(color_name)
    return color


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