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

"""Serialize protein-ligand interactions to OEB/JSON."""

import argparse
import os
import pathlib
import sys

import rich.console
from openeye import oechem
from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Serialize protein-ligand interactions to OEB/JSON."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_KEYWORDS__ = [
    "perception",
    "active-site",
    "protein-ligand",
    "interactions",
    "serialization",
]
__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]",
    )
    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",
    )

    # 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 PDB/CIF file of the ligand-protein complex",
    )
    exclusive_input_group.add_argument(
        "--design-unit",
        "--du",
        type=str,
        metavar="DU-FILE",
        help="input design unit file (.oedu)",
    )
    input_group.add_argument(
        "--mol",
        type=str,
        required=False,
        metavar="MOL-FILE",
        help="input of molecules at active site (.sdf, .oeb)",
    )

    # output options
    output_group = parser.add_argument_group("Output options")
    output_group.add_argument(
        "-o",
        "--output",
        type=str,
        required=True,
        metavar="OEB/JSON-FILE",
        help="output file with perceived interactions (.oeb or .json)",
    )

    return parser.parse_args()


def main() -> int:
    """Print interactions to console."""
    args = parse_options()

    console = rich.console.Console(record=args.save_console_svg)

    if args.complex:
        prot, lig = get_prot_lig_from_pdb(args.complex)
    elif args.design_unit:
        prot, lig = get_prot_lig_from_design_unit(args.design_unit)
    else:
        oechem.OEThrow.Fatal("Invalid input option!")

    if args.mol:
        ifs = oechem.oemolistream()
        if not ifs.open(args.mol):
            oechem.OEThrow.Fatal(f"Unable to open {args.mol} for reading!")
        mols: list[oechem.OEMolBase] = [
            oechem.OEGraphMol(m) for m in ifs.GetOEGraphMols()
        ]
        if len(mols) == 0:
            oechem.OEThrow.Fatal(f"No molecules read from {args.mol}!")
        console.print(
            f"{len(mols)} molecule(s) read from '{pathlib.Path(args.mol).name}' file!"
        )

    ofs = oechem.oemolostream()
    if not ofs.open(args.output):
        oechem.OEThrow.Fatal(f"Unable to open {args.output} for writing")
    if ofs.GetFormat() not in [oechem.OEFormat_OEB, oechem.OEFormat_JSON]:
        oechem.OEThrow.Fatal("Output file must have .oeb or .json extension")

    if args.mol:
        serialize_series_of_interactions(ofs, prot, mols, console)
    else:
        serialize_interactions(ofs, prot, lig, console)

    console.print(
        f"Protein and ligand(s) with interactions serialized to '{pathlib.Path(args.output).name}' file"
    )

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

    return os.EX_OK


def serialize_interactions(
    ofs: oechem.oemolostream,
    protein: oechem.OEMolBase,
    ligand: oechem.OEMolBase,
    console: rich.console.Console,
) -> bool:
    """Serialize interactions to output file (OEB, JSON)."""
    active_site = oechem.OEInteractionHintContainer(protein, ligand)

    if not oechem.OEIsValidActiveSite(active_site):
        console.print("[red]Invalid active site![/red]")
        return False

    oechem.OEPerceiveInteractionHints(active_site)
    if active_site.NumInteractions() == 0:
        console.print("[red]No interactions perceived![/red]")
        return False
    console.print(f"Number of interactions perceived: {active_site.NumInteractions()}")

    protein = oechem.OEGraphMol()
    ligand = oechem.OEGraphMol()
    if not oechem.OESerializeInteractionsHintContainer(protein, ligand, active_site):
        console.print("[red]Failed to serialize interactions![/red]")
        return False
    oechem.OEWriteMolecule(ofs, protein)
    oechem.OEWriteMolecule(ofs, ligand)
    return True


def serialize_series_of_interactions(
    ofs: oechem.oemolostream,
    protein: oechem.OEMolBase,
    molecules: list[oechem.OEMolBase],
    console: rich.console.Console,
) -> bool:
    """Serialize series of interactions to output file (OEB, JSON)."""
    mol_with_interactions: list[oechem.OEMolBase] = []

    # set serialization IDs for the protein to ensure that the same protein is recognized across multiple active sites
    oechem.OESetInteractionsHintSerializationIds(protein)

    with Progress(
        TextColumn("{task.description}"),
        BarColumn(bar_width=60),
        TextColumn("{task.percentage:3.1f}%"),
        TimeElapsedColumn(),
        transient=False,
        disable=len(molecules) < 10,  # noqa: PLR2004
        console=console,
    ) as progress:
        serialization_task = progress.add_task(
            "[blue]Serializing interactions", total=len(molecules)
        )
        for mol in molecules:
            progress.update(serialization_task, advance=1)

            active_site = oechem.OEInteractionHintContainer(protein, mol)
            if not oechem.OEIsValidActiveSite(active_site):
                console.print("[red]Invalid active site![/red]")
                continue

            oechem.OEPerceiveInteractionHints(active_site)
            if active_site.NumInteractions() == 0:
                console.print("[red]No interactions perceived![/red]")
                continue

            ligand = oechem.OEGraphMol()
            if not oechem.OESerializeInteractionsHintContainer(
                protein, ligand, active_site
            ):
                console.print("[red]Failed to serialize interactions![/red]")
                continue
            mol_with_interactions.append(oechem.OEGraphMol(ligand))

    mol_with_interactions.insert(
        0, protein
    )  # add the protein as the first molecule in the output file
    for mol in mol_with_interactions:
        oechem.OEWriteMolecule(ofs, mol)
    return True


def get_prot_lig_from_pdb(filename: str) -> tuple[oechem.OEMolBase, oechem.OEMolBase]:
    """Initialize active site interaction container from pdb/cif file."""
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        oechem.OEThrow.Fatal("Unable to open {filename} for reading")

    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)
    return protein, ligand


def get_prot_lig_from_design_unit(
    filename: str,
) -> tuple[oechem.OEMolBase, oechem.OEMolBase]:
    """Initialize active site from design unit."""
    du = oechem.OEDesignUnit()
    if not oechem.OEIsReadableDesignUnit(filename) or not oechem.OEReadDesignUnit(
        filename, du
    ):
        oechem.OEThrow.Fatal("Cannot read design unit.")

    protein = oechem.OEGraphMol()
    if not du.GetComponents(protein, oechem.OEDesignUnitComponents_TargetComplex):
        oechem.OEThrow.Fatal("Could not extract protein from the design unit.")

    ligand = oechem.OEGraphMol()
    if not du.GetLigand(ligand):
        oechem.OEThrow.Fatal("Could not extract ligand from the design unit.")

    return protein, ligand


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