#!/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 splitting a molecule file into N chunks or chunks of size N."""

import argparse
import os
import sys
from pathlib import Path

from openeye import oechem


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Split molecule file into N chunks or chunks of size N.",
    )
    parser.add_argument("--input", required=True, help="Input file name")
    parser.add_argument("--output", required=True, help="Output file name base")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--num", type=int, help="The number of chunks")
    group.add_argument("--size", type=int, help="The size of each chunk")

    return parser.parse_args()


def main() -> int:
    """Split molecule file into chunks."""
    args = parse_args()

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

    ifs.SetConfTest(oechem.OEIsomericConfTest(False))

    out_path = Path(args.output)
    suffixes = out_path.suffixes
    if not suffixes:
        oechem.OEThrow.Fatal("Failed to find file extension")

    ext = "".join(suffixes)
    out_base = str(out_path).removesuffix(ext)

    if args.num is not None:
        split_n_parts(ifs, args.num, out_base, ext)
    else:
        split_chunk(ifs, args.size, out_base, ext)

    return os.EX_OK


def split_n_parts(
    ifs: oechem.oemolistream, num_parts: int, out_base: str, ext: str
) -> None:
    """Split the molecule database into a fixed number of equal parts."""
    mol_database = oechem.OEMolDatabase(ifs)
    mol_count = mol_database.NumMols()

    chunk_size, lft = divmod(mol_count, num_parts)
    if lft != 0:
        chunk_size += 1
    chunk, count = 1, 0

    ofs = _new_output_stream(out_base, ext, chunk)
    for idx in range(mol_database.GetMaxMolIdx()):
        count += 1
        if count > chunk_size:
            if chunk == lft:
                chunk_size -= 1

            ofs.close()
            chunk, count = chunk + 1, 1
            ofs = _new_output_stream(out_base, ext, chunk)

        mol_database.WriteMolecule(ofs, idx)


def split_chunk(
    ifs: oechem.oemolistream, chunk_size: int, out_base: str, ext: str
) -> None:
    """Split the molecule database into chunks of a fixed size."""
    mol_database = oechem.OEMolDatabase(ifs)
    chunk, count = 1, chunk_size

    for idx in range(mol_database.GetMaxMolIdx()):
        if count == chunk_size:
            ofs = _new_output_stream(out_base, ext, chunk)
            chunk, count = chunk + 1, 0
        count += 1
        mol_database.WriteMolecule(ofs, idx)


def _new_output_stream(out_base: str, ext: str, chunk: int) -> oechem.oemolostream:
    """Create a new output molecule stream for a chunk."""
    new_name = f"{out_base}_{chunk:07d}{ext}"
    ofs = oechem.oemolostream()
    if not ofs.open(new_name):
        oechem.OEThrow.Fatal(f"Unable to open {new_name} for writing")
    return ofs


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