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

"""Print protein-ligand interactions to console."""

import argparse
import os
import pathlib
import sys
import typing
from collections import Counter

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__ = "Print protein-ligand interactions to console."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = [
    "perception",
    "active-site",
    "protein-ligand",
    "interactions",
    "display",
]
__SCRIPT_CATEGORIES__ = ["protein-ligand interactions"]


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

    # input options
    input_group = parser.add_argument_group("Input ligand-protein complex")
    exclusive_input_group = input_group.add_mutually_exclusive_group(required=True)
    exclusive_input_group.add_argument(
        "--complex",
        type=str,
        required=False,
        metavar="PDB/CIF-FILE",
        help="input file of the ligand-protein complex (.pdb, .cif)",
    )
    exclusive_input_group.add_argument(
        "--design-unit",
        "--du",
        type=str,
        metavar="DU-FILE",
        help="input design unit file (.oedu)",
    )
    exclusive_input_group.add_argument(
        "--serialized",
        type=str,
        metavar="OEB/JSON-FILE",
        help="input file with serialized interactions (.oeb or .json)",
    )
    exclusive_input_group.add_argument(
        "--protein",
        type=str,
        metavar="PDB-FILE",
        help="input protein file",
    )
    input_group.add_argument(
        "--ligand",
        type=str,
        metavar="MOL-FILE",
        help="input ligand file (required if --protein is provided)",
    )

    display_group = parser.add_argument_group("Display options")
    display_group.add_argument(
        "--disable-summery-table",
        action="store_true",
        help="skip summery interaction table",
    )
    ligand_group = parser.add_argument_group("Interaction selection")
    ligand_group.add_argument(
        "--smarts",
        type=str,
        metavar="SMARTS",
        help="show only interactions of ligand atoms matched to given pattern",
    )

    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:  # noqa: C901, PLR0912
    """Print interactions to console."""
    args = parse_options()

    if args.complex:
        active_site = get_active_site_from_pdb(args.complex)
    elif args.design_unit:
        active_site = get_active_site_from_design_unit(args.design_unit)
    elif args.serialized:
        active_site = get_active_site_from_serialized(args.serialized)
    elif args.protein:
        if not args.ligand:
            oechem.OEThrow.Fatal(
                "Please provide a ligand file (--ligand) when a protein file is provided!"
            )
        active_site = get_active_site_from_separate_files(args.protein, args.ligand)
    else:
        oechem.OEThrow.Fatal("Invalid input option!")

    console = rich.console.Console(record=args.save_console_svg)
    # Perceive interactions
    if active_site.NumInteractions() == 0 and not args.serialized:
        oechem.OEPerceiveInteractionHints(active_site)

    ligand: oechem.OEMolBase = active_site.GetMolecule(
        oechem.OELigandInteractionHintComponent()
    )
    console.print(f"Ligand detected = '{oechem.OEMolToSmiles(ligand)}'")

    only_ligand_atoms = oechem.OEAtomBondSet(ligand.GetAtoms())
    if args.smarts:
        sub_search = oechem.OESubSearch(args.smarts)
        if not sub_search.IsValid():
            oechem.OEThrow.Fatal(f"Unable to parse {args.smarts}'")
        console.print(f"Using '{args.smarts}' to select part of the ligand.")
        oechem.OEPrepareSearch(ligand, sub_search)
        unique = True
        only_ligand_atoms.Clear()
        for match in sub_search.Match(ligand, unique):
            only_ligand_atoms = oechem.OEAtomBondSet(match.GetTargetAtoms())
            break
        if only_ligand_atoms.NumAtoms() == 0:
            oechem.OEThrow.Fatal(
                "Substructure search was unsuccessful using '{smarts}'"
            )

    if not args.disable_summery_table:
        _print_summery_interaction_table(console, active_site)
    _print_ligand_interaction_table(console, active_site, only_ligand_atoms)

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

    return os.EX_OK


class InteractionHasLigandAtomPred(oechem.OEUnaryInteractionHintPred):
    """
    Predicate to identify interactions that involve only the specified ligand atoms.

    This class is used to filter interaction hints to those that contain at least one atom
    from a provided set of ligand atoms. It is useful for focusing on interactions relevant
    to a particular subset of ligand atoms, such as those matched by a SMARTS pattern.
    """

    def __init__(self, atom_set: oechem.OEAtomBondSet) -> None:
        """Initialize predicate."""
        oechem.OEUnaryInteractionHintPred.__init__(self)
        self._atoms_set = atom_set
        self._atom_pred = oechem.OEIsAtomMember(atom_set.GetAtoms())

    def __call__(self, inter: oechem.OEInteractionHint) -> bool:
        """Evaluate interaction."""
        for frag in [inter.GetBgnFragment(), inter.GetEndFragment()]:
            if frag.GetComponentType() == oechem.OELigandInteractionHintComponent():
                for a in inter.GetBgnFragment().GetAtoms():
                    if self._atom_pred(a):
                        return True
        return False

    def CreateCopy(self):  # noqa: ANN201, N802
        """Copy constructor."""
        return InteractionHasLigandAtomPred(self._atoms).__disown__()


