#!/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 B-factor of a ligand and its environment."""

import argparse
import io
import os
import sys
from pathlib import Path

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

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict B-factor of a ligand and its environment."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme"]
__SCRIPT_CATEGORIES__ = ["visualization"]


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-FILE",
        help="input PDB file of the ligand-protein complex",
    )
    exclusive_input_group.add_argument(
        "--design-unit",
        "--du",
        type=str,
        metavar="DU-FILE",
        help="input design unit file",
    )
    # image options
    image_group = parser.add_argument_group("Image options")
    image_group.add_argument(
        "--image",
        type=str,
        required=False,
        metavar="IMAGE-FILE",
        help="output image file (PNG, SVG) (required: %(required)s) -- if no output is provided the image will be displayed on the  screen",
    )
    image_group.add_argument(
        "--width",
        type=int,
        default=600,
        help="width of output image (default: %(default)s)",
    )
    image_group.add_argument(
        "--height",
        type=int,
        default=400,
        help="height of output image (default: %(default)s)",
    )
    viz_group = parser.add_argument_group("Visualization options")
    viz_group.add_argument(
        "--max-dist",
        type=float,
        default=4.0,
        help="maximum distance of receptor atoms to be considered (default: %(default)s)",
    )
    parser.add_argument("--help-image", action=HelpPreviewAction)

    return parser.parse_args()


def main() -> int:
    """Depict B-factor."""
    args = parse_options()

    _check_image_file(args)
    is_svg: bool = args.image and Path(args.image).suffix[1:].upper() == "SVG"

    if args.complex:
        protein, ligand = get_protein_and_ligand_from_pdb(args.complex)
    elif args.design_unit:
        protein, ligand = get_protein_and_ligand_from_design_unit(args.design_unit)
    else:
        oechem.OEThrow.Fatal("Invalid input option!")

    # calculate average BFactor of the whole complex
    avg_bfactor = get_average_bfactor(protein, ligand)

    console = rich.console.Console()
    console.print(f"Average B-factor in complex = {avg_bfactor:.2f}")

    # calculate minimum and maximum BFactor of the ligand and its environment
    min_bfactor, max_bfactor = get_min_and_max_bfactor(protein, ligand, args.max_dist)
    console.print(
        f"B-factor of ligand and its environment in {args.max_dist:.1f} Å in range [{min_bfactor:.2f}-{max_bfactor:.2f}]"
    )

    # attach to each ligand atom the average BFactor of the nearby protein atoms
    tag: int = oechem.OEGetTag("avg residue BFfactor")
    set_average_bfactor_of_nearby_protein_atoms(protein, ligand, tag, args.max_dist)

    image = oedepict.OEImage(args.width, args.height)

    main_frame = oedepict.OEImageFrame(
        image, args.width, args.height * 0.85, oedepict.OE2DPoint(0.0, 0.0)
    )
    legend_frame = oedepict.OEImageFrame(
        image,
        args.width,
        args.height * 0.15,
        oedepict.OE2DPoint(0.0, args.height * 0.85),
    )

    color_gradient = get_bfactor_color_gradient()

    opts = oedepict.OE2DMolDisplayOptions(
        main_frame.GetWidth(), main_frame.GetHeight(), oedepict.OEScale_AutoScale
    )
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)
    opts.SetTitleLocation(oedepict.OETitleLocation_Hidden)

    depict_bfactor(main_frame, ligand, opts, color_gradient, tag, is_svg)
    depict_color_gradient(
        legend_frame, color_gradient, min_bfactor, max_bfactor, avg_bfactor
    )

    if args.image and Path(args.image).suffix[1:].lower() == "svg":
        icon_scale = 0.5
        oedepict.OEAddInteractiveIcon(
            image, oedepict.OEIconLocation_TopRight, icon_scale
        )
    oedepict.OEDrawCurvedBorder(image, oedepict.OELightGreyPen, 10.0)

    if args.image:
        oedepict.OEWriteImage(args.image, image)
    else:
        _img = Image.open(io.BytesIO(oedepict.OEWriteImageToBytes("png", image)))
        _img.show()

    return os.EX_OK


#############################################################################
#
#############################################################################


