#!/usr/bin/env python3
# (C) 2023 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 thirst-party monomer format to OpenEye's json format."""

import argparse
import json
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Convert monomers to OpenEye's json format."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = ["monomer", "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]",
    )
    monomers_group = parser.add_argument_group("Monomer set options")
    monomers_group.add_argument(
        "--code-set",
        type=str,
        default="CUSTOM",
        help="name of the code set in the generated monomer set (default: %(default)s)",
    )
    monomers_group.add_argument(
        "--monomer-set-version",
        metavar="X.Y.Z",
        type=str,
        default="1.0.0",
        help="version of generated monomer set (default: %(default)s)",
    )

    io_group = parser.add_argument_group("Input/Output options")
    io_group.add_argument(
        "--in-json",
        metavar="JSON-FILE",
        type=str,
        required=True,
        help="input monomer file Pistoia format (json)",
    )
    io_group.add_argument(
        "--out-json",
        metavar="JSON-FILE",
        type=str,
        required=True,
        help="output molecule set file OpenEye format (json)",
    )

    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 thirst-party monomer format to OpenEye's json format."""
    args = parse_options()

    console = rich.console.Console(record=args.save_console_svg, highlight=False)

    _check_file(
        pathlib.Path(args.out_json), check_exists=False, expected_extension=".json"
    )
    code_set = args.code_set
    monomers = oechem.OEMonomerSet()
    monomers.SetVersion(args.monomer_set_version)
    monomers.AddCodeSet(code_set)
    console.print(f"[blue]Adding monomers to '{code_set}' code set![/blue]")

    if pathlib.Path(args.out_json).suffix.lower() != ".json":
        oechem.OEThrow.Fatal(
            f"Output monomer set file {args.out_json} requires json extension!"
        )

    monomer_json_data: None | list[dict] = _load_monomer_json_data(
        pathlib.Path(args.in_json), console
    )
    if monomer_json_data is None or len(monomer_json_data) == 0:
        # warning/error was already thrown
        return os.EX_DATAERR
    console.print(
        f"{len(monomer_json_data)} monomer definitions are loaded for processing!"
    )

    num_RNAs = [  # noqa: N806
        m.get(POLYMER_TYPE_TAG, "") for m in monomer_json_data
    ].count("RNA")
    console.print(f"{num_RNAs} RNA monomers are ignored!")

    for idx, m_json in enumerate(monomer_json_data):
        if (
            monomer_data := convert_to_monomer_data(idx, m_json, code_set, console)
        ) is not None:
            code = monomer_data.GetCode(code_set)
            can_smiles = monomer_data.GetCanonicalSmiles()
            if monomers.HasCodeInCodeSet(code_set, code):
                console.print(
                    f"[red]{idx:3} ❌ monomer code '{code}' already exists in '{code_set}'[/red]!"
                )
                continue
            if monomers.HasMonomer(monomer_data):
                console.print(
                    f"[red]{idx:3} ❌ monomer '{code}' {can_smiles} already exists in monomer set![/red]"
                )
                continue
            if monomers.AddMonomer(monomer_data):
                polymer_type = oechem.OEPolymerTypeToString(
                    monomer_data.GetPolymerType()
                )
                console.print(
                    f"{idx:3} ✅ {code:20s} as {polymer_type:8s} {monomer_data.GetSmiles()}"
                )
            else:
                console.print(
                    f"[red]{idx:3} ❌ adding monomer {can_smiles} to set is failed!s[/red]"
                )

    oechem.OEWriteMonomerSet(monomers, args.out_json)

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


def _check_file(
    filepath: pathlib.Path, check_exists: bool, expected_extension: str
) -> None:
    if check_exists and not filepath.exists():
        oechem.OEThrow.Fatal(f"Can not open file {filepath!s}!")
    if filepath.suffix.lower() != expected_extension:
        oechem.OEThrow.Fatal(
            f"Invalid file extension {filepath.suffix} (expected {expected_extension})!"
        )


MONOMER_SMILES_TAG = "smiles"
MONOMER_SYMBOL_TAG = "symbol"
POLYMER_TYPE_TAG = "polymerType"  # expected values PEPTIDE, CHEM, RNA
MONOMER_TYPE_TAG = "monomerType"
REQUIRED_TAGS = [
    MONOMER_SMILES_TAG,
    MONOMER_SYMBOL_TAG,
    POLYMER_TYPE_TAG,
    MONOMER_TYPE_TAG,
]
MONOMER_NAME_TAG = "name"

OPTIONAL_TAGS = [MONOMER_NAME_TAG]


