#!/usr/bin/env python3
# (C) 2023 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 the protein-ligand interactions of a list of complexes."""

import argparse
import os
import sys
import typing
from pathlib import Path

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

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict the protein-ligand interactions of a list of complexes."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme"]
__SCRIPT_CATEGORIES__ = ["visualization", "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)

    # input options
    input_group = parser.add_argument_group("Input ligand-protein complexes")
    exclusive_input_group = input_group.add_mutually_exclusive_group(required=True)
    exclusive_input_group.add_argument(
        "--complexes",
        type=str,
        nargs="+",
        required=False,
        metavar="PDB-FILE(S)",
        help="list of input PDB/CIF file of the ligand-protein complexes",
    )

    exclusive_input_group.add_argument(
        "--design-units",
        type=str,
        nargs="+",
        required=False,
        metavar="DU-FILE(S)",
        help="list of input design unit files",
    )
    input_group.add_argument(
        "--reference",
        type=str,
        required=False,
        metavar="MOL-FILE",
        help="input file of reference molecule used for ligand alignment",
    )

    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=2,
        choices=range(1, 3),
        metavar="N",
        help="number of complexes per page (default: %(default)s)",
    )
    report_group.add_argument(
        "--highlight-reference",
        action="store_true",
        help="if reference specified highlight match in ligand",
    )
    report_group.add_argument(
        "--page-by-page",
        action="store_true",
        help="write pages of report to separate numbered image files (default: %(default)s)",
    )
    return parser.parse_args()


def main() -> int:
    """Depict list  interactions of an active site."""
    args = parse_options()

    _check_report_file(args)

    console = rich.console.Console()

    report_opts = oedepict.OEReportOptions(args.rows, 1)
    report_opts.SetPageMargins(10)
    report_opts.SetCellGap(5)
    report = oedepict.OEReport(report_opts)

    reference_sub_search: oechem.OESubSearch | None = None
    if args.reference:
        ifs = oechem.oemolistream()
        if not ifs.open(args.reference):
            console.print(f"[red]Unable to open {args.reference} for reading[/red]")
            return os.EX_DATAERR
        reference_mol = oechem.OEGraphMol()
        if not oechem.OEReadMolecule(ifs, reference_mol):
            console.print(
                f"[red]Unable to read reference molecule from {args.reference}[/red]"
            )
            return os.EX_DATAERR
        if reference_mol.GetDimension() != 2:  # noqa: PLR2004
            console.print("[red]Reference molecule has to have 2D coordinates.[/red]")
            return os.EX_DATAERR
        reference_sub_search = oechem.OESubSearch(
            reference_mol,
            oechem.OEExprOpts_DefaultAtoms,
            oechem.OEExprOpts_DefaultBonds,
        )
        if not reference_sub_search.IsValid():
            console.print(
                f"[red]Invalid reference substructure search for {args.reference}[/red]"
            )
            return os.EX_DATAERR

    depict_options = oegrapheme.OE2DActiveSiteDisplayOptions(
        report.GetCellWidth(), report.GetCellHeight() * 0.8
    )
    if reference_sub_search is not None:
        depict_options.SetLigandAlignerFunctor(
            LigandSubSearchAligner(reference_sub_search)
        )
        depict_options.SetOptimizeLigandOrientation(False)

    highlight = oedepict.OEHighlightByCogwheel(oechem.OELightBlue)
    highlight.SetInnerContour(False)

    for active_site in get_active_sites(args, console):
        # depict each active site
        cell = report.NewCell()
        active_site_frame = oedepict.OEImageFrame(
            cell, cell.GetWidth(), cell.GetHeight() * 0.80, oedepict.OE2DPoint(0, 0)
        )
        legend_frame = oedepict.OEImageFrame(
            cell,
            cell.GetWidth(),
            cell.GetHeight() * 0.20,
            oedepict.OE2DPoint(0, cell.GetHeight() * 0.80),
        )

        oegrapheme.OEPrepareActiveSiteDepiction(active_site)
        active_site_disp = oegrapheme.OE2DActiveSiteDisplay(active_site, depict_options)
        if reference_sub_search is not None and args.highlight_reference:
            ligand = active_site_disp.GetDisplayedLigand()
            for match in reference_sub_search.Match(ligand):
                oegrapheme.OEAddLigandHighlighting(active_site_disp, highlight, match)
                break
        oegrapheme.OERenderActiveSite(active_site_frame, active_site_disp)

        rows, cols = 2, 5
        legend_options = oegrapheme.OE2DActiveSiteLegendDisplayOptions(rows, cols)
        oegrapheme.OEDrawActiveSiteLegend(
            legend_frame, active_site_disp, legend_options
        )

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

    return os.EX_OK