def get_active_site_from_pdb(filename: str) -> oechem.OEInteractionHintContainer:
    """
    Initialize an active site interaction container from a PDB or CIF file.

    This function reads a ligand-protein complex from the specified file, separates the ligand and protein,
    and constructs an OEInteractionHintContainer for interaction analysis.
    """
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        oechem.OEThrow.Fatal("Unable to open {filename} for reading")

    if ifs.GetFormat() not in [oechem.OEFormat_PDB, oechem.OEFormat_CIF]:
        oechem.OEThrow.Fatal("Input file must have .pdb or .cif extension")

    complex_mol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ifs, complex_mol):
        oechem.OEThrow.Fatal("Unable to read complex from {filename}")

    if not oechem.OEHasResidues(complex_mol):
        oechem.OEPerceiveResidues(complex_mol, oechem.OEPreserveResInfo_All)

    # separate ligand and protein
    split_opts = oechem.OESplitMolComplexOptions()
    ligand = oechem.OEGraphMol()
    protein = oechem.OEGraphMol()
    water = oechem.OEGraphMol()
    other = oechem.OEGraphMol()

    split_opts.SetProteinFilter(
        oechem.OEOrRoleSet(split_opts.GetProteinFilter(), split_opts.GetWaterFilter())
    )
    split_opts.SetWaterFilter(
        oechem.OEMolComplexFilterFactory(oechem.OEMolComplexFilterCategory_Nothing)
    )

    oechem.OESplitMolComplex(ligand, protein, water, other, complex_mol, split_opts)

    if ligand.NumAtoms() == 0:
        oechem.OEThrow.Fatal("Cannot separate complex!")

    active_site = oechem.OEInteractionHintContainer(protein, ligand)

    if not oechem.OEIsValidActiveSite(active_site):
        oechem.OEThrow.Fatal("Cannot initialize active site!")

    return active_site


def get_active_site_from_design_unit(
    filename: str,
) -> oechem.OEInteractionHintContainer:
    """
    Initialize an active site interaction container from a design unit file.

    This function reads a design unit file, extracts the protein and ligand components,
    and constructs an OEInteractionHintContainer for interaction analysis.
    """
    design_unit = oechem.OEDesignUnit()
    if not oechem.OEIsReadableDesignUnit(filename) or not oechem.OEReadDesignUnit(
        filename, design_unit
    ):
        oechem.OEThrow.Fatal("Cannot read design unit.")

    active_site = oechem.OEInteractionHintContainer(design_unit)

    if not oechem.OEIsValidActiveSite(active_site):
        oechem.OEThrow.Fatal("Cannot initialize active site!")

    return active_site


def get_active_site_from_serialized(
    filename: str,
) -> oechem.OEInteractionHintContainer:
    """Initialize an active site interaction container from a serialized OEB/JSON file."""
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        oechem.OEThrow.Fatal(f"Unable to open {filename} for reading")

    if ifs.GetFormat() not in [oechem.OEFormat_OEB, oechem.OEFormat_JSON]:
        oechem.OEThrow.Fatal("Input file must have .oeb or .json extension")

    protein = oechem.OEGraphMol()
    ligand = oechem.OEGraphMol()
    # protein is expected to be first
    if not oechem.OEReadMolecule(ifs, protein) or not oechem.OEReadMolecule(
        ifs, ligand
    ):
        oechem.OEThrow.Fatal(f"Unable to read serialized interactions from {filename}")

    for mol in [protein, ligand]:
        if not oechem.OEHasInteractionsHintSerializationData(mol):
            oechem.OEThrow.Fatal("Interaction serialized data is missing!")

    active_site = oechem.OEInteractionHintContainer()
    if not oechem.OEConstructInteractionHintContainer(active_site, protein, ligand):
        oechem.OEThrow.Fatal("Failed to construct active site from serialized data!")
    if active_site.NumInteractions() == 0:
        oechem.OEThrow.Fatal("No interactions found in the serialized data!")
    return active_site


