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

"""Print monomer list of HELMs to console."""

import argparse
import json
import os
import pathlib
import sys
from typing import NewType

import rich.box
import rich.console
import rich.table
import rich.text
from openeye import oechem
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Print monomers to console."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = ["HELM", "peptide", "sequence", "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 options")
    io_group.add_argument(
        "--helm",
        "--helm-file",
        metavar="HELM-FILE",
        type=str,
        required=True,
        help="input file of HELM string",
    )
    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()


MonomerSequence = NewType("MonomerSequence", list[str])


def main() -> int:
    """Print monomer list of HELMs to console."""
    args = parse_options()

    console = rich.console.Console(record=args.save_console_svg)
    monomers: oechem.OEMonomerSet = _get_monomer_collection(args)
    code_set = args.code_set or monomers.GetPrimaryCodeSet()

    helms_with_titles: list[tuple[str, str]] = _read_helms_with_title(args.helm)
    has_titles = any(title for _, title in helms_with_titles)
    console.print(
        f"{len(helms_with_titles)} HELMS read from {pathlib.Path(args.helm).name}!"
    )

    monomer_sequences_with_titles = _get_monomer_sequences(helms_with_titles)
    max_sequence_length = max(len(s) for s, _ in monomer_sequences_with_titles)
    monomer_color_dict: dict[str, str] = _get_monomer_colors(
        monomer_sequences_with_titles, monomers, code_set
    )

    table = rich.table.Table(
        title=f"[bold]Number of sequences: {len(monomer_sequences_with_titles)}[/bold]",
        box=rich.box.SIMPLE,
        collapse_padding=True,
        pad_edge=False,
    )
    table.add_column("idx", justify="right")
    for pos in range(1, max_sequence_length + 1):
        table.add_column(f"P{pos:02d}", justify="center")
    if has_titles:
        table.add_column("Title", justify="left")

    for idx, (sequence, title) in enumerate(monomer_sequences_with_titles, start=1):
        sequence_repr: list[rich.text.Text] = [
            rich.text.Text(str(idx)),
            *_get_colored_sequence(sequence, monomer_color_dict),
            rich.text.Text(title),
        ]
        table.add_row(*sequence_repr)

    console.print(table)

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


def _read_helms_with_title(helm_filename: str) -> list[tuple[str, str]]:
    # script will terminate if reading HELM string from file fails
    helm_filepath = pathlib.Path(helm_filename)
    if not helm_filepath.exists():
        oechem.OEThrow.Fatal(f"{helm_filename} file does not exist!")
    if helm_filepath.suffix.lower() != ".helm":
        oechem.OEThrow.Fatal("Invalid file extension expected .helm!")

    helms: list[tuple[str, str]] = []
    try:
        with helm_filepath.open() as helm_file:
            for line in helm_file:
                helm, _, title = line.rstrip().partition(" ")
                helms.append((helm, title))
    except OSError:
        oechem.OEThrow.Fatal(f"Can not open {helm_filename} input file!")

    if len(helms) == 0:
        oechem.OEThrow.Fatal(f"No helm string read from {helm_filename}!")
        return []
    return helms


def _get_monomer_colors(
    monomer_sequences_with_titles: list[tuple[MonomerSequence, str]],
    monomers: oechem.OEMonomerSet,
    code_set: str,
) -> dict[str, str]:
    unique_monomer_codes: set[str] = {
        c for s, _ in monomer_sequences_with_titles for c in s
    }

    monomer_colors: dict[str, str] = {}
    for code in unique_monomer_codes:
        if (monomer := monomers.GetMonomer(code_set, code)) is not None:
            color = _get_monomer_analog_color(monomer)
            monomer_colors[code] = color
    return monomer_colors


def _get_monomer_analog_color(monomer: oechem.OEMonomer) -> str:  # noqa: PLR0911
    if monomer.GetPolymerType() != oechem.OEPolymerType_Peptide:
        return "#000000"
    analog = oechem.OEGetStandardAnalog(monomer.GetCanonicalSmiles())
    if analog == oechem.OEResidueIndex_UNK:
        return "#AAAAAA"
    match analog:
        case oechem.OEResidueIndex_CYS | oechem.OEResidueIndex_MET:
            return "#e4e488"
        case (
            oechem.OEResidueIndex_ALA
            | oechem.OEResidueIndex_GLY
            | oechem.OEResidueIndex_ILE
            | oechem.OEResidueIndex_LEU
            | oechem.OEResidueIndex_PRO
            | oechem.OEResidueIndex_VAL
        ):
            return "#c09071"
        case (
            oechem.OEResidueIndex_PHE
            | oechem.OEResidueIndex_TRP
            | oechem.OEResidueIndex_TYR
        ):
            return "#5faf5f"
        case oechem.OEResidueIndex_ASP | oechem.OEResidueIndex_GLU:
            return "#e0ac70"
        case (
            oechem.OEResidueIndex_ARG
            | oechem.OEResidueIndex_HIS
            | oechem.OEResidueIndex_LYS
        ):
            return "#85b4e6"
        case oechem.OEResidueIndex_SER | oechem.OEResidueIndex_THR:
            return "#ff8787"
        case oechem.OEResidueIndex_ASN | oechem.OEResidueIndex_GLN:
            return "#5493EA"
    return "#AAAAAA"


def _get_monomer_sequences(
    helms_with_title: list[tuple[str, str]],
) -> list[tuple[MonomerSequence, str]]:
    monomer_sequences: list[tuple[MonomerSequence, str]] = []

    chain_name = "PEPTIDE1"
    for helm, title in helms_with_title:
        monomer_codes = oechem.OEGetHelmMonomerCodes(helm, chain_name)
        if len(monomer_codes) > 0:
            monomer_sequences.append((monomer_codes, title))

    return monomer_sequences


def _get_colored_sequence(
    monomer_sequence: MonomerSequence, monomer_color_dict: dict[str, str]
) -> list[rich.text.Text]:
    colored_sequence: list[rich.text.Text] = []
    for code in monomer_sequence:
        bg_color = monomer_color_dict.get(code, "white")
        style = rich.text.Style(color="black", bgcolor=bg_color, bold=True)
        colored_sequence.append(rich.text.Text(code, style, justify="center"))
    return colored_sequence


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


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