#!/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 ROC curve from actives list and scored molecules."""

import argparse
import os
import pathlib
import sys
from operator import itemgetter
from pathlib import Path

import matplotlib.pyplot as plt
from rich.console import Console
from rich_argparse import HelpPreviewAction, RichHelpFormatter
from sklearn.metrics import auc

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Plot ROC curve from actives list and scored molecules."
__SCRIPT_TOOLKITS__: list[str] = []


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(
        "--actives",
        metavar="ACTIVES-FILE",
        type=str,
        required=True,
        help="input file with active molecule identifiers (one per line)",
    )
    io_group.add_argument(
        "--scores",
        metavar="SCORES-FILE",
        type=str,
        required=True,
        help="input file with molecule scores (label on first line, id score pairs)",
    )
    io_group.add_argument(
        "--image",
        metavar="IMAGE-FILE",
        type=str,
        required=True,
        help="output image file (png, svg, pdf)",
    )

    plot_group = parser.add_argument_group("Plot options")
    plot_group.add_argument(
        "--color",
        type=str,
        default="#008000",
        help="hex color code for the ROC curve (default: %(default)s)",
    )
    plot_group.add_argument(
        "--no-random-line",
        default=False,
        action="store_true",
        help="do not draw the diagonal random baseline",
    )

    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:
    """Plot ROC curve."""
    args = parse_options()

    console = Console(record=args.save_console_svg)

    image_file = Path(args.image)
    if not _is_supported_image_type(image_file.suffix):
        console.print(f'[red]Format "{image_file.suffix}" is not supported![/red]')
        return os.EX_USAGE

    actives_file = Path(args.actives)
    actives = load_actives(actives_file)
    console.print(f"Loaded [green]{len(actives)}[/green] actives from {actives_file}")

    scores_file = Path(args.scores)
    label, scores = load_scores(scores_file)
    console.print(
        f"Loaded [green]{len(scores)}[/green] {label} scores from {scores_file}"
    )

    sorted_scores = sorted(scores, key=itemgetter(1))

    console.print("Plotting ROC Curve ...")
    depict_roc_curve(
        actives,
        sorted_scores,
        label,
        args.color,
        str(image_file),
        random_line=not args.no_random_line,
    )

    console.print(f"Saved ROC curve to [green]{image_file}[/green]")

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

    return os.EX_OK


def load_actives(filepath: Path) -> list[str]:
    """Load active molecule identifiers from a file."""
    with filepath.open() as f:
        return [line.strip() for line in f if line.strip()]


def load_scores(filepath: Path) -> tuple[str, list[tuple[str, float]]]:
    """Load scored molecules from a file."""
    with filepath.open() as f:
        label = f.readline().strip()
        scores: list[tuple[str, float]] = []
        for line in f:
            mol_id, score = line.strip().split()
            scores.append((mol_id, float(score)))
    return label, scores


def get_rates(
    actives: list[str], scores: list[tuple[str, float]]
) -> tuple[list[float], list[float]]:
    """Compute true positive and false positive rates."""
    tpr = [0.0]
    fpr = [0.0]
    num_actives = len(actives)
    num_decoys = len(scores) - num_actives

    found_actives = 0.0
    found_decoys = 0.0
    for mol_id, _score in scores:
        if mol_id in actives:
            found_actives += 1.0
        else:
            found_decoys += 1.0

        tpr.append(found_actives / float(num_actives))
        fpr.append(found_decoys / float(num_decoys))

    return tpr, fpr


def setup_roc_curve_plot() -> None:
    """Configure ROC curve plot axes and title."""
    plt.xlabel("FPR", fontsize=14)
    plt.ylabel("TPR", fontsize=14)
    plt.title("ROC Curve", fontsize=14)


def save_roc_curve_plot(filename: str, random_line: bool) -> 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, 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(filename)


def depict_roc_curve(
    actives: list[str],
    scores: list[tuple[str, float]],
    label: str,
    color: str,
    filename: str,
    random_line: bool,
) -> None:
    """Generate and save a ROC curve plot."""
    plt.figure(figsize=(4, 4), dpi=80)

    setup_roc_curve_plot()
    add_roc_curve(actives, scores, color, label)
    save_roc_curve_plot(filename, random_line)


def add_roc_curve(
    actives: list[str],
    scores: list[tuple[str, float]],
    color: str,
    label: str,
) -> None:
    """Add a ROC curve to the current matplotlib plot."""
    tpr, fpr = get_rates(actives, scores)
    roc_auc = auc(fpr, tpr)

    roc_label = f"{label} (AUC={roc_auc:.3f})"
    plt.plot(fpr, tpr, color=color, linewidth=2, label=roc_label)


def _is_supported_image_type(ext: str) -> bool:
    """Check if the image format is supported by matplotlib."""
    fig = plt.figure()
    supported = ext.lstrip(".") in fig.canvas.get_supported_filetypes()
    plt.close(fig)
    return supported


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

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