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

"""Depicts 2D molecule similarity using fingerprint overlaps."""

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

import numpy as np
from openeye import oechem, oedepict, oegrapheme, oegraphsim
from PIL import Image
from rich.console import Console
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict 2D molecule similarity using fingerprint overlaps"
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme", "oegraphsim"]
__SCRIPT_KEYWORDS__ = ["similarity", "fingerprints", "depiction"]
__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]",
    )

    io_group = parser.add_argument_group("Input/output options")
    io_group.add_argument(
        "--query",
        type=str,
        metavar="MOL-FILE",
        required=True,
        help="input query molecule file",
    )
    io_group.add_argument(
        "--target",
        type=str,
        metavar="MOL-FILE",
        required=True,
        help="input target molecule file",
    )
    io_group.add_argument(
        "--image",
        type=str,
        metavar="IMAGE-FILE",
        required=False,
        help="output image file (PNG, SVG) (required: %(required)s) -- if no output is provided the image will be displayed on the  screen",
    )

    fingerprint_group = parser.add_argument_group("Fingerprint options")
    fingerprint_group.add_argument(
        "--fp-type",
        "--fingerprint-type",
        type=FingerprintType,
        default=FingerprintType.Tree,
        choices=list(FingerprintType),
        help="type of fingerprint to use for similarity calculation (default: %(default)s)",
    )

    depict_group = parser.add_argument_group("Depiction options")
    depict_group.add_argument(
        "--width",
        type=float,
        default=800.0,
        help="image width (default: %(default)s)",
    )
    depict_group.add_argument(
        "--height",
        type=float,
        default=400.0,
        help="image height (default: %(default)s)",
    )

    parser.add_argument("--help-image", action=HelpPreviewAction)
    parser.add_argument(
        "--save-console-svg",
        default=False,
        action="store_true",
        help=f"run command and capture console output in {__SCRIPT_NAME__}.svg file",
    )
    return parser.parse_args()


def main() -> int:
    """Depict molecule similarity."""
    args = parse_options()
    console = Console(record=args.save_console_svg)

    _check_image_file(args)

    query_mol = get_molecule(args.query)
    target_mol = get_molecule(args.target)

    # get fingerprint type

    fp_type = _get_fingerprint_type(args.fp_type)
    console.print(
        f"Using fingerprint type:{fp_type.GetFPTypeString()}", highlight=False
    )

    # create image

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

    # setup depiction options

    opts = oedepict.OE2DMolDisplayOptions(
        args.width, args.height, oedepict.OEScale_AutoScale
    )
    opts.SetBondWidthScaling(True)

    # depict molecules with overlaps

    depict_molecule_overlaps(image, query_mol, target_mol, fp_type, opts)

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

    if args.save_console_svg:
        console.save_svg(f"{__SCRIPT_NAME__}.svg", title="output")
    return os.EX_OK


def _set_fingerprint_similarity(
    query_mol: oechem.OEMolBase,
    target_mol: oechem.OEMolBase,
    fp_type: oegraphsim.OEFPTypeBase,
    tag: int,
    max_value: int = 0,
) -> int:
    """
    Calculate fingerprint overlap and store per-bond overlap counts.

    Computes the fingerprint overlap between query and target molecules and
    stores the number of overlapping substructure matches on each bond.
    """
    query_bonds = np.zeros(query_mol.GetMaxBondIdx(), dtype=np.uint32)
    target_bonds = np.zeros(target_mol.GetMaxBondIdx(), dtype=np.uint32)

    for match in oegraphsim.OEGetFPOverlap(query_mol, target_mol, fp_type):
        for bond in match.GetPatternBonds():
            query_bonds[bond.GetIdx()] += 1
        for bond in match.GetTargetBonds():
            target_bonds[bond.GetIdx()] += 1

    max_value = max(max_value, int(np.max(query_bonds)))
    max_value = max(max_value, int(np.max(target_bonds)))

    for bond in query_mol.GetBonds():
        bond.SetData(tag, int(query_bonds[bond.GetIdx()]))
    for bond in target_mol.GetBonds():
        bond.SetData(tag, int(target_bonds[bond.GetIdx()]))

    return max_value


class ColorBondByOverlapScore(oegrapheme.OEBondGlyphBase):
    """Bond glyph that colors bonds by fingerprint overlap score."""

    def __init__(self, color_gradient: oechem.OELinearColorGradient, tag: int) -> None:
        """Initialize the glyph with a color gradient and SD tag."""
        oegrapheme.OEBondGlyphBase.__init__(self)
        self._color_gradient = color_gradient
        self._tag = tag

    def RenderGlyph(  # noqa: N802
        self, disp: oedepict.OE2DMolDisplay, bond: oechem.OEBondBase
    ) -> bool:
        """Render the bond glyph."""
        bond_disp = disp.GetBondDisplay(bond)
        if bond_disp is None or not bond_disp.IsVisible():
            return False

        if not bond.HasData(self._tag):
            return False

        linewidth = disp.GetScale() / 3.0
        color = self._color_gradient.GetColorAt(bond.GetData(self._tag))
        pen = oedepict.OEPen(color, color, oedepict.OEFill_Off, linewidth)

        atom_disp_bgn = disp.GetAtomDisplay(bond.GetBgn())
        atom_disp_end = disp.GetAtomDisplay(bond.GetEnd())

        layer = disp.GetLayer(oedepict.OELayerPosition_Below)
        layer.DrawLine(atom_disp_bgn.GetCoords(), atom_disp_end.GetCoords(), pen)

        return True

    def ColorBondByOverlapScore(self) -> oegrapheme.OEBondGlyphBase:  # noqa: N802
        """Create a copy of this glyph."""
        return ColorBondByOverlapScore(self._color_gradient, self._tag).__disown__()


