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

"""Enumerate carbon-to-nitrogen substitutions in a molecule."""

import argparse
import os
import pathlib
import sys
from itertools import combinations

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Enumerate carbon-to-nitrogen substitutions in a molecule."
__SCRIPT_TOOLKITS__ = ["oechem"]
__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(
        "--in",
        dest="input",
        metavar="INPUT-FILE",
        type=str,
        required=True,
        help="input molecule file",
    )
    io_group.add_argument(
        "--out",
        dest="output",
        metavar="OUTPUT-FILE",
        type=str,
        required=True,
        help="output molecule file",
    )

    sub_group = parser.add_argument_group("Substitution options")
    sub_group.add_argument(
        "--min-subs",
        metavar="N",
        type=int,
        required=True,
        help="minimum number of carbon-to-nitrogen substitutions",
    )
    sub_group.add_argument(
        "--max-subs",
        metavar="N",
        type=int,
        required=True,
        help="maximum number of carbon-to-nitrogen substitutions",
    )

    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 enumerate_substitutions(
    mol: oechem.OEMolBase,
    min_subs: int,
    max_subs: int,
    atom_pred: oechem.OEUnaryAtomPred | None = None,
) -> list[oechem.OEMolBase]:
    """Enumerate all unique carbon-to-nitrogen substitution combinations."""
    if atom_pred is None:
        atom_pred = oechem.OEIsCarbon()
    oechem.OESuppressHydrogens(mol)
    mol.Sweep()
    atom_indices = [atom.GetIdx() for atom in mol.GetAtoms(atom_pred)]

    substituted_mols: list[oechem.OEMolBase] = []
    for n in range(min_subs, max_subs + 1):
        for atom_comb in combinations(atom_indices, n):
            sub_mol = oechem.OEGraphMol(mol)
            if perform_substitution(sub_mol, atom_comb):
                substituted_mols.append(oechem.OEGraphMol(sub_mol))

    return substituted_mols


def perform_substitution(mol: oechem.OEMolBase, atom_comb: tuple[int, ...]) -> bool:
    """Substitute selected carbon atoms with nitrogen."""
    nitrogen_valence = 3

    for atom_idx in atom_comb:
        atom = mol.GetAtom(oechem.OEHasAtomIdx(atom_idx))
        if atom is None:
            return False
        atom.SetAtomicNum(oechem.OEElemNo_N)

        explicit_val = atom.GetExplicitValence()
        if nitrogen_valence > explicit_val:
            atom.SetImplicitHCount(nitrogen_valence - explicit_val)
            atom.SetFormalCharge(0)
        elif nitrogen_valence < explicit_val:
            atom.SetImplicitHCount(0)
            atom.SetFormalCharge(explicit_val - nitrogen_valence)
        else:
            atom.SetImplicitHCount(0)
            atom.SetFormalCharge(0)

    return True


def write_unique_molecules(
    mols: list[oechem.OEMolBase], ofs: oechem.oemolostream
) -> int:
    """Write only unique molecules (by canonical SMILES) to the output stream."""
    unique_smiles: set[str] = set()
    count = 0
    for mol in mols:
        smi = oechem.OEMolToSmiles(mol)
        if smi not in unique_smiles:
            unique_smiles.add(smi)
            oechem.OEWriteMolecule(ofs, mol)
            count += 1
    return count


def main() -> int:
    """Enumerate carbon-to-nitrogen substitutions."""
    args = parse_options()

    console = Console(record=args.save_console_svg)

    ifs = oechem.oemolistream()
    if not ifs.open(args.input):
        oechem.OEThrow.Fatal(f"Cannot open input file: {args.input}")

    ofs = oechem.oemolostream()
    if not ofs.open(args.output):
        oechem.OEThrow.Fatal(f"Cannot open output file: {args.output}")

    mol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ifs, mol):
        oechem.OEThrow.Fatal(f"Cannot read molecule from: {args.input}")

    min_subs = max(1, min(args.min_subs, args.max_subs))
    max_subs = min(mol.NumAtoms(), max(args.min_subs, args.max_subs))

    substituted_mols = enumerate_substitutions(mol, min_subs, max_subs)
    count = write_unique_molecules(substituted_mols, ofs)

    console.print(f"Generated [green]{count}[/green] unique substituted molecules.")

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

    return os.EX_OK


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

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