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

"""Generate random peptides (HELM or/and structure)."""

import argparse
import contextlib
import json
import os
import pathlib
import random
import sys
from collections.abc import Iterator

import rich.console
from openeye import oechem
from rich.markdown import Markdown
from rich.progress import BarColumn, Progress, TextColumn
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Generate random peptides."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = ["HELM", "monomer", "peptide", "peptide-informatics", "random"]
__SCRIPT_CATEGORIES__ = ["peptide-informatics"]


MONOMER_SET_STANDARD = "Standard"
MONOMER_SET_OPENEYE = "OpenEye"
MONOMER_SET_JSON_FILENAME = "JSON-FILENAME"

__PATTERN_EXAMPLES__ = """
----------------------------
Examples:


- PEPTIDE1{[X1].A.[X3]}$$$$                                - to generate random tri-peptides with Alanine in the middle
- PEPTIDE1{[X1].[X2].[X3].[X4].[X5]}$$$$                   - to generate random palindrome penta-peptides
- PEPTIDE1{[X1].[X2].[X3].[X4].[X5]}$PEPTIDE1,PEPTIDE1,1:R1-5:R2$$$      - generate random cyclic penta-peptides
- PEPTIDE1{[X1].C.[X3].[X4].[X5].C.[X7]}$PEPTIDE1,PEPTIDE1,2:R3-6:R3$$$  - generate random disulfide bridged hepta-peptides
"""


