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

"""Depict connected fragment combinations in a multi-page report."""

import argparse
import enum
import os
import pathlib
import sys
from collections.abc import Callable, Iterator
from itertools import combinations

from openeye import oechem, oedepict, oegrapheme, oemedchem
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict connected fragment combinations in a multi-page report."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme", "oemedchem"]
__SCRIPT_CATEGORIES__ = ["depiction"]


def parse_args() -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
    )

    input_group = parser.add_argument_group("Input options")
    input_group.add_argument(
        "--mol",
        type=str,
        required=True,
        metavar="MOL-FILE",
        help="input molecule file",
    )

    report_group = parser.add_argument_group("Report options")
    report_group.add_argument(
        "--report",
        type=str,
        required=True,
        metavar="REPORT-FILE",
        help="output report file (PDF)",
    )
    report_group.add_argument(
        "--rows",
        type=int,
        default=3,
        choices=range(2, 6),
        metavar="N",
        help="number of rows per page (default: %(default)s)",
    )
    report_group.add_argument(
        "--cols",
        type=int,
        default=2,
        choices=range(1, 3),
        metavar="N",
        help="number of columns per page (default: %(default)s)",
    )
    report_group.add_argument(
        "--page-by-page",
        action="store_true",
        help="write pages of report to separate numbered image files",
    )

    frag_group = parser.add_argument_group("Fragmentation options")
    frag_group.add_argument(
        "--frag-type",
        type=FragmentationType,
        default=FragmentationType.FunctionalGroup,
        choices=list(FragmentationType),
        help="fragmentation type (default: %(default)s)",
    )

    parser.add_argument("--help-image", action=HelpPreviewAction)
    return parser.parse_args()


def main() -> int:
    """Depict connected fragment combinations in a multi-page report."""
    args = parse_args()
    _check_report_file(args)

    # open input file
    input_stream = oechem.oemolistream()
    if not input_stream.open(args.mol):
        oechem.OEThrow.Fatal("Cannot open input file!")

    # read a molecule
    mol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(input_stream, mol):
        oechem.OEThrow.Fatal("Cannot read input file!")
    oedepict.OEPrepareDepiction(mol)

    # initialize fragmentation function
    frag_func = _get_fragmentation_function(args.frag_type)

    # initialize multi-page report
    report_options = oedepict.OEReportOptions(args.rows, args.cols)
    report_options.SetFooterHeight(25.0)
    report_options.SetHeaderHeight(report_options.GetPageHeight() / 4.0)
    report = oedepict.OEReport(report_options)

    # setup depiction options
    display_options = oedepict.OE2DMolDisplayOptions()
    cell_width, cell_height = report.GetCellWidth(), report.GetCellHeight()
    display_options.SetDimensions(cell_width, cell_height, oedepict.OEScale_AutoScale)
    display_options.SetTitleLocation(oedepict.OETitleLocation_Hidden)
    display_options.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)
    display_options.SetAtomLabelFontScale(1.2)

    # depict molecule with fragment combinations
    depict_molecule_with_fragment_combinations(report, mol, frag_func, display_options)

    if args.page_by_page:
        oedepict.OEWriteReportPageByPage(args.report, report)
    else:
        oedepict.OEWriteReport(args.report, report)

    return os.EX_OK


