#!/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 polar surface area visualization."""

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

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

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict molecule with polar surface area visualization."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme", "oemolprop"]
__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 (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 polar surface area visualization."""
    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!")

    oedepict.OEPrepareDepiction(mol)

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

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

    # depict molecule with polar surface area
    depict_molecule_with_psa(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 depict_molecule_with_psa(
    image: oedepict.OEImageBase,
    mol: oechem.OEMolBase,
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """
    Depict a molecule with polar surface area visualization.

    Calculates topological PSA per atom and renders eyelash-style
    surface arcs colored by PSA contribution on the 2D depiction.

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

    """
    scale = oegrapheme.OEGetMoleculeSurfaceScale(mol, opts)
    opts.SetScale(scale)

    tag = oechem.OEGetTag("PSA")
    min_value, max_value = set_atom_properties(mol, tag, s_and_p=True)

    negative_color = oechem.OEColorStop(min_value, oechem.OEWhite)
    positive_color = oechem.OEColorStop(max_value, oechem.OEDarkBlue)
    color_gradient = oechem.OELinearColorGradient(negative_color, positive_color)

    arc_fxn = PSAArcFxn(color_gradient, tag, opts.GetDefaultBondPen())
    for atom in mol.GetAtoms():
        oegrapheme.OESetSurfaceArcFxn(mol, atom, arc_fxn)

    disp = oedepict.OE2DMolDisplay(mol, opts)
    oegrapheme.OEDraw2DSurface(disp)
    oedepict.OERenderMolecule(image, disp)


def set_atom_properties(
    mol: oechem.OEMolBase,
    tag: int,
    min_value: float = float("inf"),
    max_value: float = float("-inf"),
    s_and_p: bool = True,  # noqa: FBT002
) -> tuple[float, float]:
    """
    Calculate per-atom PSA values and store them as generic data.

    Computes the topological polar surface area contribution for each
    atom and stores it under the given tag. Updates and returns the
    running min/max values across calls.
    """
    atom_values = oechem.OEFloatArray(mol.GetMaxAtomIdx())
    psa = oemolprop.OEGet2dPSA(mol, atom_values, s_and_p)

    mol.SetTitle(f"{mol.GetTitle()}Topological Polar Surface Area = {psa:.2f}")

    for atom in mol.GetAtoms():
        val = atom_values[atom.GetIdx()]
        atom.SetData(tag, val)
        min_value = min(min_value, val)
        max_value = max(max_value, val)

    return min_value, max_value


class PSAArcFxn(oegrapheme.OESurfaceArcFxnBase):
    """Surface arc function for rendering PSA eyelash arcs."""

    def __init__(
        self,
        color_gradient: oechem.OELinearColorGradient,
        tag: int,
        pen: oedepict.OEPen,
    ) -> None:
        """Initialize arc function."""
        super().__init__()
        self._color_gradient = color_gradient
        self._tag = tag
        self._pen = pen

    def __call__(
        self,
        image: oedepict.OEImageBase,
        arc: oegrapheme.OESurfaceArc,
    ) -> bool:
        """Draw arc."""
        atom_disp = arc.GetAtomDisplay()
        if atom_disp is None or not atom_disp.IsVisible():
            return False

        atom = atom_disp.GetAtom()
        atom_psa = atom.GetData(self._tag)
        if atom_psa == 0.0:
            return True

        pen = oedepict.OEPen(self._pen)
        color = self._color_gradient.GetColorAt(atom_psa)
        pen.SetForeColor(color)

        center = arc.GetCenter()
        bgn_angle = arc.GetBgnAngle()
        end_angle = arc.GetEndAngle()
        radius = arc.GetRadius()

        edge_angle = 10.0
        pattern_direction = oegrapheme.OEPatternDirection_Outside
        pattern_angle = 10.0
        min_pattern_width_ratio = 0.05
        max_pattern_width_ratio = 0.70
        act_pattern_width_ratio = min(
            max_pattern_width_ratio, atom_psa * (max_pattern_width_ratio / 40.0)
        )
        oegrapheme.OEDrawEyelashSurfaceArc(
            image,
            center,
            bgn_angle,
            end_angle,
            radius,
            pen,
            edge_angle,
            pattern_direction,
            pattern_angle,
            min_pattern_width_ratio,
            act_pattern_width_ratio,
        )
        return True

    def CreateCopy(self):  # noqa: ANN201, N802
        """Copy constructor."""
        return PSAArcFxn(self._color_gradient, self._tag, self._pen).__disown__()


def _check_image_file(args: argparse.Namespace) -> None:
    """Validate image output file extension."""
    if args.image is None:
        return
    ext = Path(args.image).suffix[1:]
    if not oedepict.OEIsRegisteredImageFile(ext):
        oechem.OEThrow.Fatal("Unknown image type!")


def _get_molecule(args: argparse.Namespace) -> oechem.OEMolBase:
    """Return a molecule from the input arguments."""
    mol = oechem.OEGraphMol()
    if args.smiles:
        if not oechem.OESmilesToMol(mol, args.smiles):
            oechem.OEThrow.Fatal(f"Cannot parse SMILES: {args.smiles}")
    else:
        ifs = oechem.oemolistream()
        if not ifs.open(args.mol):
            oechem.OEThrow.Fatal(f"Cannot open input file: {args.mol}")
        if not oechem.OEReadMolecule(ifs, mol):
            oechem.OEThrow.Fatal(f"Cannot read molecule from {args.mol} input file!")
    return mol


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