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

"""Performs 2D fingerprint similarity search and depicts molecule similarities."""

import argparse
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = (
    "Perform 2D fingerprint similarity search and depict molecule similarities"
)
__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 options")
    io_group.add_argument(
        "--query",
        type=str,
        metavar="MOL-FILE",
        required=True,
        help="input query molecule file",
    )
    io_group.add_argument(
        "--mol",
        "--mol-file",
        type=str,
        metavar="MOL-FILE",
        required=True,
        help="input molecule database file",
    )
    io_group.add_argument(
        "--fp-database",
        type=str,
        metavar="FP-FILE",
        required=True,
        help="input fast fingerprint database file",
    )

    output_group = parser.add_argument_group("Output options")
    output_group.add_argument(
        "--report",
        type=str,
        metavar="REPORT-FILE",
        required=True,
        help="output report file (.pdf)",
    )
    output_group.add_argument(
        "--page-by-page",
        action="store_true",
        help="write pages of report to separate numbered image files (default: %(default)s)",
    )

    fp_group = parser.add_argument_group("Fingerprint database options")
    fp_group.add_argument(
        "--num-hits",
        type=int,
        default=10,
        help="number of hits to return (default: %(default)s)",
    )
    fp_group.add_argument(
        "--mem-type",
        type=str,
        default="memory-mapped",
        choices=["memory-mapped", "in-memory"],
        help="fingerprint database memory type (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:
    """Perform fingerprint similarity search and depict results."""
    args = parse_options()

    console = Console()

    _check_report_file(args)

    query_mol = get_molecule(args.query)

    # initialize databases
    mol_db = oechem.OEMolDatabase()
    if not mol_db.Open(args.mol):
        console.print(f"[red]Cannot open molecule database: {args.mol}[/red]")
        return os.EX_IOERR

    if args.mem_type == "in-memory":
        mem_type = oegraphsim.OEFastFPDatabaseMemoryType_InMemory
    else:
        mem_type = oegraphsim.OEFastFPDatabaseMemoryType_MemoryMapped

    fp_db = oegraphsim.OEFastFPDatabase(args.fp_database, mem_type)
    if not fp_db.IsValid():
        console.print(
            f"[red]Cannot open fingerprint database: {args.fp_database}[/red]"
        )
        return os.EX_IOERR

    if not oegraphsim.OEAreCompatibleDatabases(mol_db, fp_db):
        console.print("[red]Databases are not compatible![/red]")
        return 1

    fp_type = fp_db.GetFPTypeBase()
    console.print(f"Using fingerprint type: {fp_type.GetFPTypeString()}")

    # create report

    report_opts = oedepict.OEReportOptions()
    report_opts.SetHeaderHeight(120.0)
    report_opts.SetFooterHeight(25.0)
    report = oedepict.OEReport(report_opts)

    # search fingerprint database

    opts = oegraphsim.OEFPDatabaseOptions(
        args.num_hits, oegraphsim.OESimMeasure_Tanimoto
    )
    scores = list(fp_db.GetSortedScores(query_mol, opts))

    console.print(f"Found {len(scores)} hits in fingerprint database")

    depict_molecule_similarities(report, query_mol, mol_db, scores, fp_type)

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

    console.print(f"[green]Wrote report to {args.report}[/green]")

    return os.EX_OK


def depict_molecule_similarities(
    report: oedepict.OEReport,
    query_mol: oechem.OEMolBase,
    mol_db: oechem.OEMolDatabase,
    scores: list,
    fp_type: oegraphsim.OEFPTypeBase,
) -> None:
    """
    Depict fingerprint similarity between query and hit molecules.

    Renders a multi-page report with the query molecule in the header
    and hit molecules in grid cells, with bonds colored by fingerprint
    overlap score.
    """
    prep_opts = oedepict.OEPrepareDepictionOptions()
    prep_opts.SetDepictOrientation(oedepict.OEDepictOrientation_Horizontal)
    oedepict.OEPrepareDepiction(query_mol, prep_opts)

    width, height = report.GetHeaderWidth(), report.GetHeaderHeight()
    depict_opts = oedepict.OE2DMolDisplayOptions(
        width, height, oedepict.OEScale_AutoScale
    )
    query_disp = oedepict.OE2DMolDisplay(query_mol, depict_opts)
    scale = query_disp.GetScale()

    depict_opts.SetDimensions(report.GetCellWidth(), report.GetCellHeight(), scale)
    depict_opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)

    self_score = _get_max_bond_self_similarity_score(query_mol, fp_type)

    tag = oechem.OEGetTag("fp-overlap")

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

    bond_glyph = ColorBondByOverlapScore(color_gradient, tag)

    hit = oechem.OEGraphMol()
    for si in scores:
        if mol_db.GetMolecule(hit, si.GetIdx()):
            oedepict.OEPrepareDepiction(hit)
            _set_fingerprint_similarity(query_mol, hit, fp_type, tag)
            hit.SetTitle(f"Score = {si.GetScore():1.3f}")

            overlaps = oegraphsim.OEGetFPOverlap(query_mol, hit, fp_type)
            oedepict.OEPrepareMultiAlignedDepiction(hit, query_mol, overlaps)
            disp = oedepict.OE2DMolDisplay(hit, depict_opts)
            oegrapheme.OEAddGlyph(disp, bond_glyph, oechem.IsTrueBond())

            cell = report.NewCell()
            oedepict.OERenderMolecule(cell, disp)

    for header in report.GetHeaders():
        oedepict.OERenderMolecule(header, query_disp)
        oedepict.OEDrawCurvedBorder(header, oedepict.OELightGreyPen, 10)

    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Bold,
        12,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )
    for page_num, footer in enumerate(report.GetFooters(), start=1):
        text = f"Page {page_num} of {report.NumPages()}"
        oedepict.OEDrawTextToCenter(footer, text, font)


def _get_max_bond_self_similarity_score(
    mol: oechem.OEMolBase,
    fp_type: oegraphsim.OEFPTypeBase,
) -> int:
    """Calculate the maximum bond self-similarity score."""
    ref_bonds = np.zeros(mol.GetMaxBondIdx(), dtype=np.uint32)

    for match in oegraphsim.OEGetFPOverlap(mol, mol, fp_type):
        for bond in match.GetPatternBonds():
            ref_bonds[bond.GetIdx()] += 1

    return int(np.max(ref_bonds))


def _set_fingerprint_similarity(
    ref_mol: oechem.OEMolBase,
    fit_mol: oechem.OEMolBase,
    fp_type: oegraphsim.OEFPTypeBase,
    tag: int,
) -> None:
    """Store per-bond fingerprint overlap counts on the fit molecule."""
    fit_bonds = np.zeros(fit_mol.GetMaxBondIdx(), dtype=np.uint32)

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

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


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 for bond scores."""
        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 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_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_KEYWORDS__", __SCRIPT_KEYWORDS__)
setattr(main, "__SCRIPT_CATEGORIES__", __SCRIPT_CATEGORIES__)

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