#!/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.  CADENCE 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 shape and color overlap between a 3D reference and pre-aligned fit molecules."""

import argparse
import os
import pathlib
import sys

import numpy as np
import rich.console
from openeye import oechem, oedepict, oegrapheme, oegraphsim, oeshape
from rich.progress import track
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = (
    "Depict shape and color overlap between a 3D reference and fit molecules"
)
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme", "oegraphsim", "oeshape"]
__SCRIPT_CATEGORIES__ = ["visualization"]


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(
        "--mol",
        metavar="MOL-FILE",
        type=str,
        required=True,
        help="input molecule file (first molecule is the reference)",
    )
    io_group.add_argument(
        "--report",
        metavar="REPORT-FILE",
        type=str,
        required=True,
        help="output report file (.pdf)",
    )

    gen_group = parser.add_argument_group("General options")
    gen_group.add_argument(
        "--max-hits",
        type=int,
        default=0,
        metavar="N",
        help="maximum number of hits to depict; 0 means no limit (default: %(default)s)",
    )
    gen_group.add_argument(
        "--depict-sim",
        action="store_true",
        help="calculate and depict 2D molecule similarity (default: %(default)s)",
    )

    report_group = parser.add_argument_group("Report options")
    report_group.add_argument(
        "--page-by-page",
        action="store_true",
        help="write pages of report to separate numbered image files (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 shape and color overlap in a PDF report."""
    args = parse_options()

    console = rich.console.Console(record=args.save_console_svg, highlight=False)

    mol_path = pathlib.Path(args.mol)
    if not mol_path.exists():
        console.print(f"[red]Error: Cannot open input file '{mol_path}'![/red]")
        return os.EX_NOINPUT

    ifs = oechem.oemolistream()
    if not ifs.open(str(mol_path)):
        console.print(f"[red]Error: Cannot open input file '{mol_path}'![/red]")
        return os.EX_NOINPUT

    _check_report_file(args)

    ref_mol = oechem.OEMol()
    if not oechem.OEReadMolecule(ifs, ref_mol):
        console.print("[red]Error: Cannot read reference molecule![/red]")
        return os.EX_DATAERR

    report_opts = oedepict.OEReportOptions(3, 1)
    report_opts.SetHeaderHeight(40.0)
    report_opts.SetFooterHeight(20.0)
    report = oedepict.OEReport(report_opts)

    cff = oeshape.OEColorForceField()
    cff.Init(oeshape.OEColorFFType_ImplicitMillsDean)
    cff_display = oegrapheme.OEColorForceFieldDisplay(cff)

    query_opts = _get_shape_query_display_options(args.depict_sim)

    ref_mol_displays: dict[str, oegrapheme.OEShapeQueryDisplay] = {}
    _init_multi_query_displays(ref_mol_displays, ref_mol, cff, query_opts)
    console.print(
        f"Shape overlaps will be generated for the reference with {len(ref_mol_displays)} conformations."
    )

    fit_mols: list[oechem.OEMolBase] = [
        oechem.OEGraphMol(m)
        for i, m in enumerate(ifs.GetOEGraphMols())
        if args.max_hits <= 0 or i < args.max_hits
    ]

    depict_shape_color_graphsim_overlaps(
        report, ref_mol, ref_mol_displays, fit_mols, args.depict_sim
    )

    cff_opts = oegrapheme.OEColorForceFieldLegendDisplayOptions(1, 6)
    for header in report.GetHeaders():
        oegrapheme.OEDrawColorForceFieldLegend(header, cff_display, cff_opts)
        oedepict.OEDrawCurvedBorder(header, oedepict.OELightGreyPen, 10.0)

    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Default,
        12,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )
    for idx, footer in enumerate(report.GetFooters()):
        oedepict.OEDrawTextToCenter(footer, f"- {idx + 1} -", font)

    if args.page_by_page:
        oedepict.OEWriteReportPageByPage(args.report, report)
    else:
        oedepict.OEWriteReport(args.report, report)

    console.print(
        f"[green]Report written to '{pathlib.Path(args.report).name}'[/green]"
    )
    if args.save_console_svg:
        console.save_svg(f"{__SCRIPT_NAME__}.svg", title="output")

    return os.EX_OK