def get_active_site_from_separate_files(
    pro_filename: str, lig_filename: str
) -> oechem.OEInteractionHintContainer:
    """Read protein and ligand from separate PDB and MOL files."""
    protein = oechem.OEGraphMol()
    ligand = oechem.OEGraphMol()

    for filename, mol in [(pro_filename, protein), (lig_filename, ligand)]:
        ifs = oechem.oemolistream()
        if not ifs.open(filename):
            oechem.OEThrow.Fatal(f"Unable to open {filename} for reading")
        if not oechem.OEReadMolecule(ifs, mol):
            oechem.OEThrow.Fatal(f"Unable to read molecule from {filename}")
        oechem.OESetDimensionFromCoords(mol)

    active_site = oechem.OEInteractionHintContainer(protein, ligand)

    if not oechem.OEIsValidActiveSite(active_site):
        oechem.OEThrow.Fatal("Cannot initialize active site!")

    return active_site


def _print_summery_interaction_table(
    console: rich.console.Console, active_site: oechem.OEInteractionHintContainer
) -> None:
    columns = ["interaction name", "count"]
    columns += [
        "inter",
        "intra (pro)",
        "intra (lig)",
        "unpaired",
    ]
    table = rich.table.Table(
        *columns,
        title=f"[bold]Number of all interactions: {active_site.NumInteractions()} [/bold]",
    )

    for interaction_type in oechem.OEGetActiveSiteInteractionHintTypes():

        num_interactions, num_inter = 0, 0
        num_intra_pro, num_intra_lig, num_unpaired = 0, 0, 0

        for i in active_site.GetInteractions(
            oechem.OEHasInteractionHintType(interaction_type)
        ):
            num_interactions = num_interactions + 1
            component_type = i.GetBgnFragment().GetComponentType()
            if i.IsIntra():
                if i.GetBgnFragment().GetAtom(
                    oechem.OEIsTrueAtom()
                ) == i.GetEndFragment().GetAtom(oechem.OEIsTrueAtom()):
                    num_unpaired = num_unpaired + 1
                elif component_type == oechem.OELigandInteractionHintComponent():
                    num_intra_lig = num_intra_lig + 1
                elif component_type == oechem.OEProteinInteractionHintComponent():
                    num_intra_pro = num_intra_pro + 1
            if i.IsInter():
                num_inter = num_inter + 1

        if num_interactions != 0:
            inter_name = interaction_type.GetName().removeprefix("bio:active-site:")
            row_data: list[str] = [inter_name, str(num_interactions)]
            for inter_count in [num_inter, num_intra_pro, num_intra_lig, num_unpaired]:
                row_data.append("✅" if inter_count != 0 else "")  # noqa: PERF401

            table.add_row(*row_data)

    console.print(table)


class _InteractionSymbol(typing.NamedTuple):
    """Data structure representing a mapping between an interaction predicate, its name, and a display symbol."""

    pred: oechem.OEUnaryInteractionHintPred
    name: str
    symbol: str


class IsLigandAcceptsHBondInterInteraction(oechem.OEUnaryInteractionHintPred):
    """
    Predicate to identify hydrogen bond inter-molecular interaction where ligand accepts proton from protein atom.

    LigandAccepts - same as ProteinDonates
    NonIdealLigandAccepts - same as NonIdealProteinDonates
    """

    def __call__(self, inter: oechem.OEInteractionHint) -> bool:
        """Evaluate interaction."""
        if inter.IsIntra():
            return False
        accepted_types = [
            oechem.OEHBondInteractionHintType_LigandAccepts,
            oechem.OEHBondInteractionHintType_NonIdealLigandAccepts,
        ]
        return any(
            inter.GetInteractionType() == oechem.OEHBondInteractionHint(a)
            for a in accepted_types
        )


class IsLigandDonatesHBondInterInteraction(oechem.OEUnaryInteractionHintPred):
    """
    Predicate to identify hydrogen bond inter-molecular interaction where ligand donates proton from protein atom.

    LigandDonates - same as ProteinAccepts
    NonIdealLigandDonates - same as NonIdealProteinAccepts
    """

    def __call__(self, inter: oechem.OEInteractionHint) -> bool:
        """Evaluate interaction."""
        if inter.IsIntra():
            return False
        accepted_types = [
            oechem.OEHBondInteractionHintType_LigandDonates,
            oechem.OEHBondInteractionHintType_NonIdealLigandDonates,
        ]
        return any(
            inter.GetInteractionType() == oechem.OEHBondInteractionHint(a)
            for a in accepted_types
        )


