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

"""Search a fast fingerprint database for similar molecules."""

import argparse
import datetime
import os
import sys
from pathlib import Path

import humanize
import rich.table
from openeye import oechem, oegraphsim
from rich.console import Console
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Search a fast fingerprint database for similar molecules."
__SCRIPT_TOOLKITS__ = ["oechem", "oegraphsim"]
__SCRIPT_KEYWORDS__ = ["fingerprints", "similarity", "search"]
__SCRIPT_CATEGORIES__ = ["cheminformatics"]


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,
        required=True,
        metavar="MOL-FILE",
        help="input query molecule file",
    )
    io_group.add_argument(
        "--mol",
        type=str,
        required=True,
        metavar="MOL-FILE",
        help="input molecule database file",
    )
    io_group.add_argument(
        "--fp-database",
        type=str,
        required=True,
        metavar="BINARY-FILE",
        help="input fast fingerprint database file (.fpbin)",
    )
    io_group.add_argument(
        "--hits",
        type=str,
        required=True,
        metavar="MOL-FILE",
        help="output molecule file of similarity search hits",
    )

    search_group = parser.add_argument_group("Search options")
    search_group.add_argument(
        "--num-hits",
        type=int,
        default=10,
        metavar="N",
        help="number of hits to return (default: %(default)s)",
    )
    search_group.add_argument(
        "--cutoff",
        type=float,
        default=0.0,
        metavar="SCORE",
        help="minimum similarity score cutoff (default: %(default)s)",
    )
    search_group.add_argument(
        "--memory-mode",
        type=str,
        default="memory-mapped",
        choices=["memory-mapped", "in-memory"],
        help="memory mode for fingerprint database (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",
    )
    parser.add_argument(
        "--display-top-hits",
        default=False,
        action="store_true",
        help="display top ten hits with scores in console",
    )

    return parser.parse_args()


def main() -> int:
    """Search a fast fingerprint database for similar molecules."""
    args = parse_options()

    console = Console(record=args.save_console_svg)

    timer = oechem.OEWallTimer()

    query: oechem.OEGraphMol = _get_query_molecule(args.query)
    console.print(
        f"Query molecule: {oechem.OEMolToSmiles(query)}", markup=False, highlight=False
    )

    ofs = oechem.oemolostream()
    if not ofs.open(args.hits):
        oechem.OEThrow.Fatal("Cannot open output file!")

    mol_database = oechem.OEMolDatabase()
    if not mol_database.Open(args.mol):
        oechem.OEThrow.Fatal("Cannot open molecule database!")

    timer.Start()
    memory_type = (
        oegraphsim.OEFastFPDatabaseMemoryType_MemoryMapped
        if args.memory_mode == "memory-mapped"
        else oegraphsim.OEFastFPDatabaseMemoryType_InMemory
    )
    fp_database = oegraphsim.OEFastFPDatabase(args.fp_database, memory_type)
    if not fp_database.IsValid():
        oechem.OEThrow.Fatal("Cannot open fingerprint database!")
    num_fingerprints = fp_database.NumFingerPrints()

    if not oegraphsim.OEAreCompatibleDatabases(mol_database, fp_database):
        oechem.OEThrow.Fatal("Databases are not compatible!")

    delta = datetime.timedelta(seconds=timer.Elapsed())
    console.print(
        f"[blue]{humanize.precisedelta(delta)}[/blue] secs to initialize databases"
        f" [blue]{fp_database.GetMemoryTypeString()}[/blue] mode"
    )

    fp_type = fp_database.GetFPTypeBase()
    console.print(
        f"Using fingerprint type {fp_type.GetFPTypeString()}", highlight=False
    )

    opts = oegraphsim.OEFPDatabaseOptions(
        args.num_hits, oegraphsim.OESimMeasure_Tanimoto
    )
    opts.SetCutoff(args.cutoff)

    timer.Start()
    scores = fp_database.GetSortedScores(query, opts)
    delta = datetime.timedelta(seconds=timer.Elapsed())

    console.print(
        f"[blue]{humanize.precisedelta(delta)}[/blue] to search {humanize.intcomma(num_fingerprints)} fingerprints"
    )

    columns = ["score", "smiles"]
    table = rich.table.Table(*columns)

    hit = oechem.OEGraphMol()
    hit_count = 0
    for si in scores:
        if mol_database.GetMolecule(hit, si.GetIdx()):
            oechem.OESetSDData(hit, "Similarity score", f"{si.GetScore():.4f}")
            oechem.OEWriteMolecule(ofs, hit)
            if args.display_top_hits and hit_count < 10:  # noqa: PLR2004
                table.add_row(f"{si.GetScore():.4f}", oechem.OEMolToSmiles(hit))
            hit_count += 1

    if args.display_top_hits:
        console.print(table)

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


def _get_query_molecule(filename: str) -> oechem.OEGraphMol:
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        oechem.OEThrow.Fatal("Cannot open query file!")
    query = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ifs, query):
        oechem.OEThrow.Fatal("Cannot read query molecule!")
    return query


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