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

"""Convert a molecule file into a HELM file."""

import argparse
import collections
import json
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Convert molecule file into HELM file."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = ["HELM", "peptide", "peptide-informatics"]
__SCRIPT_CATEGORIES__ = ["peptide-informatics"]


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(
        "--mol",
        "--mol-file",
        metavar="MOL-FILE",
        type=str,
        required=True,
        help="input molecule file",
    )
    io_group.add_argument(
        "--helm",
        "--helm-file",
        metavar="HELM-FILE",
        type=str,
        required=True,
        help="outout file of HELM string",
    )
    io_group.add_argument(
        "--failures",
        metavar="MOL-FILE",
        type=str,
        required=False,
        help="output file for failed HELM generation (required: %(required)s)",
    )

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

    helm_gen_group = parser.add_argument_group("HELM generation options")
    _add_helm_generation_options(helm_gen_group)

    verbose_group = parser.add_argument_group("Verbose options")
    verbose_group.add_argument(
        "--show-summary",
        default=False,
        action="store_true",
        help="show summary of HELM generation attempts and unmapped fragments (default: %(default)s)",
    )

    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:
    """Convert molecule file to HELM file."""
    args = parse_options()
    console = rich.console.Console(record=args.save_console_svg)

    monomers = _get_monomer_collection(args)
    code_set = args.code_set or monomers.GetPrimaryCodeSet()

    helm_filepath = pathlib.Path(args.helm)
    if helm_filepath.suffix.lower() != ".helm":
        oechem.OEThrow.Fatal("Invalid file extension expected .helm!")

    failures_ofs: oechem.oemolostream | None = None
    if args.failures:
        failures_ofs = oechem.oemolostream()
        if not failures_ofs.open(args.failures):
            console.print(
                f"[red]Error: Unable to open failures file '{args.failures}'![/red]"
            )
            return os.EX_DATAERR

    options = oechem.OEHelmGenerationOptions(code_set)
    options.SetAllowUnspecifiedStereo(args.allow_unspecified_stereo)
    options.SetAllowUnmatchedFragments(args.allow_unmatched_fragments)

    mol_database = oechem.OEMolDatabase()
    if not mol_database.Open(args.mol):
        console.print(f"[red]Error: Unable to open molecule file '{args.mol}'![/red]")
        return os.EX_DATAERR

    unmapped_fragments: list[str] = []
    return_codes: list[int] = []
    num_failures = 0
    num_molecules = mol_database.NumMols()

    with (
        Progress(
            TextColumn("{task.description}"),
            BarColumn(bar_width=60),
            TextColumn("{task.percentage:3.1f}%"),
            transient=True,
            disable=num_molecules < 10,  # noqa: PLR2004
            console=console,
        ) as progress,
        helm_filepath.open("w") as helm_file,
    ):
        conversion = progress.add_task("[blue]HELM generation", total=num_molecules)
        for idx in range(num_molecules):
            mol = oechem.OEGraphMol()
            if not mol_database.GetMolecule(mol, idx):
                console.print(
                    f"[red]Error: Unable to get molecule at index {idx}![/red]"
                )
                continue
            result = oechem.OEHelmGenerationResult()
            if helm := oechem.OEMolToHelm(mol, monomers, options, result):
                helm_file.write(helm + "\n")
            else:
                num_failures += 1
                unmapped_fragments.extend(list(result.GetUnmappedFragments()))
                if failures_ofs is not None:
                    oechem.OEWriteMolecule(failures_ofs, mol)

            return_codes.append(result.GetReturnCode())

            progress.update(conversion, advance=1)

    num_successes = mol_database.NumMols() - num_failures
    num_monomers = oechem.OECount(monomers, oechem.OEIsInMonomerCodeSet(code_set))
    console.print(
        f"[yellow]HELM generation completed with {num_failures} failures and {num_successes} successes "
        f"using {num_monomers} monomers from `{code_set}` code-set.[/yellow]"
    )
    if args.show_summary:
        _show_summary(return_codes, unmapped_fragments, console)

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

    return os.EX_OK


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 = ["Standard", "OpenEye", "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 ["Standard", "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
        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]")
            return False
        return True


def _add_monomer_collection(arg_group: argparse._ArgumentGroup) -> None:
    arg_group.add_argument(
        "-m",
        "--monomers",
        type=str,
        default="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 _add_helm_generation_options(arg_group: argparse._ArgumentGroup) -> None:
    arg_group.add_argument(
        "--allow-unspecified-stereo",
        default=False,
        action="store_true",
        help="allow unspecified stereo in input molecule (default: %(default)s)",
    )
    arg_group.add_argument(
        "--allow-unmatched-fragments",
        default=False,
        action="store_true",
        help="allow unmatched fragments in input molecule -- embedded SMILES in HELM (default: %(default)s)",
    )


def _show_summary(
    return_codes: list[int],
    unmapped_fragments: list[str],
    console: rich.console.Console,
) -> None:
    return_code_map = {
        oechem.OEHelmGenerationReturnCode_Success: "[green]success[/green]",
        oechem.OEHelmGenerationReturnCode_UnspecifiedAtomStereo: "[yellow]unspecified atom stereo[/yellow]",
        oechem.OEHelmGenerationReturnCode_UnspecifiedBondStereo: "[yellow]unspecified bond stereo[/yellow]",
        oechem.OEHelmGenerationReturnCode_UnmappedFragments: "[yellow]unmapped fragments[/yellow]",
        oechem.OEHelmGenerationReturnCode_NoSequence: "[red]no sequence[/red]",
    }
    for code, message in return_code_map.items():
        count = return_codes.count(code)
        if count != 0:
            console.print(
                f"{count:5d} ({count / len(return_codes) * 100:5.2f}%) {message}"
            )

    # summarize unmapped fragments
    if not unmapped_fragments:
        return

    fragment_counts = collections.Counter(unmapped_fragments)
    total = sum(fragment_counts.values())

    console.print("\n[yellow]Unmapped fragments summary:[/yellow]")
    for frag, count in fragment_counts.most_common():
        console.print(f"{count:5d} ({count / total * 100:5.2f}%) {frag}")


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