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

"""Mutate peptide (HELM or/and structure)."""

import argparse
import contextlib
import itertools
import json
import math
import os
import pathlib
import re
import sys
from collections.abc import Iterator

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Mutate peptide (HELM or/and structure)."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = ["HELM", "monomer", "peptide", "peptide-informatics", "mutate"]
__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 mutate tri-peptides at two positions
- PEPTIDE1{[A].[P].[X3].[C].[F]}$PEPTIDE1,PEPTIDE1,1:R1-5:R2$$$    - to mutate penta-peptide at one position
"""


def parse_options() -> tuple[argparse.Namespace, list[str]]:
    """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",
        type=str,
        required=True,
        help="HELM pattern used for mutation (default: %(default)s)",
    )

    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_known_args()


def _parse_replace_args(unknown_args: list[str]) -> dict[str, list[str]]:
    """Parse unknown args like ['--X1', 'A', 'C', '--X2', 'D,', 'F'] into {'X1': ['A', 'C'], 'X2': ['D', 'F']}."""
    replace_dict: dict[str, list[str]] = {}
    current_key: None | str = None
    for arg in unknown_args:
        if re.match(r"^-{1,2}X\d+$", arg):
            current_key = arg.lstrip("-")
            replace_dict[current_key] = []
        elif current_key is not None:
            for code in arg.split(","):
                if code:
                    replace_dict[current_key].append(code.strip())
    return replace_dict


def main() -> int:
    """Generate random peptides and write requested outputs."""
    args, unknown_args = parse_options()
    replace_args = _parse_replace_args(unknown_args)

    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

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

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

    replace_code_dict = _build_replace_code_dict(
        replace_codes, replace_args, helm_pattern, monomers, code_set, console
    )
    if replace_code_dict is None:
        return os.EX_DATAERR

    num_combinations = math.prod(len(codes) for codes in replace_code_dict.values())
    console.print(
        f"Number of possible mutations: [bold green]{num_combinations}[/bold green]"
    )

    with (
        Progress(
            TextColumn("{task.description}"),
            TextColumn("{task.percentage:3.1f}%"),
            transient=True,
            disable=num_combinations < 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,
    ):
        task = progress.add_task("[blue]Generating mutations", total=num_combinations)
        for helm, mol in generate_all_mutations(
            helm_pattern, replace_code_dict, monomers
        ):
            if helm_file is not None:
                helm_file.write(helm + "\n")
            if mol_file is not None:
                oechem.OEWriteMolecule(mol_file, mol)
            progress.update(task, advance=1)

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


def generate_all_mutations(
    helm_pattern: str,
    replace_code_dict: dict[str, list[str]],
    monomers: oechem.OEMonomerSet,
) -> Iterator[tuple[str, oechem.OEMolBase]]:
    """Generate all possible peptide mutations from a HELM pattern."""
    mol = oechem.OEGraphMol()
    replace_codes = list(replace_code_dict.keys())
    code_lists = [replace_code_dict[code] for code in replace_codes]
    for combination in itertools.product(*code_lists):
        helm = helm_pattern
        for replace_code, chosen_code in zip(replace_codes, combination, strict=True):
            helm = helm.replace(
                f"[{replace_code}]",
                f"[{chosen_code}]" if len(chosen_code) > 1 else chosen_code,
            )
        if oechem.OEHelmToMol(mol, helm, monomers):
            yield (helm, oechem.OEGraphMol(mol))


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"Mutate 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 _build_replace_code_dict(
    replace_codes: list[str],
    replace_args: dict[str, list[str]],
    helm_pattern: str,
    monomers: oechem.OEMonomerSet,
    code_set: str,
    console: rich.console.Console,
) -> None | dict[str, list[str]]:
    """Build a dictionary mapping replace codes to their equivalent monomer codes."""
    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
        )
        if code in replace_args:
            specified_codes = replace_args[code]
            equivalent_codes = [c for c in equivalent_codes if c in specified_codes]
            if not equivalent_codes:
                console.print(
                    f"[red]Error: No valid monomer codes found for replace code '{code}' with specified codes {specified_codes}![/red]"
                )
                return None
        codes_display = (
            f"[{', '.join(equivalent_codes)}]"
            if len(equivalent_codes) <= 10  # noqa: PLR2004
            else f"{len(equivalent_codes)} codes"
        )
        console.print(
            f"[bold blue]{code}[/bold blue] can be replaced with {codes_display}"
        )
        replace_code_dict[code] = equivalent_codes
    return replace_code_dict


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())