def depict_shape_color_graphsim_overlaps(
    report: oedepict.OEReport,
    ref_mol: oechem.OEMCMolBase,
    ref_mol_displays: dict[str, oegrapheme.OEShapeQueryDisplay],
    fit_mols: list[oechem.OEMolBase],
    depict_sim: bool,
) -> None:
    """Depict shape, color, and optionally 2D similarities for fit molecules."""
    fp_tag = oechem.OEGetTag("fp_overlap")
    fp_type: oegraphsim.OEFPTypeBase | None = None
    ref_mol_fp: oegraphsim.OEFingerPrint | None = None
    bond_glyph: oegrapheme.OEBondGlyphBase | None = None
    if depict_sim and ref_mol.GetMaxBondIdx() > 0:
        fp_type = oegraphsim.OEGetFPType(oegraphsim.OEFPType_Tree)
        ref_mol_fp = oegraphsim.OEFingerPrint()
        oegraphsim.OEMakeFP(ref_mol_fp, ref_mol, fp_type)
        if fp_type:
            fp_color_gradient = _get_fingerprint_color_gradient(
                _get_max_bond_self_similarity_score(ref_mol, fp_type)
            )
        bond_glyph = ColorBondByOverlapScore(fp_color_gradient, fp_tag)

    shape_opts = _get_shape_overlap_display_options()
    color_opts = _get_color_overlap_display_options()

    fit_table_opts = _get_fit_table_options(depict_sim)
    ref_table_opts = _get_ref_table_options()

    score_font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Bold,
        9,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )

    for fit_idx, fit_mol in enumerate(
        track(fit_mols, description="Generating overlays"), start=1
    ):
        if not oechem.OEHasSDData(fit_mol, "ROCS_ShapeQuery"):
            oechem.OEThrow.Warning(
                f"Shape query reference is not available for molecule '{fit_mol.GetTitle()}'"
            )
            continue

        ref_title = oechem.OEGetSDData(fit_mol, "ROCS_ShapeQuery")

        if ref_title not in ref_mol_displays:
            oechem.OEThrow.Warning(
                f"Shape query reference '{ref_title}' is not valid for molecule '{fit_mol.GetTitle()}'"
            )
            continue
        ref_disp = ref_mol_displays[ref_title]

        cell = report.NewCell()
        fittable = oedepict.OEImageTable(cell, fit_table_opts)

        fittable.DrawText(fittable.GetCell(1, 1), f"Hit: {fit_mol.GetTitle()}")

        ref_table = oedepict.OEImageTable(fittable.GetCell(2, 1), ref_table_opts)

        ref_table.DrawText(ref_table.GetCell(1, 1), f"Rank: {fit_idx}")
        _render_score(
            ref_table.GetCell(2, 1),
            fit_mol,
            "ROCS_TanimotoCombo",
            "Tanimoto Combo",
            score_font,
        )

        sim_score: float | None = None
        if fp_type and ref_mol_fp and fp_tag:
            sim_score = _calc_fingerprint_similarity(
                ref_mol, ref_mol_fp, fit_mol, fp_type, fp_tag
            )
        _render_score_radial(ref_table.GetCell(3, 1), fit_mol, sim_score)

        oegrapheme.OERenderShapeQuery(ref_table.GetCell(4, 1), ref_disp)
        ref_table.DrawText(ref_table.GetCell(5, 1), f"query : {ref_title}")

        overlap_disp = oegrapheme.OEShapeOverlapDisplay(
            ref_disp, fit_mol, shape_opts, color_opts
        )

        # shape overlap
        _render_score(
            fittable.GetHeaderCell(1),
            fit_mol,
            "ROCS_ShapeTanimoto",
            "Shape Tanimoto",
            score_font,
        )
        oegrapheme.OERenderShapeOverlap(fittable.GetCell(2, 2), overlap_disp)

        # color overlap
        _render_score(
            fittable.GetHeaderCell(2),
            fit_mol,
            "ROCS_ColorTanimoto",
            "Color Tanimoto",
            score_font,
        )
        oegrapheme.OERenderColorOverlap(fittable.GetCell(2, 3), overlap_disp)

        # 2D similarity
        if bond_glyph:
            sim_title = f"2D Graph Tanimoto = {sim_score:4.3f}"
            oedepict.OEDrawTextToCenter(
                fittable.GetHeaderCell(3), sim_title, score_font
            )
            _depict_molecule_similarity(
                fittable.GetCell(2, 4), fit_mol, ref_disp, bond_glyph
            )


def _init_multi_query_displays(
    query_map: dict,
    ref_mol: oechem.OEMCMolBase,
    cff: oeshape.OEColorForceField,
    query_opts: oegrapheme.OEShapeQueryDisplayOptions,
) -> int:
    """Generate shape display objects for each conformation of the reference molecule."""
    for conf in ref_mol.GetConfs():
        title = conf.GetTitle()
        query_map[title] = oegrapheme.OEShapeQueryDisplay(conf, cff, query_opts)
    return len(query_map)