def depict_bfactor(
    image: oedepict.OEImageBase,
    ligand: oechem.OEMolBase,
    opts: oedepict.OE2DMolDisplayOptions,
    color_gradient: oechem.OEColorGradientBase,
    tag: int,
    is_svg: bool,
) -> None:
    """Depict B-factor."""
    # prepare ligand for depiction

    oegrapheme.OEPrepareDepictionFrom3D(ligand)

    clear_coords, suppress_hydrogens = False, False
    prep_opts = oedepict.OEPrepareDepictionOptions(clear_coords, suppress_hydrogens)
    prep_opts.SetDepictOrientation(oedepict.OEDepictOrientation_Horizontal)
    oedepict.OEPrepareDepiction(ligand, prep_opts)

    arc_fxn = BFactorArcFxn(color_gradient, tag)
    for atom in ligand.GetAtoms():
        oegrapheme.OESetSurfaceArcFxn(ligand, atom, arc_fxn)
    opts.SetScale(oegrapheme.OEGetMoleculeSurfaceScale(ligand, opts))

    # render ligand and visualize BFactor

    disp = oedepict.OE2DMolDisplay(ligand, opts)

    if is_svg:
        font = oedepict.OEFont(
            oedepict.OEFontFamily_Default,
            oedepict.OEFontStyle_Default,
            14,
            oedepict.OEAlignment_Center,
            oechem.OEBlack,
        )
        for atom_disp in disp.GetAtomDisplays():
            atom = atom_disp.GetAtom()
            if not oechem.OEHasResidue(atom):
                continue
            res = oechem.OEAtomGetResidue(atom)
            hover_text = f"bfactor={res.GetBFactor():.2f}"
            oedepict.OEDrawSVGHoverText(disp, atom_disp, hover_text, font)

    color_bfactor = ColorLigandAtomByBFactor(color_gradient)
    oegrapheme.OEAddGlyph(disp, color_bfactor, oechem.OEIsTrueAtom())

    oegrapheme.OEDraw2DSurface(disp)

    oedepict.OERenderMolecule(image, disp)


def depict_color_gradient(
    image: oedepict.OEImageBase,
    color_gradient: oechem.OEColorGradientBase,
    min_bfactor: float,
    max_bfactor: float,
    avg_bfactor: float,
) -> None:
    """Depicts color gradient."""
    opts = oegrapheme.OEColorGradientDisplayOptions()
    opts.SetColorStopPrecision(1)
    opts.AddMarkedValue(avg_bfactor)
    opts.SetBoxRange(min_bfactor, max_bfactor)

    oegrapheme.OEDrawColorGradient(image, color_gradient, opts)


def get_average_bfactor(protein: oechem.OEMolBase, ligand: oechem.OEMolBase) -> float:
    """Calculate the average b-factor for the all ligand and protein atoms."""
    num_atoms, sum_bfactor = 0, 0.0
    for mol in [protein, ligand]:
        for atom in mol.GetAtoms():
            if not oechem.OEHasResidue(atom):
                continue
            res = oechem.OEAtomGetResidue(atom)
            sum_bfactor += res.GetBFactor()
            num_atoms += 1
    return sum_bfactor / num_atoms


class NotHydrogenOrWater(oechem.OEUnaryAtomPred):
    """Predicate used to identify heady atoms but ignore OH2."""

    def __call__(self, atom: oechem.OEAtomBase) -> bool:
        """Evaluate atom."""
        if atom.GetAtomicNum() == oechem.OEElemNo_H:
            return False
        if not oechem.OEHasResidue(atom):
            return False

        water_pred = oechem.OEIsWater()
        return not water_pred(atom)


def get_min_and_max_bfactor(
    protein: oechem.OEMolBase, ligand: oechem.OEMolBase, max_distance: float
) -> tuple[float, float]:
    """Calculate the range of the b-factor."""
    min_bfactor, max_bfactor = float("inf"), float("-inf")

    # ligand atoms

    for atom in ligand.GetAtoms(oechem.OEIsHeavy()):
        if not oechem.OEHasResidue(atom):
            continue
        res = oechem.OEAtomGetResidue(atom)
        min_bfactor = min(min_bfactor, res.GetBFactor())
        max_bfactor = max(max_bfactor, res.GetBFactor())

    # protein atoms close to ligand atoms
    consider_bfactor = NotHydrogenOrWater()

    nn = oechem.OENearestNbrs(protein, max_distance)
    for lig_atom in ligand.GetAtoms(oechem.OEIsHeavy()):
        for neigh in nn.GetNbrs(lig_atom):
            prot_atom = neigh.GetBgn()

            if consider_bfactor(prot_atom):
                res = oechem.OEAtomGetResidue(prot_atom)
                min_bfactor = min(min_bfactor, res.GetBFactor())
                max_bfactor = max(max_bfactor, res.GetBFactor())

    return min_bfactor, max_bfactor


def set_average_bfactor_of_nearby_protein_atoms(
    protein: oechem.OEMolBase, ligand: oechem.OEMolBase, tag: int, max_distance: float
) -> None:
    """Set average b-factor on protein atoms close to ligand."""
    consider_bfactor = NotHydrogenOrWater()
    nn = oechem.OENearestNbrs(protein, max_distance)
    for ligand_atom in ligand.GetAtoms(oechem.OEIsHeavy()):
        sum_bfactor = 0.0
        neighs = []
        for neigh in nn.GetNbrs(ligand_atom):
            pro_atom = neigh.GetBgn()
            if consider_bfactor(pro_atom):
                res = oechem.OEAtomGetResidue(pro_atom)
                sum_bfactor += res.GetBFactor()
                neighs.append(pro_atom)

        avg_bfactor = 0.0
        if len(neighs) > 0:
            avg_bfactor = sum_bfactor / len(neighs)
        ligand_atom.SetDoubleData(tag, avg_bfactor)