def parse_options() -> argparse.Namespace:
    """Parse command-line options for random peptide generation."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
        epilog=Markdown(__PATTERN_EXAMPLES__),  # type: ignore  # noqa: PGH003
    )

    io_group = parser.add_argument_group("Output options")
    io_group.add_argument(
        "--helm",
        metavar="HELM-FILE",
        type=str,
        required=False,
        help="output file of HELM string",
    )
    io_group.add_argument(
        "--mol",
        metavar="MOL-FILE",
        type=str,
        required=False,
        help="output molecule file",
    )

    generation_group = parser.add_argument_group("Peptide generation options")
    generation_group.add_argument(
        "--helm-pattern",
        metavar="HELM-PATTERN",
        default="PEPTIDE1{[X1].[X2].[X3]}$$$$",
        type=str,
        required=False,
        help="HELM pattern to generate peptides (default: %(default)s)",
    )
    generation_group.add_argument(
        "--num-peptides",
        metavar="N",
        default=10,
        type=int,
        required=False,
        help="number of peptides to generate (default: %(default)s)",
    )
    generation_group.add_argument(
        "--random-seed",
        metavar="SEED",
        type=int,
        required=False,
        help="random seed for peptide generation",
    )

    monomers_group = parser.add_argument_group("Monomer set options")
    _add_monomer_collection(monomers_group)

    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:
    """Generate random peptides and write requested outputs."""
    args = parse_options()

    console = rich.console.Console(record=args.save_console_svg)
    if args.helm is None and args.mol is None:
        console.print(
            "[red]Error: No output file specified, use --helm or --mol.[/red]"
        )
        return os.EX_USAGE

    monomers = _get_monomer_collection(args)
    code_set = monomers.GetPrimaryCodeSet() if args.code_set is None else args.code_set

    mol_file: None | oechem.oemolostream = None
    if args.mol is not None:
        mol_file = oechem.oemolostream(args.mol)

    helm_pattern = args.helm_pattern
    if not (
        replace_codes := _validate_helm_pattern(
            helm_pattern, monomers, code_set, console
        )
    ):
        return os.EX_DATAERR
    num_peptides = args.num_peptides

    replace_code_dict: dict[str, list[str]] = {}
    for code in replace_codes:
        equivalent_codes = _get_equivalent_monomer_codes(
            code, helm_pattern, monomers, code_set, replace_codes
        )
        console.print(
            f"[bold blue][{code:2s}][/bold blue] can be randomly replaced with {len(equivalent_codes)} monomers"
        )
        replace_code_dict[code] = equivalent_codes

    rand = random.Random(args.random_seed)  # noqa: S311
    with (
        Progress(
            TextColumn("{task.description}"),
            BarColumn(bar_width=60),
            TextColumn("{task.percentage:3.1f}%"),
            transient=True,
            disable=num_peptides < 10,  # noqa: PLR2004
            console=console,
        ) as progress,
        (
            pathlib.Path(args.helm).open("w")
            if args.helm is not None
            else contextlib.nullcontext()
        ) as helm_file,
    ):
        conversion = progress.add_task("[blue]HELM conversion", total=num_peptides)
        for mol, helm in get_next_random_peptide(
            helm_pattern, num_peptides, replace_code_dict, monomers, rand
        ):
            if helm_file is not None:
                helm_file.write(helm + "\n")
            if mol_file is not None:
                oechem.OEWriteMolecule(mol_file, mol)
            progress.update(conversion, advance=1)

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


def get_next_random_peptide(
    helm_pattern: str,
    num_peptides: int,
    replace_code_dict: dict[str, list[str]],
    monomers: oechem.OEMonomerSet,
    rand: random.Random,
) -> Iterator[tuple[oechem.OEMolBase, str]]:
    """Generate random peptides as OEGraphMol and HELM string tuples."""
    mol = oechem.OEGraphMol()
    while num_peptides > 0:
        helm = helm_pattern
        for replace_code, equivalent_codes in replace_code_dict.items():
            random_code = equivalent_codes[rand.randint(0, len(equivalent_codes) - 1)]
            helm = helm.replace(
                f"[{replace_code}]",
                f"[{random_code}]" if len(random_code) > 1 else random_code,
            )
        if not oechem.OEHelmToMol(mol, helm, monomers):
            # skip invalid HELM patterns that cannot be converted to molecules with the given monomer set
            continue
        num_peptides -= 1
        yield (oechem.OEGraphMol(mol), helm)


class MonomerSetParameter:  # noqa: PLW1641
    """Utility class to handle both built-in and user defined monomer sets."""

    def __init__(self) -> None:  # noqa: D107
        self._monomer_sets = [
            MONOMER_SET_STANDARD,
            MONOMER_SET_OPENEYE,
            MONOMER_SET_JSON_FILENAME,
        ]

    def __repr__(self) -> str:  # noqa: D105
        return ",".join(self._monomer_sets)

    def __eq__(self, param: object) -> bool:  # noqa: D105
        if not isinstance(param, str):
            return False
        if param in [MONOMER_SET_STANDARD, MONOMER_SET_OPENEYE]:
            return True

        console = rich.console.Console()
        monomer_set_filepath = pathlib.Path(param)
        if (
            not monomer_set_filepath.exists()
            or monomer_set_filepath.suffix.lower() != ".json"
        ):
            console.print(f"[red]Invalid monomer set file '{param}' ![/red]")
            return False

        is_valid = True
        try:
            with monomer_set_filepath.open("r") as json_file:
                json.load(json_file)
        except json.JSONDecodeError as e:
            console.print(f"[red]Invalid monomer set file '{param}' ![/red]")
            console.print(f"[red]Error decoding JSON: {e} ![/red]")
            is_valid = False
        return is_valid


def _add_monomer_collection(arg_group: argparse._ArgumentGroup) -> None:
    arg_group.add_argument(
        "-m",
        "--monomers",
        type=str,
        default=MONOMER_SET_STANDARD,
        choices=[MonomerSetParameter()],
        help="built-in monomer-set type or json file of monomers",
    )
    arg_group.add_argument(
        "--code-set",
        type=str,
        metavar="CODE-SET",
        required=False,
        default=None,
        help="code-set, if not specified primary code-set is used",
    )


def _get_monomer_collection(args: argparse.Namespace) -> oechem.OEMonomerSet:
    monomers = oechem.OEMonomerSet()
    match args.monomers:
        case "Standard":
            oechem.OELoadStandardMonomerSet(monomers)
        case "OpenEye":
            oechem.OELoadOpenEyeMonomerSet(monomers)
        case _:
            oechem.OEReadMonomerSet(monomers, args.monomers)
    return monomers


def _validate_helm_pattern(
    helm_pattern: str,
    monomers: oechem.OEMonomerSet,
    code_set: str,
    console: rich.console.Console,
) -> None | list[str]:
    if oechem.OEGetHelmChainNames(helm_pattern) != ["PEPTIDE1"]:
        console.print(
            f"[red]Invalid HELM pattern: {helm_pattern}; it should only have the PEPTIDE1 chain![/red]"
        )
        return None
    codes = set(oechem.OEGetHelmMonomerCodes(helm_pattern, "PEPTIDE1"))
    replace_codes: list[str] = [
        code for code in codes if monomers.GetMonomer(code_set, code) is None
    ]
    replace_codes.sort()
    console.print(
        f"Randomly replace monomer codes [blue]{replace_codes}[/blue] in `{helm_pattern}`"
    )

    # generate "fake" monomers to test helm_pattern
    for repeat, code in enumerate(replace_codes, start=1):
        fake_monomer_smiles = "[H:1]N[C@@H](C" + repeat * "[U]" + "S[H:3])C(=O)[OH:2]"
        monomer_data = oechem.OEMonomerData(fake_monomer_smiles)
        monomer_data.AddCode(code_set, code)
        if not oechem.OEIsValidMonomerData(monomer_data):
            return None
        monomers.AddMonomer(monomer_data)

    mol = oechem.OEGraphMol()
    result = oechem.OEHelmParsingResult()
    if not oechem.OEHelmToMol(mol, helm_pattern, monomers, result):
        console.print(
            f"[red]Invalid HELM pattern: {helm_pattern} : {result.GetWarning()}![/red]"
        )
        return None
    return replace_codes


def _get_equivalent_monomer_codes(
    replace_code: str,
    helm_pattern: str,
    monomers: oechem.OEMonomerSet,
    code_set: str,
    ignore_codes: list[str],
) -> list[str]:
    equivalent_monomer_codes = []
    mol = oechem.OEGraphMol()
    for monomer in monomers.GetMonomers(oechem.OEIsInMonomerCodeSet(code_set)):
        code = monomer.GetCode(code_set)
        if code in ignore_codes:
            continue
        if monomer.GetPolymerType() != oechem.OEPolymerType_Peptide:
            continue
        test_helm = helm_pattern.replace(
            f"[{replace_code}]", f"[{code}]" if len(code) > 1 else code
        )
        if oechem.OEHelmToMol(mol, test_helm, monomers):
            equivalent_monomer_codes.append(code)

    return equivalent_monomer_codes


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

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