def depict_molecule_with_fragment_combinations(
    report: oedepict.OEReport,
    mol: oechem.OEMolBase,
    frag_func: Callable,
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """
    Depict molecule with all adjacent fragment combinations.

    Fragments the molecule, generates all connected fragment pair combinations,
    and renders each combination into a report cell with highlighted fragments
    and faded non-fragment regions. The original fragmentation is shown in
    each page header.
    """
    # fragment molecule
    frags = list(frag_func(mol))

    # assign fragment indexes
    str_tag = "fragment idx"
    int_tag = oechem.OEGetTag(str_tag)
    for frag_idx, frag in enumerate(frags):
        for bond in frag.GetBonds():
            bond.SetData(int_tag, frag_idx)

    # setup depiction styles
    num_frags = len(frags)
    colors = list(oechem.OEGetLightColors())
    if len(colors) < num_frags:
        colors = list(
            oechem.OEGetColors(oechem.OEYellowTint, oechem.OEDarkOrange, num_frags)
        )

    bond_glyph = ColorBondByFragmentIndex(colors, int_tag)

    line_width_scale = 0.75
    fade_highlight = oedepict.OEHighlightByColor(oechem.OEGrey, line_width_scale)

    # generate adjacent fragment combinations
    frag_combs = get_fragment_atom_bond_set_combinations(frags)

    # depict each fragment combination
    for frag in frag_combs:
        cell = report.NewCell()
        disp = oedepict.OE2DMolDisplay(mol, opts)

        frag_atoms = oechem.OEIsAtomMember(frag.GetAtoms())
        frag_bonds = oechem.OEIsBondMember(frag.GetBonds())

        not_frag_atoms = oechem.OENotAtom(frag_atoms)
        not_frag_bonds = oechem.OENotBond(frag_bonds)

        oedepict.OEAddHighlighting(disp, fade_highlight, not_frag_atoms, not_frag_bonds)
        oegrapheme.OEAddGlyph(disp, bond_glyph, frag_bonds)

        oedepict.OERenderMolecule(cell, disp)

    # depict original fragmentation in each header
    cell_width, cell_height = report.GetHeaderWidth(), report.GetHeaderHeight()
    opts.SetDimensions(cell_width, cell_height, oedepict.OEScale_AutoScale)
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)
    disp = oedepict.OE2DMolDisplay(mol, opts)
    oegrapheme.OEAddGlyph(disp, bond_glyph, oechem.IsTrueBond())

    header_pen = oedepict.OEPen(
        oechem.OEWhite, oechem.OELightGrey, oedepict.OEFill_Off, 2.0
    )
    for header in report.GetHeaders():
        oedepict.OERenderMolecule(header, disp)
        oedepict.OEDrawBorder(header, header_pen)


class ColorBondByFragmentIndex(oegrapheme.OEBondGlyphBase):
    """Bond glyph that colors bonds by their fragment index."""

    def __init__(self, color_list: list, tag: int) -> None:
        """Initialize with color list and fragment tag."""
        oegrapheme.OEBondGlyphBase.__init__(self)
        self.color_list = color_list
        self.tag = tag

    def RenderGlyph(  # noqa: N802
        self, disp: oedepict.OE2DMolDisplay, bond: oechem.OEBondBase
    ) -> bool:
        """Render a colored line glyph on a bond."""
        bond_disp = disp.GetBondDisplay(bond)
        if bond_disp is None or not bond_disp.IsVisible():
            return False

        if not bond.HasData(self.tag):
            return False

        linewidth = disp.GetScale() / 2.0
        color = self.color_list[bond.GetData(self.tag)]
        pen = oedepict.OEPen(color, color, oedepict.OEFill_Off, linewidth)

        atom_disp_bgn = disp.GetAtomDisplay(bond.GetBgn())
        atom_disp_end = disp.GetAtomDisplay(bond.GetEnd())

        layer = disp.GetLayer(oedepict.OELayerPosition_Below)
        layer.DrawLine(atom_disp_bgn.GetCoords(), atom_disp_end.GetCoords(), pen)

        return True

    def CreateCopy(self):  # noqa: ANN201, N802
        """Create a copy of this glyph."""
        return ColorBondByFragmentIndex(self.color_list, self.tag).__disown__()


def _is_adjacent_atom_bond_sets(
    frag_a: oechem.OEAtomBondSet,
    frag_b: oechem.OEAtomBondSet,
) -> bool:
    """Check if two atom/bond sets share an adjacent atom pair."""
    for atom_a in frag_a.GetAtoms():
        for atom_b in frag_b.GetAtoms():
            if atom_a.GetBond(atom_b) is not None:
                return True
    return False