def get_bfactor_color_gradient() -> oechem.OELinearColorGradient:
    """Initialise color gradient used to visualize b-factor values."""
    color_gradient = oechem.OELinearColorGradient()
    color_gradient.AddStop(oechem.OEColorStop(0.0, oechem.OEDarkBlue))
    color_gradient.AddStop(oechem.OEColorStop(10.0, oechem.OELightBlue))
    color_gradient.AddStop(oechem.OEColorStop(25.0, oechem.OEYellowTint))
    color_gradient.AddStop(oechem.OEColorStop(50.0, oechem.OERed))
    color_gradient.AddStop(oechem.OEColorStop(100.0, oechem.OEDarkRose))
    return color_gradient


class BFactorArcFxn(oegrapheme.OESurfaceArcFxnBase):
    """Surface drawer around ligand."""

    def __init__(self, color_gradient: oechem.OEColorGradientBase, tag: int) -> None:
        """Initialize."""
        oegrapheme.OESurfaceArcFxnBase.__init__(self)
        self._color_gradient = color_gradient
        self._tag = tag

    def __call__(
        self, image: oedepict.OEImageBase, arc: oegrapheme.OESurfaceArc
    ) -> bool:
        """Draw arc."""
        atom_disp = arc.GetAtomDisplay()
        if atom_disp is None or not atom_disp.IsVisible():
            return False

        atom = atom_disp.GetAtom()
        if atom is None:
            return False

        avg_residue_bfactor = atom.GetDoubleData(self._tag)
        if avg_residue_bfactor == 0.0:
            return True
        color = self._color_gradient.GetColorAt(avg_residue_bfactor)

        pen = oedepict.OEPen(color, color, oedepict.OEFill_Off, 5.0)

        center = arc.GetCenter()
        bgn_angle, end_angle = arc.GetBgnAngle(), arc.GetEndAngle()
        radius = arc.GetRadius()

        oegrapheme.OEDrawDefaultSurfaceArc(
            image, center, bgn_angle, end_angle, radius, pen
        )

        return True

    def CreateCopy(self):  # noqa: ANN201, N802
        """Copy constructor."""
        return BFactorArcFxn(self._color_gradient, self._tag).__disown__()


class ColorLigandAtomByBFactor(oegrapheme.OEAtomGlyphBase):
    """Class used to color ligand atoms based on their b-factor."""

    def __init__(self, color_gradient: oechem.OEColorGradientBase) -> None:
        """Initialize."""
        oegrapheme.OEAtomGlyphBase.__init__(self)
        self._color_gradient = color_gradient

    def RenderGlyph(  # noqa: N802
        self, disp: oedepict.OE2DMolDisplay, atom: oechem.OEAtomBase
    ) -> bool:
        """Highlight atom."""
        atom_disp = disp.GetAtomDisplay(atom)
        if atom_disp is None or not atom_disp.IsVisible():
            return False

        if not oechem.OEHasResidue(atom):
            return False

        res = oechem.OEAtomGetResidue(atom)
        bfactor = res.GetBFactor()
        color = self._color_gradient.GetColorAt(bfactor)

        pen = oedepict.OEPen(color, color, oedepict.OEFill_On, 1.0)
        radius = disp.GetScale() / 3.0

        layer = disp.GetLayer(oedepict.OELayerPosition_Below)
        circle_style = oegrapheme.OECircleStyle_Default
        oegrapheme.OEDrawCircle(layer, circle_style, atom_disp.GetCoords(), radius, pen)
        return True

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


def get_protein_and_ligand_from_pdb(
    pdb_filename: str,
) -> tuple[oechem.OEMolBase, oechem.OEMolBase]:
    """Read protein and and ligand from from pdb/cif file."""
    ifs = oechem.oemolistream()
    if not ifs.open(pdb_filename):
        oechem.OEThrow.Fatal(f"Unable to open {pdb_filename} for reading")

    complex_mol = oechem.OEGraphMol()
    if not oechem.OEReadMolecule(ifs, complex_mol):
        oechem.OEThrow.Fatal(f"Unable to read complex from {pdb_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!")

    return protein, ligand


def get_protein_and_ligand_from_design_unit(
    filename: str,
) -> tuple[oechem.OEMolBase, oechem.OEMolBase]:
    """Read protein and and ligand from from design unit file."""
    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)


def _check_image_file(args: argparse.Namespace) -> None:
    # script will terminate if there is some issues
    if not args.image:
        # image will be displayed on the screen
        return
    ext = Path(args.image).suffix[1:].upper()
    if not oedepict.OEIsRegisteredImageFile(ext):
        oechem.OEThrow.Fatal("Unknown image output type!")

    ofs = oechem.oeofstream()
    if not ofs.open(args.image):
        oechem.OEThrow.Fatal("Cannot open output image file!")


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