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

"""Generate plot showing monomer similarity."""

import argparse
import json
import os
import pathlib
import sys

import numpy as np
import rich.console
from bokeh.models import ColumnDataSource, HoverTool, LabelSet
from bokeh.plotting import figure, output_file, save, show
from openeye import oechem, oedepict, oegrapheme, oegraphsim
from rich.progress import BarColumn, Progress, TextColumn
from rich_argparse import HelpPreviewAction, RichHelpFormatter
from sklearn.decomposition import PCA

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Plot similarity of monomers."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme", "oegraphsim"]
__SCRIPT_KEYWORDS__ = [
    "monomer",
    "peptide",
    "peptide-informatics",
    "similarity",
    "visualization",
]
__SCRIPT_CATEGORIES__ = ["peptide-informatics"]


def parse_options() -> argparse.Namespace:
    """Parse main options."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
    )

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

    image_group = parser.add_argument_group("Output options")
    image_group.add_argument(
        "--html",
        type=str,
        required=False,
        metavar="HTML-FILE",
        help="output HTML file (required: %(required)s) -- if no output is provided the image will be displayed in the browser",
    )
    image_group.add_argument(
        "--width",
        type=int,
        default=900,
        help="width of output image (default: %(default)s)",
    )
    image_group.add_argument(
        "--height",
        type=int,
        default=900,
        help="height of output image (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:
    """Console script for monomer explorer."""
    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()

    monomer_pred = oechem.OEAndMonomer(
        oechem.OEHasPolymerType(oechem.OEPolymerType_Peptide),
        oechem.OEIsInMonomerCodeSet(code_set),
    )

    num_peptide_monomers = oechem.OECount(monomers, monomer_pred)
    fp_type = oegraphsim.OEGetFPType(oegraphsim.OEFPType_Tree)

    monomer_list: list[oechem.OEMonomer] = list(monomers.GetMonomers(monomer_pred))
    similarity_bits: np.ndarray = np.empty(shape=(num_peptide_monomers, 4096))
    svg_images: list[str] = []

    with Progress(
        TextColumn("{task.description}"),
        BarColumn(bar_width=60),
        TextColumn("{task.percentage:3.1f}%"),
        transient=True,
        console=console,
    ) as progress:
        monomer_prep = progress.add_task(
            "[blue]Preparing monomers", total=num_peptide_monomers
        )
        for idx, monomer in enumerate(monomers.GetMonomers(monomer_pred)):
            image, fingerprint = _generate_monomer_image_and_fingerprint(
                monomer, fp_type
            )
            svg_images.append(image)
            bits = [
                1 if fingerprint.IsBitOn(b) else 0 for b in range(fingerprint.GetSize())
            ]
            similarity_bits[idx] = np.array(bits)
            progress.update(monomer_prep, advance=1)

    # # prepare coordinates for monomer plot
    coords = PCA(n_components=2, random_state=0).fit_transform(similarity_bits)
    coords_x = [x for x, _ in coords]
    coords_y = [y for _, y in coords]

    data = {
        "x": coords_x,
        "y": coords_y,
        "code": [m.GetCode(code_set) for m in monomer_list],
        "color": [
            (
                "#1E7614"
                if oechem.OEIsStandardMonomer(m.GetCanonicalSmiles())
                else "#3b3b3b"
            )
            for m in monomer_list
        ],
        "marker": [
            "square" if oechem.OEIsStandardMonomer(m.GetCanonicalSmiles()) else "circle"
            for m in monomer_list
        ],
        "size": [
            20 if oechem.OEIsStandardMonomer(m.GetCanonicalSmiles()) else 5
            for m in monomer_list
        ],
        "label": [
            (
                m.GetCode(code_set)
                if oechem.OEIsStandardMonomer(m.GetCanonicalSmiles())
                else ""
            )
            for m in monomer_list
        ],
        "svg_image": svg_images,
    }
    source = ColumnDataSource(data=data)

    tooltips = """
    <div>
       <div>
       <span style="font-size: 20px;">Code: @code</span>
       </div>
       <img style="float: left; margin: 5px 5px 5px 5px; width:250px;">
          @svg_image{safe}
       </img>
    </div>
    """

    fig = figure(
        title=f"{num_peptide_monomers} peptide monomers in '{code_set}' code set",
        width=args.width,
        height=args.height,
        margin=(50, 50, 50, 50),
    )

    fig.scatter(
        x="x",
        y="y",
        fill_color="color",
        line_color=None,
        marker="marker",
        size="size",
        fill_alpha=0.5,
        source=source,
    )

    labels = LabelSet(
        x="x",
        y="y",
        text="label",
        text_align="center",
        text_font_size="16pt",
        x_offset=0,
        y_offset=6,
        source=source,
    )
    fig.add_layout(labels, place="above")
    fig.add_tools(HoverTool(tooltips=tooltips))

    if not args.html or pathlib.Path(args.html).suffix.lower() != ".html":
        show(fig)
    else:
        console.print(f"[blue]Saving plot to {args.html} ...[/blue]")
        output_file(args.html, title="Monomer Similarity", mode="inline")
        save(fig, args.html)

    return os.EX_OK


def _generate_monomer_image_and_fingerprint(
    monomer: oechem.OEMonomer, fp_type: oegraphsim.OEFPTypeBase
) -> tuple[str, oegraphsim.OEFingerPrint]:
    mol: oechem.OEMolBase = oechem.OEGraphMol()
    oechem.OESmilesToMol(mol, monomer.GetCanonicalSmiles())
    fingerprint: oegraphsim.OEFingerPrint = oegraphsim.OEFingerPrint()
    oegraphsim.OEMakeFP(fingerprint, mol, fp_type)

    image = oedepict.OEImage(600, 600)
    oegrapheme.OEDrawMonomer(image, monomer)
    svg = oedepict.OEWriteImageToString("svg", image).decode("utf-8")

    return svg, fingerprint


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