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

"""Plot fingerprint ROC curves."""

import argparse
import os
import sys
from pathlib import Path

import matplotlib.pyplot as plt
from numpy import interp, linspace, maximum, mean, minimum, std
from openeye import oechem, oegraphsim
from rich.console import Console
from rich_argparse import HelpPreviewAction, RichHelpFormatter
from sklearn.metrics import auc, roc_curve

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Plot fingerprint ROC curves."
__SCRIPT_TOOLKITS__ = ["oechem", "oegraphsim"]


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
    )

    input_group = parser.add_argument_group("Input options")
    input_group.add_argument(
        "--active",
        type=str,
        required=True,
        metavar="FP-DATABASE-FILE",
        help="input binary fingerprint database file of actives (.fpbin)",
    )
    input_group.add_argument(
        "--decoy",
        type=str,
        required=True,
        metavar="FP-DATABASE-FILE",
        help="input binary fingerprint database file of decoys (.fpbin)",
    )

    output_group = parser.add_argument_group("Output options")
    output_group.add_argument(
        "--image",
        type=str,
        required=True,
        metavar="IMAGE-FILE",
        help="output image file (e.g. PNG, PDF, SVG)",
    )

    parser.add_argument(
        "--mem-type",
        type=str,
        default="in-memory",
        choices=["in-memory", "on-disk"],
        help="database memory type (default: %(default)s)",
    )

    parser.add_argument("--help-image", action=HelpPreviewAction)
    return parser.parse_args()


def _get_memory_type(mem_type_str: str) -> int:
    """Return the OEFPDatabaseMemoryType constant for the given string."""
    if mem_type_str == "on-disk":
        return oegraphsim.OEFastFPDatabaseMemoryType_MemoryMapped
    return oegraphsim.OEFastFPDatabaseMemoryType_InMemory


def plot_roc_curves(
    active_fp_database: oegraphsim.OEFastFPDatabase,
    decoy_fp_database: oegraphsim.OEFastFPDatabase,
) -> None:
    """
    Plot ROC curves and calculate AUC for each active fingerprint.

    For each fingerprint in the active database, scores are computed against
    both the active and decoy databases. Individual ROC curves are plotted
    in light green, and the mean ROC curve with +/- 1 std. dev. band is
    overlaid.
    """
    tprs = []
    aucs = []
    mean_fpr = linspace(0, 1, 100)

    limit = 0  # zero means no limit, all scores returned
    opts = oegraphsim.OEFPDatabaseOptions(limit, oegraphsim.OESimMeasure_Tanimoto)
    fp = oegraphsim.OEFingerPrint()
    for fp_idx in range(active_fp_database.NumFingerPrints()):
        if not active_fp_database.GetFingerPrint(fp, fp_idx):
            continue

        decoy_scores = [si.GetScore() for si in decoy_fp_database.GetScores(fp, opts)]
        active_scores = [si.GetScore() for si in active_fp_database.GetScores(fp, opts)]

        active_flags = [1] * len(decoy_scores) + [0] * len(active_scores)
        scores = decoy_scores + active_scores
        tpr, fpr, _threshold = roc_curve(active_flags, scores)
        tprs.append(interp(mean_fpr, fpr, tpr))
        roc_auc = auc(fpr, tpr)
        aucs.append(roc_auc)
        plt.plot(fpr, tpr, color="green", linewidth=1, alpha=0.10, label=None)

    mean_tpr = mean(tprs, axis=0)
    mean_tpr[-1] = 1.0
    mean_auc = auc(mean_fpr, mean_tpr)
    std_auc = std(aucs)

    plt.plot(
        mean_fpr,
        mean_tpr,
        color="green",
        label=rf"Mean ROC (AUC = {mean_auc:0.2f} $\pm$ {std_auc:0.2f})",
        linewidth=2.0,
        alpha=0.80,
    )

    std_tpr = std(tprs, axis=0)
    tprs_upper = minimum(mean_tpr + std_tpr, 1)
    tprs_lower = maximum(mean_tpr - std_tpr, 0)
    plt.fill_between(
        mean_fpr,
        tprs_lower,
        tprs_upper,
        color="grey",
        alpha=0.33,
        label=r"$\pm$ 1 std. dev.",
    )


def _setup_roc_curve_plot(
    fp_type: oegraphsim.OEFPTypeBase,
    active_name: str,
    decoy_name: str,
) -> None:
    """Configure ROC curve plot labels and title."""
    plt.xlabel("FPR", fontsize=14)
    plt.ylabel("TPR", fontsize=14)
    fp_type_str_comps = fp_type.GetFPTypeString().split(",")
    title_lines = [
        "ROC Curve",
        " ".join(fp_type_str_comps[ci] for ci in [0, 2]),
        " ".join(fp_type_str_comps[ci] for ci in [4, 5]),
        f"actives: {Path(active_name).name}",
        f"decoys: {Path(decoy_name).name}",
    ]
    plt.title("\n".join(title_lines), fontsize=10, weight="bold", linespacing=2.0)


def _save_roc_curve_plot(fname: str, *, random_line: bool = True) -> None:
    """Save the ROC curve plot to a file."""
    if random_line:
        x = [0.0, 1.0]
        plt.plot(x, x, linestyle="dashed", color="red", linewidth=2.0, label="random")

    plt.xlim(0.0, 1.0)
    plt.ylim(0.0, 1.0)
    plt.legend(fontsize=10, loc="best")
    plt.tight_layout()
    plt.savefig(fname)


def main() -> int:
    """Plot fingerprint ROC curves."""
    args = parse_args()
    console = Console()

    mem_type = _get_memory_type(args.mem_type)

    # load active fingerprint database
    active_fp_database = oegraphsim.OEFastFPDatabase(args.active, mem_type)
    if not active_fp_database.IsValid():
        oechem.OEThrow.Fatal(
            f"Cannot open fingerprint database of actives {args.active}!"
        )
    active_fp_type = active_fp_database.GetFPTypeBase()
    console.print(
        f"Loaded {active_fp_database.NumFingerPrints():4d} fingerprints of actives "
        f"({active_fp_type.GetFPTypeString()})"
    )

    # load decoy fingerprint database
    decoy_fp_database = oegraphsim.OEFastFPDatabase(args.decoy, mem_type)
    if not decoy_fp_database.IsValid():
        oechem.OEThrow.Fatal(
            f"Cannot open fingerprint database of decoys {args.decoy}!"
        )
    decoy_fp_type = decoy_fp_database.GetFPTypeBase()
    console.print(
        f"Loaded {decoy_fp_database.NumFingerPrints():4d} fingerprints of decoys "
        f"({decoy_fp_type.GetFPTypeString()})"
    )

    if active_fp_type != decoy_fp_type:
        oechem.OEThrow.Fatal("Fingerprint type mismatch!")

    # plot fingerprint ROCs
    plt.figure(figsize=(6, 6), dpi=80)
    _setup_roc_curve_plot(active_fp_type, args.active, args.decoy)
    plot_roc_curves(active_fp_database, decoy_fp_database)
    _save_roc_curve_plot(args.image, random_line=True)

    console.print(f"Image written to {args.image}")
    return os.EX_OK


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

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