def _is_adjacent_atom_bond_set_combination(
    frag_list: tuple[oechem.OEAtomBondSet, ...],
) -> bool:
    """Check if a combination of fragments forms a single connected component."""
    parts = [0] * len(frag_list)
    num_parts = 0

    for idx, frag in enumerate(frag_list):
        if parts[idx] != 0:
            continue

        num_parts += 1
        parts[idx] = num_parts
        _traverse_fragments(frag, frag_list, parts, num_parts)

    return num_parts == 1


def _traverse_fragments(
    act_frag: oechem.OEAtomBondSet,
    frag_list: tuple[oechem.OEAtomBondSet, ...],
    parts: list[int],
    num_parts: int,
) -> None:
    """Recursively traverse adjacent fragments to find connected components."""
    for idx, frag in enumerate(frag_list):
        if parts[idx] != 0:
            continue

        if not _is_adjacent_atom_bond_sets(act_frag, frag):
            continue

        parts[idx] = num_parts
        _traverse_fragments(frag, frag_list, parts, num_parts)


def _combine_and_connect_atom_bond_sets(
    frag_list: tuple[oechem.OEAtomBondSet, ...],
) -> oechem.OEAtomBondSet:
    """Combine fragment atom/bond sets and add connecting bonds."""
    combined = oechem.OEAtomBondSet()
    for frag in frag_list:
        for atom in frag.GetAtoms():
            combined.AddAtom(atom)
        for bond in frag.GetBonds():
            combined.AddBond(bond)

    # add connecting bonds
    for atom_a in combined.GetAtoms():
        for atom_b in combined.GetAtoms():
            if atom_a.GetIdx() < atom_b.GetIdx():
                continue

            bond = atom_a.GetBond(atom_b)
            if bond is None:
                continue
            if combined.HasBond(bond):
                continue

            combined.AddBond(bond)

    return combined


def get_fragment_atom_bond_set_combinations(
    frag_list: list[oechem.OEAtomBondSet],
) -> list[oechem.OEAtomBondSet]:
    """Generate all adjacent connected fragment combinations."""
    frag_combs: list[oechem.OEAtomBondSet] = []

    num_frags = len(frag_list)
    for n in range(2, num_frags):
        for frag_comb in combinations(frag_list, n):
            if _is_adjacent_atom_bond_set_combination(frag_comb):
                frag = _combine_and_connect_atom_bond_sets(frag_comb)
                frag_combs.append(frag)

    return frag_combs


class FragmentationType(enum.Enum):
    """Molecule fragmentation type."""

    FunctionalGroup = "func-group"
    RingChain = "ring-chain"
    RingLinkerSideChain = "ring-linker-sidechain"

    def __str__(self) -> str:
        """Convert to string representation."""
        return self.value


def _get_fragmentation_function(
    frag_type: FragmentationType,
) -> Callable[[oechem.OEMolBase], Iterator[oechem.OEAtomBondSet]]:
    """Return the fragmentation function for the given type."""
    match frag_type:
        case FragmentationType.RingChain:
            return oemedchem.OEGetRingChainFragments
        case FragmentationType.RingLinkerSideChain:
            return oemedchem.OEGetRingLinkerSideChainFragments
    return oemedchem.OEGetFuncGroupFragments


def _check_report_file(args: argparse.Namespace) -> None:
    """Validate report file format and adjust page-by-page if needed."""
    ext = pathlib.Path(args.report).suffix[1:]
    if not args.page_by_page and not oedepict.OEIsRegisteredMultiPageImageFile(ext):
        oechem.OEThrow.Warning("Report will be generated into separate pages!")
        args.page_by_page = True


setattr(main, "__SCRIPT_NAME__", __SCRIPT_NAME__)
setattr(main, "__SCRIPT_DESC__", __SCRIPT_DESC__)
setattr(main, "__SCRIPT_TOOLKITS__", __SCRIPT_TOOLKITS__)
setattr(main, "__SCRIPT_CATEGORIES__", __SCRIPT_CATEGORIES__)

if __name__ == "__main__":
    sys.exit(main())