class IsStackingInteraction(oechem.OEUnaryInteractionHintPred):
    """Predicate to identify Pi/T stacking inter molecular interaction."""

    def __init__(self, stacking_type: int) -> None:
        """Initialize predicate."""
        self._stacking_type = stacking_type

    def __call__(self, inter: oechem.OEInteractionHint) -> bool:
        """Evaluate interaction."""
        if inter.IsIntra():
            return False
        return inter.GetInteractionType() == oechem.OEStackingInteractionHint(
            oechem.OEStackingInteractionHintType_Pi
        )

    def CreateCopy(self):  # noqa: ANN201, N802
        """Copy constructor."""
        return IsStackingInteraction(self._stacking_type).__disown__()


INTERACTION_SYMBOLS = [
    _InteractionSymbol(oechem.OEIsClashInteractionHint(), "clash", "❌"),
    _InteractionSymbol(oechem.OEIsContactInteractionHint(), "contact", "🤝"),
    _InteractionSymbol(
        IsLigandAcceptsHBondInterInteraction(), "lig-acc", "[blue]⬅H[/blue]"
    ),
    _InteractionSymbol(
        IsLigandDonatesHBondInterInteraction(), "lig-don", "[red]H⮕[/red]"
    ),
    _InteractionSymbol(oechem.OEIsSaltBridgeInteractionHint(), "salt-bridge", "🧂"),
    _InteractionSymbol(
        IsStackingInteraction(oechem.OEStackingInteractionHintType_T),
        "t-stack",
        "[green]T⏣[/green]",
    ),
    _InteractionSymbol(
        IsStackingInteraction(oechem.OEStackingInteractionHintType_Pi),
        "pi-stack",
        "[green]π⏣[/green]",
    ),
    _InteractionSymbol(oechem.OEIsHalogenBondInteractionHint(), "halogen", "💡"),
    _InteractionSymbol(
        oechem.OEIsCationPiInteractionHint(), "cation-pi", "[blue]+π[/blue]"
    ),
    _InteractionSymbol(oechem.OEIsChelatorInteractionHint(), "chelator", "🟨"),
]


def _print_ligand_interaction_table(
    console: rich.console.Console,
    active_site: oechem.OEInteractionHintContainer,
    only_ligand_atoms: oechem.OEAtomBondSet,
) -> None:
    columns = [
        "ligand atom",
        "interactions",
        "contacts/clashes",
        "interacting with residue(s)",
    ]
    caption = "; ".join(f"{s} {n}" for _, n, s in INTERACTION_SYMBOLS)
    table = rich.table.Table(
        *columns,
        title="Ligand Protein interactions",
        caption=caption,
        caption_style="bold",
    )

    for atom in active_site.GetMolecule(
        oechem.OELigandInteractionHintComponent()
    ).GetAtoms():
        if not only_ligand_atoms.HasAtom(atom):
            continue

        residue_names: set[str] = set()
        symbols: list[str] = []

        for inter in active_site.GetInteractions(
            oechem.OEAndInteractionHint(
                oechem.OEIsInterInteractionHint(), oechem.OEHasInteractionHint(atom)
            )
        ):
            if prot_frag := inter.GetFragment(
                oechem.OEProteinInteractionHintComponent()
            ):
                for prot_atom in prot_frag.GetAtoms():
                    residue = oechem.OEAtomGetResidue(prot_atom)
                    residue_names.add(_get_residue_str(residue))

            symbol = next(
                (symbol for pred, _, symbol in INTERACTION_SYMBOLS if pred(inter)), "❓"
            )
            symbols.append(symbol)

        if len(residue_names) != 0:
            symbol_counts: dict[str, int] = dict(Counter(symbols))
            num_clashes, num_contacts = symbol_counts.pop("❌", 0), symbol_counts.pop(
                "🤝", 0
            )
            clash_contact_symbols = "" if num_clashes == 0 else num_clashes * "❌"
            clash_contact_symbols += "" if num_contacts == 0 else num_contacts * "🤝"
            other_symbols = " ".join(
                c * s if c == 1 else f"{c}x{s}" for (s, c) in symbol_counts.items()
            )
            row_data = [
                str(atom),
                other_symbols,
                clash_contact_symbols,
                "   ".join(sorted(residue_names)),
            ]

            table.add_row(*row_data)

    console.print(table)


def _get_residue_str(residue: oechem.OEResidue) -> str:
    return (
        f"{residue.GetName():3s} {residue.GetResidueNumber():4d} {residue.GetChainID()}"
    )


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