#!/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 a PDF report of library generation products from an MDL reaction."""

import argparse
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = (
    "Generate a PDF report of library generation products from an MDL reaction."
)
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict"]
__SCRIPT_CATEGORIES__ = ["depiction"]


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(
        "--rxn",
        metavar="REACTION-FILE",
        type=str,
        required=True,
        help="input MDL reaction file (rxn)",
    )
    io_group.add_argument(
        "--reactants",
        metavar="REACTANT-FILE",
        type=str,
        nargs="+",
        required=True,
        help="input reactant molecule files (.sdf, .mol, etc.) (one per reactant in reaction)",
    )

    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,
        help="number of rows per page (default: %(default)s)",
    )
    report_group.add_argument(
        "--cols",
        type=int,
        default=2,
        help="number of columns per page (default: %(default)s)",
    )
    report_group.add_argument(
        "--page-by-page",
        default=False,
        action="store_true",
        help="write individual numbered separate pages",
    )

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


def main() -> int:
    """Generate a PDF report of library generation products."""
    args = parse_options()
    _check_report_file(args)
    console = Console()
    page_by_page = args.page_by_page
    ext = oechem.OEGetFileExtension(args.report)
    if not page_by_page and not oedepict.OEIsRegisteredMultiPageImageFile(ext):
        console.print("[yellow]Report will be generated into separate pages.[/yellow]")
        page_by_page = True

    tag = oechem.OEGetTag("reactant_idx")

    reaction_mol = get_reaction(args.rxn)
    lib_gen = get_library_gen(reaction_mol, args.reactants, tag)
    sub_src = get_product_sub_search(reaction_mol)

    report_opts = oedepict.OEReportOptions(args.rows, args.cols)
    report_opts.SetHeaderHeight(report_opts.GetPageHeight() * 0.15)
    report = oedepict.OEReport(report_opts)

    width, height = report.GetCellWidth(), report.GetCellHeight()
    opts = oedepict.OE2DMolDisplayOptions(width, height, oedepict.OEScale_AutoScale)
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)
    opts.SetTitleLocation(oedepict.OETitleLocation_Hidden)

    colors = list(oechem.OEGetLightColors())

    depict_products(report, lib_gen, sub_src, tag, opts, colors)
    depict_reaction_in_headers(report, reaction_mol, colors)

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

    console.print(f"Report saved to [green]{args.report}[/green]")

    return os.EX_OK


def get_reaction(reaction_file: str) -> oechem.OEGraphMol:
    """Read an MDL reaction from file."""
    ifs = oechem.oemolistream()
    if not ifs.open(reaction_file):
        oechem.OEThrow.Fatal(f"Cannot open input reaction file: {reaction_file}")

    reaction_mol = oechem.OEGraphMol()
    if not oechem.OEReadMDLReactionQueryFile(ifs, reaction_mol):
        oechem.OEThrow.Fatal(f"Cannot read reaction from: {reaction_file}")

    return reaction_mol


def add_starting_materials(
    lib_gen: oechem.OELibraryGen, filename: str, reactant_idx: int, tag: int
) -> None:
    """Add starting materials from a file to the library generator."""
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        oechem.OEThrow.Fatal(f"Cannot open reactant file: {filename}")

    unique = True
    for mol in ifs.GetOEGraphMols():
        for atom in mol.GetAtoms():
            atom.SetData(tag, reactant_idx)
        for bond in mol.GetBonds():
            bond.SetData(tag, reactant_idx)
        lib_gen.AddStartingMaterial(mol, reactant_idx, unique)


def get_library_gen(
    reaction_mol: oechem.OEMolBase, reactant_files: list[str], tag: int
) -> oechem.OELibraryGen:
    """Initialize a library generator from a reaction and reactant files."""
    reaction = oechem.OEQMol()
    opts = oechem.OEMDLQueryOpts_ReactionQuery | oechem.OEMDLQueryOpts_SuppressExplicitH
    oechem.OEBuildMDLQueryExpressions(reaction, reaction_mol, opts)

    lib_gen = oechem.OELibraryGen()
    if not lib_gen.Init(reaction):
        oechem.OEThrow.Fatal("Failed to initialize library generator")
    lib_gen.SetValenceCorrection(True)
    lib_gen.SetExplicitHydrogens(False)

    if lib_gen.NumReactants() != len(reactant_files):
        oechem.OEThrow.Fatal(
            "Number of reactant files does not match number of reactants in reaction!"
        )

    for r_idx, fname in enumerate(reactant_files):
        add_starting_materials(lib_gen, fname, r_idx, tag)

    return lib_gen