def _add_common_display_options(opts: oedepict.OE2DMolDisplayOptions) -> None:
    """Set common display options shared across query, shape, and color displays."""
    opts.SetTitleLocation(oedepict.OETitleLocation_Hidden)
    opts.SetAtomLabelFontScale(1.5)
    pen = oedepict.OEPen(oechem.OEBlack, oechem.OEBlack, oedepict.OEFill_Off, 1.5)
    opts.SetDefaultBondPen(pen)


def _get_shape_query_display_options(
    depict_sim: bool,
) -> oegrapheme.OEShapeQueryDisplayOptions:
    """Create shape query display options."""
    query_opts = oegrapheme.OEShapeQueryDisplayOptions()
    _add_common_display_options(query_opts)
    arc_pen = oedepict.OEPen(oedepict.OELightGreyPen)
    query_opts.SetSurfaceArcFxn(oegrapheme.OEDefaultArcFxn(arc_pen))
    if depict_sim:
        query_opts.SetDepictOrientation(oedepict.OEDepictOrientation_Vertical)
    else:
        query_opts.SetDepictOrientation(oedepict.OEDepictOrientation_Square)
    query_opts.SetBackgroundColor(oechem.OETransparentColor)
    return query_opts


def _get_shape_overlap_display_options() -> oegrapheme.OEShapeOverlapDisplayOptions:
    """Create shape overlap display options."""
    shape_opts = oegrapheme.OEShapeOverlapDisplayOptions()
    _add_common_display_options(shape_opts)
    arc_pen = oedepict.OEPen(
        oechem.OEGrey, oechem.OEGrey, oedepict.OEFill_Off, 1.0, 0x1111
    )
    shape_opts.SetQuerySurfaceArcFxn(oegrapheme.OEDefaultArcFxn(arc_pen))
    shape_opts.SetOverlapColor(oechem.OEColor(110, 110, 190))
    shape_opts.SetOverlapDisplayStyle(
        oegrapheme.OEShapeOverlapDisplayStyle_PropertyCloud
    )
    shape_opts.SetBackgroundColor(oechem.OETransparentColor)
    return shape_opts


def _get_color_overlap_display_options() -> oegrapheme.OEColorOverlapDisplayOptions:
    """Create color overlap display options."""
    color_opts = oegrapheme.OEColorOverlapDisplayOptions()
    _add_common_display_options(color_opts)
    arc_pen = oedepict.OEPen(
        oechem.OEGrey, oechem.OEGrey, oedepict.OEFill_Off, 1.0, 0x1111
    )
    color_opts.SetQuerySurfaceArcFxn(oegrapheme.OEDefaultArcFxn(arc_pen))
    color_opts.SetBackgroundColor(oechem.OETransparentColor)
    return color_opts


def _get_fit_table_options(depict_sim: bool) -> oedepict.OEImageTableOptions:
    """Create table options for fit molecule cells."""
    rows = 2
    cols = 4 if depict_sim else 3
    table_opts = oedepict.OEImageTableOptions(
        rows, cols, oedepict.OEImageTableStyle_LightGrey
    )
    table_opts.SetHeader(True)
    table_opts.SetStubColumn(True)
    table_opts.SetRowHeights([10, 90])
    if depict_sim:
        table_opts.SetColumnWidths([20, 30, 30, 30])
    else:
        table_opts.SetColumnWidths([18, 30, 30])
    table_opts.SetCellColor(oechem.OEWhite, False)
    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Bold,
        9,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )
    table_opts.SetHeaderFont(font)
    return table_opts


def _get_ref_table_options() -> oedepict.OEImageTableOptions:
    """Create table options for the reference molecule stub column."""
    table_opts = oedepict.OEImageTableOptions(5, 1, oedepict.OEImageTableStyle_NoStyle)
    table_opts.SetHeader(False)
    table_opts.SetStubColumn(False)
    table_opts.SetRowHeights([6, 6, 25, 60, 6])
    table_opts.SetMargins(0)
    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Default,
        8,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )
    table_opts.SetCellFont(font)
    return table_opts


def _get_score(mol: oechem.OEMolBase, sd_tag: str) -> float:
    """Return an SD data value as a float, or 0.0 if not present."""
    if oechem.OEHasSDData(mol, sd_tag):
        return float(oechem.OEGetSDData(mol, sd_tag))
    return 0.0


def _render_score_radial(
    image: oedepict.OEImageBase,
    mol: oechem.OEMolBase,
    fp_score: float | None = None,
) -> None:
    """Render a radial ROCS score chart."""
    shape_score = max(min(_get_score(mol, "ROCS_ShapeTanimoto"), 1.0), 0.0)
    color_score = max(min(_get_score(mol, "ROCS_ColorTanimoto"), 1.0), 0.0)
    if shape_score > 0.0 or color_score > 0.0:
        if fp_score is None:
            scores = oechem.OEDoubleVector([shape_score, color_score])
        else:
            scores = oechem.OEDoubleVector([shape_score, color_score, fp_score])
        oegrapheme.OEDrawROCSScores(image, scores)


