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


"""Code snippet for randomly reordering molecules and obtaining a random subset."""

import argparse
import os
import sys
from random import Random

from openeye import oechem


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Randomly reorder molecules and optionally obtain a random subset.",
    )
    parser.add_argument("-i", "--input", required=True, help="Input file name")
    parser.add_argument("-o", "--output", default=None, help="Output file name")
    parser.add_argument(
        "--seed", type=int, default=None, help="Random seed (default: system time)"
    )

    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "-p", "--percent", type=float, help="Percentage of output molecules"
    )
    group.add_argument("-n", "--number", type=int, help="Number of output molecules")

    return parser.parse_args()


def main() -> int:
    """Randomly reorder and sample molecules."""
    args = parse_args()

    ifs = oechem.oemolistream()
    if not ifs.open(args.input):
        oechem.OEThrow.Fatal(f"Unable to open {args.input} for reading")

    ofs = oechem.oemolostream(".ism")
    if args.output and not ofs.open(args.output):
        oechem.OEThrow.Fatal(f"Unable to open {args.output} for writing")

    rand = Random(args.seed)  # noqa: S311

    if args.number is not None:
        randomize_n(ifs, ofs, args.number, rand)
    elif args.percent is not None:
        randomize_percent(ifs, ofs, args.percent, rand)
    else:
        randomize(ifs, ofs, rand)

    return os.EX_OK


def randomize(ifs: oechem.oemolistream, ofs: oechem.oemolostream, rand: Random) -> None:
    """Randomly reorder all molecules in the database."""
    randomize_percent(ifs, ofs, 100.0, rand)


def randomize_percent(
    ifs: oechem.oemolistream,
    ofs: oechem.oemolostream,
    percent: float,
    rand: Random,
) -> None:
    """Randomly sample a percentage of molecules from the database."""
    mol_database = oechem.OEMolDatabase(ifs)
    indices = range(mol_database.GetMaxMolIdx())
    size = max(1, int(percent * 0.01 * mol_database.GetMaxMolIdx()))
    mol_indices = rand.sample(indices, size)
    _write_database(mol_database, ofs, mol_indices)


def randomize_n(
    ifs: oechem.oemolistream,
    ofs: oechem.oemolostream,
    count: int,
    rand: Random,
) -> None:
    """Randomly sample a fixed number of molecules from the database."""
    mol_database = oechem.OEMolDatabase(ifs)
    indices = range(mol_database.GetMaxMolIdx())
    mol_indices = rand.sample(indices, count)
    _write_database(mol_database, ofs, mol_indices)


def _write_database(
    mol_database: oechem.OEMolDatabase, ofs: oechem.oemolostream, mol_indices: list[int]
) -> None:
    """Write selected molecules to the output stream."""
    for mol_idx in mol_indices:
        mol_database.WriteMolecule(ofs, mol_idx)


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