def get_product_sub_search(reaction_mol: oechem.OEMolBase) -> oechem.OESubSearch:
    """Create a substructure search from the product side of a reaction."""
    oedepict.OEPrepareDepiction(reaction_mol)

    atom_pred = oechem.OEIsAtomMember(reaction_mol.GetAtoms(oechem.OEAtomIsInProduct()))
    product_mol = oechem.OEGraphMol()
    oechem.OESubsetMol(product_mol, reaction_mol, atom_pred, True)

    product = oechem.OEQMol()
    opts = oechem.OEMDLQueryOpts_ReactionQuery | oechem.OEMDLQueryOpts_SuppressExplicitH
    oechem.OEBuildMDLQueryExpressions(product, product_mol, opts)

    return oechem.OESubSearch(product)


def _get_product_atoms(
    reaction_mol: oechem.OEMolBase, reactant: oechem.OEAtomBondSet
) -> oechem.OEAtomBondSet:
    """Map reactant atoms to their corresponding product atoms."""
    product = oechem.OEAtomBondSet()
    for reactant_atom in reactant.GetAtoms():
        map_idx = reactant_atom.GetMapIdx()
        if map_idx != 0:
            product_atom = reaction_mol.GetAtom(
                oechem.OEAndAtom(
                    oechem.OEHasMapIdx(map_idx), oechem.OEAtomIsInProduct()
                )
            )
            if product_atom is not None:
                product.AddAtom(product_atom)
    return product


def _is_reactant_component(
    reaction_mol: oechem.OEMolBase, part_pred: oechem.OEPartPredAtom
) -> bool:
    """Check whether a component belongs to the reactant side."""
    atom = reaction_mol.GetAtom(
        oechem.OEAndAtom(part_pred, oechem.OEAtomIsInReactant())
    )
    return atom is not None


def _add_bonds_between_atoms(
    reaction_mol: oechem.OEMolBase, abset: oechem.OEAtomBondSet
) -> None:
    """Add all bonds between atoms already in the atom-bond set."""
    pred = oechem.OEIsAtomMember(abset.GetAtoms())
    for bond in reaction_mol.GetBonds():
        if pred(bond.GetBgn()) and pred(bond.GetEnd()):
            abset.AddBond(bond)


def color_reaction_by_components(
    disp: oedepict.OE2DMolDisplay, comp_colors: list[oechem.OEColor]
) -> None:
    """Color reactants and their corresponding products with matching colors."""
    reaction_mol = disp.GetMolecule()
    num_parts, parts = oechem.OEDetermineComponents(reaction_mol)
    part_pred = oechem.OEPartPredAtom(parts)

    highlight_style = oedepict.OEHighlightStyle_BallAndStick

    for part_idx, color in zip(range(1, num_parts + 1), comp_colors, strict=False):
        part_pred.SelectPart(part_idx)
        if _is_reactant_component(reaction_mol, part_pred):
            reactant = oechem.OEAtomBondSet(reaction_mol.GetAtoms(part_pred))
            _add_bonds_between_atoms(reaction_mol, reactant)
            oedepict.OEAddHighlighting(disp, color, highlight_style, reactant)

            product = _get_product_atoms(reaction_mol, reactant)
            _add_bonds_between_atoms(reaction_mol, product)
            oedepict.OEAddHighlighting(disp, color, highlight_style, product)