class LigandSubSearchAligner(oegrapheme.OELigandAlignerBase):
    """Align ligand based on substructure."""

    def __init__(self, sub_search: oechem.OESubSearch) -> None:
        """Initialize functor."""
        oegrapheme.OELigandAlignerBase.__init__(self)
        self._sub_search = oechem.OESubSearch(sub_search)
        # Configure options so we keep the ligand's existing coordinates,
        # only reorienting it to match the core's orientation.
        self._alignment_options = oedepict.OEAlignmentOptions()
        self._alignment_options.SetClearCoords(False)
        self._alignment_options.SetAddDepictionHydrogens(False)
        self._alignment_options.SetRotateAroundBonds(False)

    def __call__(self, ligand: oechem.OEMolBase) -> bool:
        """Align ligand based on substructure."""
        # This method must not modify the ligand except for changing its coordinates.
        # Adding or removing atoms would invalidate the ligand's active site.
        if not self._sub_search.SingleMatch(ligand):
            # ligand does not match core
            return False
        align_result = oedepict.OEPrepareAlignedDepiction(
            ligand, self._sub_search, self._alignment_options
        )
        return align_result.IsValid()

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


def get_active_sites(
    args: argparse.Namespace, console: rich.console.Console
) -> typing.Iterator[oechem.OEInteractionHintContainer]:
    """Yield valid active sites."""
    file_names: list[str]
    if args.complexes:
        file_names = args.complexes
        process_func = get_protein_and_ligand_from_pdb
    if args.design_units:
        file_names = args.design_units
        process_func = get_protein_and_ligand_from_design_unit

    for filename in file_names:
        result = process_func(filename, console)
        if result is not None:
            protein, ligand = result
            active_site = oechem.OEInteractionHintContainer(protein, ligand)
            if not active_site.IsValid():
                console.print(
                    f"[red]Cannot initialize active site for {filename}![/red]"
                )
                continue
            active_site.SetTitle(ligand.GetTitle())
            oechem.OEPerceiveInteractionHints(active_site)
            if active_site.NumInteractions() == 0:
                console.print(f"[red]No interaction detected for {filename}![/red]")
                continue
            yield active_site


def get_protein_and_ligand_from_pdb(
    filename: str, console: rich.console.Console
) -> tuple[oechem.OEMolBase, oechem.OEMolBase] | None:
    """Read protein and ligand from pdb/cif file."""
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        console.print(f"[red]Unable to open {filename} for reading[/red]")
        return None

    complex_mol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ifs, complex_mol):
        console.print(f"[red]Unable to read complex from {filename}[/red]")
        return None

    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:
        console.print("Cannot separate complex!")
        return None

    return protein, ligand


def get_protein_and_ligand_from_design_unit(
    filename: str, console: rich.console.Console
) -> tuple[oechem.OEMolBase, oechem.OEMolBase] | None:
    """Read protein and ligand from design unit file."""
    du = oechem.OEDesignUnit()
    if not oechem.OEIsReadableDesignUnit(filename) or not oechem.OEReadDesignUnit(
        filename, du
    ):
        console.print(f"Cannot read design unit {filename}.")
        return None

    protein = oechem.OEGraphMol()
    if not du.GetComponents(protein, oechem.OEDesignUnitComponents_TargetComplex):
        console.print(f"Could not extract protein from the design unit {filename}.")
        return None

    ligand = oechem.OEGraphMol()
    if not du.GetLigand(ligand):
        console.print(f"Could not extract ligand from the design unit {filename}.")
        return None

    return (protein, ligand)


def _check_report_file(args: argparse.Namespace) -> bool:
    ext = Path(args.report).suffix[1:]
    if not oedepict.OEIsRegisteredImageFile(ext):
        oechem.OEThrow.Fatal("Unknown image output 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())