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

"""Generate a fingerprint file for fast fingerprint search."""

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

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

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Generate a fingerprint file for fast fingerprint search."
__SCRIPT_TOOLKITS__ = ["oechem", "oegraphsim"]
__SCRIPT_KEYWORDS__ = ["fingerprints", "similarity"]
__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(
        "--mol",
        type=str,
        required=True,
        metavar="MOL-FILE",
        help="input molecule file",
    )
    io_group.add_argument(
        "--fp-database",
        type=str,
        required=True,
        metavar="BINARY-FILE",
        help="output fingerprint database file (.fpbin)",
    )

    fp_group = parser.add_argument_group("Fingerprint generation options")
    fp_group.add_argument(
        "--fp-type",
        type=str,
        default="tree",
        choices=["tree", "circular", "path"],
        help="fingerprint type (default: %(default)s)",
    )
    fp_group.add_argument(
        "--fp-size",
        type=int,
        default=4096,
        choices=[512, 1024, 2048, 4096, 8192],
        help="fingerprint size (default: %(default)s)",
    )

    parser.add_argument(
        "--num-processors",
        type=int,
        default=multiprocessing.cpu_count() - 1,
        metavar="N",
        help="number of processors (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:
    """Generate a binary fingerprint file for fast fingerprint search."""
    args = parse_options()

    console = Console(record=args.save_console_svg)
    num_processors: int = max(1, min(args.num_processors, multiprocessing.cpu_count()))

    if Path(args.fp_database).suffix != ".fpbin":
        oechem.OEThrow.Fatal(
            "Fingerprint database file should have '.fpbin' file extension!"
        )

    idx_fname = oechem.OEGetMolDatabaseIdxFileName(args.mol)

    if not Path(idx_fname).exists() and not oechem.OECreateMolDatabaseIdx(args.mol):
        oechem.OEThrow.Warning(f"Unable to create {idx_fname} molecule index file")

    console.print(f"Using {Path(idx_fname).name} index molecule file")

    mol_db = oechem.OEMolDatabase()
    if not mol_db.Open(args.mol):
        oechem.OEThrow.Fatal("Cannot open molecule database file!")
        return os.EX_DATAERR

    fp_type = _get_fp_type(args)
    if fp_type is None:
        oechem.OEThrow.Fatal(f"Unsupported fingerprint type: {args.fp_type}")
        return os.EX_DATAERR

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

    opts = oegraphsim.OECreateFastFPDatabaseOptions(fp_type)
    opts.SetNumProcessors(num_processors)
    dots = oechem.OEDots(100000, 1000, "fingerprints")
    opts.SetTracer(dots)

    console.print(f"Generating fingerprints with {opts.GetNumProcessors()} threads")

    timer = oechem.OEWallTimer()
    if not oegraphsim.OECreateFastFPDatabaseFile(args.fp_database, args.mol, opts):
        oechem.OEThrow.Fatal("Cannot create fingerprint database file!")

    delta = datetime.timedelta(seconds=timer.Elapsed())
    console.print(
        f"[blue]{humanize.precisedelta(delta)}[/blue] to generate {humanize.intcomma(mol_db.GetMaxMolIdx())} fingerprints"
    )

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


def _get_fp_type(args: argparse.Namespace) -> oegraphsim.OEFPTypeBase | None:
    """Return the fingerprint type object for the given name."""
    match args.fp_type:
        case "tree":
            return oegraphsim.OEGetTreeFPType(
                args.fp_size,
                0,
                4,
                oegraphsim.OEFPAtomType_DefaultTreeAtom,
                oegraphsim.OEFPBondType_DefaultTreeBond,
            )

        case "circular":
            return oegraphsim.OEGetCircularFPType(
                args.fp_size,
                0,
                5,
                oegraphsim.OEFPAtomType_DefaultCircularAtom,
                oegraphsim.OEFPBondType_DefaultCircularBond,
            )
        case "path":
            return oegraphsim.OEGetPathFPType(
                args.fp_size,
                0,
                5,
                oegraphsim.OEFPAtomType_DefaultPathAtom,
                oegraphsim.OEFPBondType_DefaultPathBond,
            )

    return None


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