class HasAtomReactantNumber(oechem.OEUnaryAtomPred):
    """Predicate matching atoms tagged with a specific reactant index."""

    def __init__(self, idx: int, tag: int) -> None:
        """Initialize predicate."""
        oechem.OEUnaryAtomPred.__init__(self)
        self.idx = idx
        self.tag = tag

    def __call__(self, atom: oechem.OEAtomBase) -> bool:
        """Evaluate atom."""
        return atom.HasData(self.tag) and atom.GetData(self.tag) == self.idx

    def CreateCopy(self):  # noqa: ANN201, N802
        """Create copy."""
        return HasAtomReactantNumber(self.idx, self.tag).__disown__()


class HasBondReactantNumber(oechem.OEUnaryBondPred):
    """Predicate matching bonds tagged with a specific reactant index."""

    def __init__(self, idx: int, tag: int) -> None:
        """Initialize predicate."""
        oechem.OEUnaryBondPred.__init__(self)
        self.idx = idx
        self.tag = tag

    def __call__(self, bond: oechem.OEBondBase) -> bool:
        """Evaluate bond."""
        return bond.HasData(self.tag) and bond.GetData(self.tag) == self.idx

    def CreateCopy(self):  # noqa: ANN201, N802
        """Create copy."""
        return HasBondReactantNumber(self.idx, self.tag).__disown__()


def highlight_product_by_reactant(
    disp: oedepict.OE2DMolDisplay,
    colors: list[oechem.OEColor],
    num_reactants: int,
    tag: int,
) -> None:
    """Highlight product atoms/bonds by their originating reactant."""
    highlight = oedepict.OEHighlightByBallAndStick(oechem.OEWhite)

    mol = disp.GetMolecule()
    for idx, color in zip(range(num_reactants), colors, strict=False):
        reactant = oechem.OEAtomBondSet(
            mol.GetAtoms(HasAtomReactantNumber(idx, tag)),
            mol.GetBonds(HasBondReactantNumber(idx, tag)),
        )
        highlight.SetColor(color)
        oedepict.OEAddHighlighting(disp, highlight, reactant)


def depict_products(
    report: oedepict.OEReport,
    lib_gen: oechem.OELibraryGen,
    sub_src: oechem.OESubSearch,
    tag: int,
    opts: oedepict.OE2DMolDisplayOptions,
    colors: list[oechem.OEColor],
) -> None:
    """Generate and depict all library products in the report."""
    product_list: list[oechem.OEGraphMol] = [
        oechem.OEGraphMol(mol) for mol in lib_gen.GetProducts()
    ]
    unique = True
    for mol in product_list:
        miter = sub_src.Match(mol, unique)
        if miter.IsValid():
            match = miter.Target()
            oedepict.OEPrepareAlignedDepiction(mol, sub_src.GetPattern(), match)
        else:
            oedepict.OEPrepareDepiction(mol)

    mol_scale = float("inf")
    for mol in product_list:
        mol_scale = min(mol_scale, oedepict.OEGetMoleculeScale(mol, opts))
    opts.SetScale(mol_scale * 1.6)

    num_reactants = lib_gen.NumReactants()

    for mol in product_list:
        cell = report.NewCell()
        disp = oedepict.OE2DMolDisplay(mol, opts)
        highlight_product_by_reactant(disp, colors, num_reactants, tag)
        oedepict.OERenderMolecule(cell, disp)


def depict_reaction_in_headers(
    report: oedepict.OEReport,
    reaction_mol: oechem.OEMolBase,
    colors: list[oechem.OEColor],
) -> None:
    """Depict the reaction with colored components in report headers."""
    header_width = report.GetHeaderWidth()
    header_height = report.GetHeaderHeight()
    header_opts = oedepict.OE2DMolDisplayOptions(
        header_width, header_height, oedepict.OEScale_AutoScale
    )
    header_opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)
    disp = oedepict.OE2DMolDisplay(reaction_mol, header_opts)

    color_reaction_by_components(disp, colors)

    for header in report.GetHeaders():
        oedepict.OERenderMolecule(header, disp)
        oedepict.OEDrawCurvedBorder(header, oedepict.OELightGreyPen, 10)


def _check_report_file(args: argparse.Namespace) -> bool:
    ext = pathlib.Path(args.report).suffix[1:]
    if not oedepict.OEIsRegisteredImageFile(ext):
        oechem.OEThrow.Fatal("Unknown image outout type!")

    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

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