def depict_molecule_overlaps(
    image: oedepict.OEImageBase,
    query_mol: oechem.OEMolBase,
    target_mol: oechem.OEMolBase,
    fp_type: oegraphsim.OEFPTypeBase,
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """
    Depict query and target molecules with fingerprint overlap coloring.

    Renders query and target molecules side by side in a grid, with bonds
    colored by fingerprint overlap score and a Tanimoto similarity score
    displayed below.
    """
    tag = oechem.OEGetTag("fp-overlap")
    max_value: int = _set_fingerprint_similarity(query_mol, target_mol, fp_type, tag)

    color_gradient = oechem.OELinearColorGradient()
    color_gradient.AddStop(oechem.OEColorStop(0.0, oechem.OEPinkTint))
    color_gradient.AddStop(oechem.OEColorStop(1.0, oechem.OEYellow))
    color_gradient.AddStop(oechem.OEColorStop(max_value, oechem.OEDarkGreen))
    bond_glyph = ColorBondByOverlapScore(color_gradient, tag)

    oedepict.OEPrepareDepiction(query_mol)
    overlaps = oegraphsim.OEGetFPOverlap(query_mol, target_mol, fp_type)
    oedepict.OEPrepareMultiAlignedDepiction(target_mol, query_mol, overlaps)

    grid = oedepict.OEImageGrid(image, 1, 2)
    grid.SetMargin(oedepict.OEMargin_Bottom, 10)
    opts.SetDimensions(
        grid.GetCellWidth(), grid.GetCellHeight(), oedepict.OEScale_AutoScale
    )
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)

    mol_scale = min(
        oedepict.OEGetMoleculeScale(query_mol, opts),
        oedepict.OEGetMoleculeScale(target_mol, opts),
    )
    opts.SetScale(mol_scale)

    query_disp = oedepict.OE2DMolDisplay(query_mol, opts)
    oegrapheme.OEAddGlyph(query_disp, bond_glyph, oechem.IsTrueBond())
    oedepict.OERenderMolecule(grid.GetCell(1, 1), query_disp)

    target_disp = oedepict.OE2DMolDisplay(target_mol, opts)
    oegrapheme.OEAddGlyph(target_disp, bond_glyph, oechem.IsTrueBond())
    oedepict.OERenderMolecule(grid.GetCell(1, 2), target_disp)

    qfp = oegraphsim.OEFingerPrint()
    oegraphsim.OEMakeFP(qfp, query_mol, fp_type)

    tfp = oegraphsim.OEFingerPrint()
    oegraphsim.OEMakeFP(tfp, target_mol, fp_type)

    score = oegraphsim.OETanimoto(qfp, tfp)

    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Default,
        16,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )
    center = oedepict.OE2DPoint(image.GetWidth() / 2.0, image.GetHeight() - 10)
    image.DrawText(center, f"Tanimoto score = {score:.3f}", font)


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

    Tree = "tree"
    Circular = "circular"
    Path = "path"

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


def _get_fingerprint_type(
    fp_type: FingerprintType,
) -> oegraphsim.OEFPTypeBase:
    """Return the appropriate fragmentation function based on the type."""
    match fp_type:
        case FingerprintType.Tree:
            return oegraphsim.OEGetFPType(oegraphsim.OEFPType_Tree)
        case FingerprintType.Circular:
            return oegraphsim.OEGetFPType(oegraphsim.OEFPType_Circular)
        case FingerprintType.Path:
            return oegraphsim.OEGetFPType(oegraphsim.OEFPType_Path)

    msg = f"Unsupported fingerprint type: {fp_type}"
    raise ValueError(msg)


def get_molecule(filename: str) -> oechem.OEMolBase:
    """Read a molecule from a file."""
    ifs = oechem.oemolistream()
    mol = oechem.OEGraphMol()
    if not ifs.open(filename) or not oechem.OEReadMolecule(ifs, mol):
        oechem.OEThrow.Fatal(f"Cannot read molecule from file: {filename}")

    return mol


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


setattr(main, "__SCRIPT_NAME__", __SCRIPT_NAME__)
setattr(main, "__SCRIPT_DESC__", __SCRIPT_DESC__)
setattr(main, "__SCRIPT_TOOLKITS__", __SCRIPT_TOOLKITS__)
setattr(main, "__SCRIPT_KEYWORDS__", __SCRIPT_KEYWORDS__)
setattr(main, "__SCRIPT_CATEGORIES__", __SCRIPT_CATEGORIES__)

if __name__ == "__main__":
    sys.exit(main())