def convert_to_monomer_data(  # noqa: C901, PLR0911, PLR0912
    idx: int, monomer_def: dict, code_set: str, console: rich.console.Console
) -> None | oechem.OEMonomerData:
    """Convert PistoiaHELM json format to oechem.OEMonomerData."""
    if not isinstance(monomer_def, dict):
        console.print(f"[red]{idx:2} Invalid monomer definition format![red]")
        return None

    for tag_name in REQUIRED_TAGS:
        if (monomer_def.get(tag_name, "")) == "":
            console.print(f"[red] {idx:2d} missing required {tag_name}[/red]!")
            return None

    code = monomer_def[MONOMER_SYMBOL_TAG]

    polymer_type = monomer_def[POLYMER_TYPE_TAG]
    if polymer_type not in ["PEPTIDE", "CHEM"]:
        return None
    monomer_type = monomer_def[MONOMER_TYPE_TAG]
    if monomer_type not in ["Backbone", "Terminal"]:
        console.print(
            f"[red] {idx:2d} ❌ {code:10s} monomer with '{monomer_type}' monomer type is not supported in OEChem[/red]!"
        )
        return None

    smiles = monomer_def[MONOMER_SMILES_TAG]
    monomer_mol = oechem.OEGraphMol()
    if not oechem.OESmilesToMol(monomer_mol, smiles):
        console.print(f"[red] {idx:2d} problem parsing {smiles} [/red]!")

    oechem_polymer_type = oechem.OEGetPolymerType(monomer_mol)
    oechem_monomer_type = oechem.OEGetMonomerType(monomer_mol)
    if oechem_polymer_type not in [
        oechem.OEPolymerType_Peptide,
        oechem.OEPolymerType_Chem,
    ]:
        console.print(
            f"[red]{idx:3} ❌ {code:10s} {smiles} is rejected due to wrong polymer type![/red]"
        )
        return None

    monomer_data = oechem.OEMonomerData(smiles)
    monomer_data.AddCode(code_set, code)
    monomer_data.SetPolymerType(oechem_polymer_type)
    monomer_data.SetMonomerType(oechem_monomer_type)
    if (name := monomer_def[MONOMER_NAME_TAG]) != "":
        monomer_data.SetName(name)

    result = oechem.OEMonomerValidationResult()
    if not oechem.OEIsValidMonomerData(monomer_data, result):
        if result.GetReturnCode() in [
            oechem.OEMonomerValidationReturnCode_UnspecifiedAtomStereo,
            oechem.OEMonomerValidationReturnCode_UnspecifiedBondStereo,
        ]:
            return None
        console.print(
            f"[red]{idx:3} ❌ {code:10s} {smiles} is rejected due to {result.GetWarning()}[/red]"
        )
        return None

    oechem_str = oechem.OEPolymerTypeToString(oechem_polymer_type)
    if (
        oechem_polymer_type == oechem.OEPolymerType_Peptide
        and polymer_type != "PEPTIDE"
    ):
        console.print(
            f"[#ff8787]{idx:3}    {code:10s} is supported as {oechem_str} rather than {polymer_type} {smiles} ![/#ff8787]"
        )
    if oechem_polymer_type == oechem.OEPolymerType_Chem and polymer_type != "CHEM":
        console.print(
            f"[#ff8787]{idx:3}    {code:10s} is supported as {oechem_str} rather than {polymer_type} {smiles}![/#ff8787]"
        )

    oechem_str = oechem.OEMonomerTypeToString(oechem_monomer_type)

    if (
        oechem_monomer_type == oechem.OEMonomerType_Backbone
        and monomer_type != "Backbone"
    ):
        console.print(
            f"[#ff8787]{idx:3}    {code:10s} is supported as {oechem_str} rather than {monomer_type} {smiles} ![/#ff8787]"
        )
    if (
        oechem_monomer_type == oechem.OEMonomerType_Terminal
        and monomer_type != "Terminal"
    ):
        console.print(
            f"[#ff8787]{idx:3}    {code:10s} is supported as {oechem_str} rather than {monomer_type} {smiles}![/#ff8787]"
        )
    return monomer_data


def _load_monomer_json_data(
    filepath: pathlib.Path, console: rich.console.Console
) -> None | list:
    if not filepath.exists():
        console.print(f"[red] Can not open file {filepath!s} file for reading![red]")
        return None
    if filepath.suffix.lower() != ".json":
        oechem.OEThrow.Fatal(f"Invalid file extension {filepath.suffix}!")
        return None
    try:
        with pathlib.Path(filepath).open() as in_json_file:
            monomer_json_data = json.load(in_json_file)
    except json.JSONDecodeError:
        console.print(
            f"Failed to decode JSON from the file '{filepath!s}'. Check if the file has valid JSON format."
        )
        return None

    if not isinstance(monomer_json_data, list):
        console.print("[red]Invalid monomer json file format![red]")
        return None

    for m in monomer_json_data:
        if not isinstance(m, dict):
            console.print("[red]Invalid monomer json file format![red]")
            return None
    return monomer_json_data


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