def _render_score(
    image: oedepict.OEImageBase,
    mol: oechem.OEMolBase,
    sd_tag: str,
    label: str,
    score_font: oedepict.OEFont,
) -> None:
    """Render a labeled score value centered in the image."""
    score = _get_score(mol, sd_tag)
    if score == 0.0:
        return
    oedepict.OEDrawTextToCenter(image, f"{label} = {score}", score_font)


# ---------------------------------------------------------------------------
# 2D similarity depiction
# ---------------------------------------------------------------------------


def _get_max_bond_self_similarity_score(
    mol: oechem.OEMolBase, fp_type: oegraphsim.OEFPTypeBase
) -> int:
    """Return the maximum bond self-similarity overlap count."""
    count_bonds = np.zeros(mol.GetMaxBondIdx(), dtype=np.uint32)
    for match in oegraphsim.OEGetFPOverlap(mol, mol, fp_type):
        for bond in match.GetPatternBonds():
            count_bonds[bond.GetIdx()] += 1
    return int(count_bonds.max())


def _calc_fingerprint_similarity(
    ref_mol: oechem.OEMolBase,
    ref_mol_fp: oegraphsim.OEFingerPrint,
    fit_mol: oechem.OEMolBase,
    fp_type: oegraphsim.OEFPTypeBase,
    tag: str,
) -> float:
    """Calculate fingerprint similarity and annotate bonds with overlap counts."""
    bond_counts = np.zeros(fit_mol.GetMaxBondIdx(), dtype=np.uint32)

    for match in oegraphsim.OEGetFPOverlap(ref_mol, fit_mol, fp_type):
        for bond in match.GetTargetBonds():
            bond_counts[bond.GetIdx()] += 1

    for bond in fit_mol.GetBonds():
        bond.SetData(tag, int(bond_counts[bond.GetIdx()]))

    fit_mol_fp = oegraphsim.OEFingerPrint()
    oegraphsim.OEMakeFP(fit_mol_fp, fit_mol, fp_type)
    return oegraphsim.OETanimoto(ref_mol_fp, fit_mol_fp)


def _get_fingerprint_color_gradient(self_score: int) -> oechem.OELinearColorGradient:
    """Create a color gradient for fingerprint overlap visualization."""
    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(self_score, oechem.OEDarkGreen))
    return color_gradient


def _depict_molecule_similarity(
    cell: oedepict.OEImageBase,
    mol_3d: oechem.OEMolBase,
    ref_disp: oegrapheme.OEShapeQueryDisplay,
    bond_glyph: oegrapheme.OEBondGlyphBase,
) -> None:
    """Align a molecule to the reference and depict 2D fingerprint similarity."""
    mol_2d = oechem.OEGraphMol(mol_3d)
    oegrapheme.OEPrepareAlignedDepictionFrom3D(mol_2d, mol_3d, ref_disp)

    width, height = cell.GetWidth(), cell.GetHeight()
    opts = oedepict.OE2DMolDisplayOptions(width, height, oedepict.OEScale_AutoScale)
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)
    opts.SetTitleLocation(oedepict.OETitleLocation_Hidden)
    opts.SetScale(oegrapheme.OEGetMoleculeSurfaceScale(mol_2d, opts))

    disp = oedepict.OE2DMolDisplay(mol_2d, opts)
    oegrapheme.OEAddGlyph(disp, bond_glyph, oechem.IsTrueBond())
    oedepict.OERenderMolecule(cell, disp, False)


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

    def __init__(self, color_gradient: oechem.OEColorGradientBase, tag: str) -> None:
        """Initialize the glyph with a color gradient and SD tag for bond data."""
        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 a bond glyph colored by the fingerprint overlap score stored in the bond data."""
        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_b = disp.GetAtomDisplay(bond.GetBgn())
        atom_disp_e = disp.GetAtomDisplay(bond.GetEnd())

        layer = disp.GetLayer(oedepict.OELayerPosition_Below)
        layer.DrawLine(atom_disp_b.GetCoords(), atom_disp_e.GetCoords(), pen)

        return True

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


def _check_report_file(args: argparse.Namespace) -> bool:
    ext = pathlib.Path(args.report).suffix[1:]
    if not oedepict.OEIsRegisteredImageFile(ext):
        oechem.OEThrow.Fatal("Unknown image outout type!")

    if not args.page_by_page and not oedepict.OEIsRegisteredMultiPageImageFile(ext):
        oechem.OEThrow.Warning("Report will be generated into separate pages!")
        args.page_by_page = True

    return True


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