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

"""Substructure search in monomer set."""

import argparse
import json
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Substructure search in monomer set."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = [
    "monomer",
    "peptide",
    "peptide-informatics",
    "search",
    "substructure",
]
__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]",
    )

    search_group = parser.add_argument_group("Search options")
    search_group.add_argument(
        "--smarts",
        type=str,
        metavar="SMARTS-PATTERN",
        required=True,
        help="SMARTS pattern used in substructure search",
    )
    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:
    """Search substructure in monomer set."""
    args = parse_options()

    console = rich.console.Console(record=args.save_console_svg, highlight=False)
    sub_search = oechem.OESubSearch(args.smarts)
    if not sub_search.IsValid():
        console.print(f"[red]Error: invalid SMARTS pattern '{args.smarts}'![/red]")
        return os.EX_DATAERR

    monomers = _get_monomer_collection(args)

    primary_code_set = monomers.GetPrimaryCodeSet()
    code_sets: list[str] = [
        primary_code_set,
        *[c for c in monomers.GetCodeSets() if c != primary_code_set],
    ]
    columns = ["idx", *code_sets, "monomer smiles"]
    table = rich.table.Table(*columns)

    monomer_mol = oechem.OEGraphMol()
    num_matches = 0
    for monomer in monomers.GetMonomers():
        oechem.OESmilesToMol(monomer_mol, monomer.GetCanonicalSmiles())
        oechem.OEPrepareSearch(monomer_mol, sub_search)
        if sub_search.SingleMatch(monomer_mol):
            num_matches += 1
            row_data = [
                f"{num_matches:2d}",
                *[monomer.GetCode(s) if monomer.HasCode(s) else "" for s in code_sets],
                monomer.GetCanonicalSmiles(),
            ]
            table.add_row(*row_data)

    console.print(table, markup=False)

    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",
    )


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