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

"""Visualizes the dihedral angles of multi-conformer molecule."""

import argparse
import math
import os
import sys
import uuid
from pathlib import Path

from openeye import oechem, oedepict, oegrapheme
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict dihedral angles of multi-conformer molecule."
__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]",
    )
    parser.add_argument("--help-image", action=HelpPreviewAction)

    # input options
    input_group = parser.add_argument_group("Input options")
    input_group.add_argument(
        "--mol",
        type=str,
        required=True,
        metavar="MOL-FILE",
        help="input multi-conformer molecule file (oeb, sdf)",
    )
    input_group.add_argument(
        "--ref-mol",
        type=str,
        required=False,
        metavar="MOL-FILE",
        help="input reference molecule (required: %(required)s)",
    )
    image_group = parser.add_argument_group("Image options")
    image_group.add_argument(
        "--image",
        type=str,
        required=True,
        metavar="IMAGE-FILE",
        help="output image file (SVG)",
    )
    image_group.add_argument(
        "--width",
        type=int,
        default=800,
        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)",
    )
    visual_group = parser.add_argument_group("visualization options")
    visual_group.add_argument(
        "--num-bins",
        type=int,
        default=24,
        metavar="N",
        choices=range(12, 49),
        help="number of bins in the dihedral angle histogram (default: %(default)s)",
    )
    visual_group.add_argument(
        "--flexibility",
        default=False,
        action="store_true",
        help="visualize dihedral angle flexibility",
    )
    return parser.parse_args()


def main() -> int:
    """Visualizes the dihedral angles of multi-conformer molecule."""
    args = parse_options()

    _check_image_file(args)

    compere_title = True
    conf_test = oechem.OEIsomericConfTest(not compere_title)

    mol = _get_molecule(args.mol, conf_test)
    if mol.NumConfs() == 1:
        oechem.OEThrow.Fatal("Multi conformations are required!")

    ref_mol: oechem.OEMCMolBase | None = None
    if args.ref_mol:
        ref_mol = _get_molecule(args.ref_mol, conf_test)
        for m in [mol, ref_mol]:
            oechem.OESuppressHydrogens(m)
            oechem.OECanonicalOrderAtoms(m)
            oechem.OECanonicalOrderBonds(m)
            m.Sweep()

        if not conf_test.CompareMols(mol, ref_mol):
            oechem.OEThrow.Warning("Unmatched reference molecule is ignored!")
            ref_mol = None

    tag = oechem.OEGetTag("dihedral_histogram")

    find_dihedrals(mol, tag)
    set_dihedral_histograms(mol, tag, args.num_bins)

    if ref_mol:
        find_dihedrals(ref_mol, tag)
        set_dihedral(ref_mol, tag)

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

    main_offset = oedepict.OE2DPoint(0, 0)
    main_frame = oedepict.OEImageFrame(image, width * 0.70, height, main_offset)
    dihedral_offset = oedepict.OE2DPoint(main_frame.GetWidth(), height * 0.30)
    dihedral_frame = oedepict.OEImageFrame(
        image, width * 0.30, height * 0.5, dihedral_offset
    )

    color_gradient = get_color_gradient(args.num_bins, args.flexibility)

    opts = oedepict.OE2DMolDisplayOptions(
        main_frame.GetWidth(), main_frame.GetHeight(), oedepict.OEScale_AutoScale
    )

    depict_dihedrals(
        main_frame,
        dihedral_frame,
        mol,
        ref_mol,
        opts,
        tag,
        args.num_bins,
        color_gradient,
    )

    if args.flexibility:
        legend_opts = _get_legend_display_options()
        legend = oedepict.OELegendLayout(image, "Legend", legend_opts)

        legend_area = legend.GetLegendArea()
        draw_color_gradient(legend_area, color_gradient)

        oedepict.OEDrawLegendLayout(legend)

    icon_scale = 0.5
    oedepict.OEAddInteractiveIcon(image, oedepict.OEIconLocation_TopRight, icon_scale)
    oedepict.OEDrawCurvedBorder(image, oedepict.OELightGreyPen, 10.0)

    oedepict.OEWriteImage(args.image, image)

    return os.EX_OK


def _get_molecule(
    filename: str, conf_test: oechem.OEConfTestBase
) -> oechem.OEMCMolBase:
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        oechem.OEThrow.Fatal(f"Cannot open {filename} input file!")

    if ifs.GetFormat() != oechem.OEFormat_OEB:
        ifs.SetConfTest(conf_test)

    mol = oechem.OEMol()
    if not oechem.OEReadMolecule(ifs, mol):
        oechem.OEThrow.Fatal(f"Cannot read molecule from {filename} input file!")

    if mol.GetDimension() != 3:  # noqa: PLR2004
        oechem.OEThrow.Fatal("3D coordinates are requires!")

    return mol


def _check_image_file(args: argparse.Namespace) -> None:
    # script will terminate if there is some issues
    if Path(args.image).suffix[1:].upper() != "SVG":
        oechem.OEThrow.Fatal(
            "This script only accepts SVG as the output image file format!"
        )

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


def _get_legend_display_options() -> oedepict.OELegendLayoutOptions:
    legend_opts = oedepict.OELegendLayoutOptions(
        oedepict.OELegendLayoutStyle_HorizontalTopLeft,
        oedepict.OELegendColorStyle_LightBlue,
        oedepict.OELegendInteractiveStyle_Hover,
    )
    legend_opts.SetButtonWidthScale(1.2)
    legend_opts.SetButtonHeightScale(1.2)
    legend_opts.SetMargin(oedepict.OEMargin_Right, 40.0)
    legend_opts.SetMargin(oedepict.OEMargin_Bottom, 80.0)
    return legend_opts


class IsRotatableOrMacroCycleBond(oechem.OEUnaryBondPred):
    """Predicate to identify rotatable bonds and single bonds in macro-cycles."""

    def __call__(self, bond: oechem.OEBondBase) -> bool:
        """Evaluate bond."""
        if bond.GetOrder() != 1:
            return False
        if bond.IsAromatic():
            return False

        rotor_pred = oechem.OEIsRotor()
        if rotor_pred(bond):
            return True

        return oechem.OEBondGetSmallestRingSize(bond) >= 10  # noqa: PLR2004


def find_dihedrals(mol: oechem.OEMCMolBase, tag: int) -> int:
    """
    Identify dihedral atoms.

    Iterate over rotatable bonds and identifies their dihedral
    atoms. These atoms are added to the molecule in a group
    using the given tag.
    """
    num_dihedrals = 0
    for bond in mol.GetBonds(IsRotatableOrMacroCycleBond()):
        atom_bgn = bond.GetBgn()
        atom_end = bond.GetEnd()

        neigh_bgn: oechem.OEAtomBase | None = None
        neigh_end: oechem.OEAtomBase | None = None

        for atom in atom_bgn.GetAtoms(oechem.OEIsHeavy()):
            if atom != atom_end:
                neigh_bgn = atom
                break
        for atom in atom_end.GetAtoms(oechem.OEIsHeavy()):
            if atom != atom_bgn:
                neigh_end = atom
                break

        if neigh_bgn is None or neigh_end is None:
            continue

        atom_order = [neigh_bgn, atom_bgn, atom_end, neigh_end]
        bond_order = [
            mol.GetBond(neigh_bgn, atom_bgn),
            bond,
            mol.GetBond(neigh_end, atom_end),
        ]

        if neigh_bgn.GetIdx() < neigh_end.GetIdx():
            atom_order.reverse()
            bond_order.reverse()

        atoms = oechem.OEAtomVector(atom_order)
        bonds = oechem.OEBondVector(bond_order)

        num_dihedrals += 1
        mol.NewGroup(tag, atoms, bonds)  # type: ignore[attr-defined]

    return num_dihedrals


def set_dihedral_histograms(mol: oechem.OEMCMolBase, tag: int, num_bins: int) -> None:
    """
    Set dihedral angle histogram data.

    Iterates over the dihedral groups and bins the torsional
    angles for each conformation. The histogram data is then
    attached to the groups with the given tag.
    """
    angle_inc = 360.0 / float(num_bins)

    for group in mol.GetGroups(oechem.OEHasGroupType(tag)):
        atoms = oechem.OEAtomVector()
        for atom in group.GetAtoms():
            atoms.append(atom)
        histogram = [0] * num_bins

        for conf in mol.GetConfs():
            rad = oechem.OEGetTorsion(conf, atoms[0], atoms[1], atoms[2], atoms[3])
            deg = math.degrees(rad)
            deg = (deg + 360.0) % 360.0
            bin_idx: int = math.floor(deg / angle_inc)
            histogram[bin_idx] += 1

        group.SetData(tag, histogram)


def set_dihedral(mol: oechem.OEMCMolBase, tag: int) -> None:
    """
    Set dihedral angle data.

    Iterates over the dihedral groups of the molecule, calculates the
    dihedral angle amd attaches it to the to the group.
    """
    for group in mol.GetGroups(oechem.OEHasGroupType(tag)):
        atoms = oechem.OEAtomVector()
        for atom in group.GetAtoms():
            atoms.append(atom)

        rad = oechem.OEGetTorsion(mol, atoms[0], atoms[1], atoms[2], atoms[3])
        deg = math.degrees(rad)
        deg = (deg + 360.0) % 360.0
        group.SetData(tag, deg)


def depict_dihedrals(  # noqa: PLR0913, PLR0915, PLR0917
    image: oedepict.OEImageBase,
    dihedral_image: oedepict.OEImageBase,
    mol: oechem.OEMCMolBase,
    ref_mol: oechem.OEMCMolBase | None,
    opts: oedepict.OE2DMolDisplayOptions,
    tag: int,
    num_bins: int,
    color_gradient: oechem.OEColorGradientBase,
) -> None:
    """
    Depict dihedrals.

    Highlights the dihedral atoms of a torsion and the depicts the
    corresponding dihedral angle histogram when hovering over
    the center of the torsion on the molecule display.
    """
    num_confs = mol.NumConfs()
    center = oedepict.OEGetCenter(dihedral_image)
    radius = min(dihedral_image.GetWidth(), dihedral_image.GetHeight()) * 0.40

    draw_dihedral_circle(dihedral_image, center, radius, num_bins, num_confs)

    suppress_hydrogens = True
    oegrapheme.OEPrepareDepictionFrom3D(mol, suppress_hydrogens)
    if ref_mol:
        oegrapheme.OEPrepareDepictionFrom3D(ref_mol, suppress_hydrogens)

    disp = oedepict.OE2DMolDisplay(mol, opts)

    mol_dihedral_groups: list[oechem.OEGroupBase] = []
    ref_dihedrals: list[oechem.OEGroupBase] | None = [] if ref_mol else None
    dihedral_centers: list[oedepict.OE2DPoint] = []
    group_ones = []
    svg_dihedral_groups: list[oedepict.OESVGClass] = []

    num_dihedrals = 0
    for group in mol.GetGroups(oechem.OEHasGroupType(tag)):
        uniqueid = uuid.uuid4().hex
        group_one = image.NewSVGGroup("torsion_area_" + uniqueid)
        svg_group = image.NewSVGGroup("torsion_data_" + uniqueid)
        oedepict.OEAddSVGHover(group_one, svg_group)

        mol_dihedral_groups.append(group)
        if (
            ref_mol
            and ref_dihedrals is not None
            and (ref_group := get_reference_dihedral(group, ref_mol, tag))
        ):
            ref_dihedrals.append(ref_group)

        dihedral_centers.append(get_dihedral_center(disp, group))
        group_ones.append(group_one)
        svg_dihedral_groups.append(svg_group)
        num_dihedrals += 1

    for d_idx in range(num_dihedrals):
        image.PushGroup(svg_dihedral_groups[d_idx])

        dihedral = mol_dihedral_groups[d_idx]
        ab_set = oechem.OEAtomBondSet(dihedral.GetAtoms(), dihedral.GetBonds())
        draw_highlight(image, disp, ab_set)
        dihedral_histogram = dihedral.GetData(tag)
        draw_dihedral_histogram(
            dihedral_image, dihedral_histogram, center, radius, num_bins, num_confs
        )

        if ref_mol and ref_dihedrals:
            draw_reference_dihedral(
                dihedral_image, ref_dihedrals[d_idx], tag, center, radius
            )

        image.PopGroup(svg_dihedral_groups[d_idx])

    clear_background = True
    oedepict.OERenderMolecule(image, disp, not clear_background)

    mark_pen = oedepict.OEPen(oechem.OEBlack, oechem.OEWhite, oedepict.OEFill_On, 1.0)
    far_pen = oedepict.OEPen(oechem.OEBlack, oechem.OERed, oedepict.OEFill_Off, 2.0)

    angle_inc = 360.0 / float(num_bins)

    for d_idx in range(num_dihedrals):
        image.PushGroup(group_ones[d_idx])

        dihedral = mol_dihedral_groups[d_idx]
        dihedral_histogram = dihedral.GetData(tag)
        flexibility = determine_flexibility(dihedral_histogram)
        color = color_gradient.GetColorAt(flexibility)
        mark_pen.SetBackColor(color)

        mark_radius = disp.GetScale() / 8.0
        image.DrawCircle(dihedral_centers[d_idx], mark_radius, mark_pen)

        if (
            ref_mol
            and ref_dihedrals
            and (
                get_closest_dihedral_angle(mol, dihedral, ref_dihedrals[d_idx], tag)
                > angle_inc
            )
        ):
            image.DrawCircle(dihedral_centers[d_idx], mark_radius, far_pen)

        radius = disp.GetScale() / 4.0
        image.DrawCircle(dihedral_centers[d_idx], radius, oedepict.OESVGAreaPen)

        image.PopGroup(group_ones[d_idx])


def get_closest_dihedral_angle(
    mol: oechem.OEMCMolBase,
    dihedral: oechem.OEGroupBase,
    ref_dihedral: oechem.OEGroupBase,
    tag: int,
) -> float:
    """Return the closest torsion angle difference to the reference."""
    closest_angle = float("inf")

    for conf in mol.GetConfs():
        atoms = list(dihedral.GetAtoms())
        rad = oechem.OEGetTorsion(conf, atoms[0], atoms[1], atoms[2], atoms[3])
        deg = math.degrees(rad)
        angle_diff = (abs(deg - ref_dihedral.GetData(tag)) + 360) % 360
        closest_angle = min(closest_angle, angle_diff)

    return closest_angle


def get_dihedral_center(
    disp: oedepict.OE2DMolDisplay, dihedral_group: oechem.OEGroupBase
) -> oedepict.OE2DPoint:
    """
    Return the center of a dihedral angle in the 2D image.

    The dihedral angle (stored in a group) on the
    molecule display.
    """
    center_bgn: oechem.OEAtomBase | None = None
    center_end: oechem.OEAtomBase | None = None
    for bond in dihedral_group.GetBonds():
        atom_bgn = bond.GetBgn()
        atom_end = bond.GetEnd()

        num_neighs_bgn = 0
        for neigh in atom_bgn.GetAtoms():
            if dihedral_group.HasAtom(neigh):
                num_neighs_bgn += 1
        num_neighs_end = 0
        for neigh in atom_end.GetAtoms():
            if dihedral_group.HasAtom(neigh):
                num_neighs_end += 1

        if num_neighs_bgn == 2 and num_neighs_end == 2:  # noqa: PLR2004
            center_bgn = atom_bgn
            center_end = atom_end
            break

    if not center_bgn or not center_end:
        msg = "Can not determine dihedral angle center"
        raise ValueError(msg)

    atom_disp_bgn = disp.GetAtomDisplay(center_bgn)
    atom_disp_end = disp.GetAtomDisplay(center_end)
    return (atom_disp_bgn.GetCoords() + atom_disp_end.GetCoords()) / 2.0


def draw_dihedral_circle(
    image: oedepict.OEImageBase,
    center: oedepict.OE2DPoint,
    radius: float,
    num_bins: int,
    num_confs: int,
) -> None:
    """Draw the base radial histogram."""
    grey = oechem.OEColor(210, 210, 210)
    pen = oedepict.OEPen(grey, grey, oedepict.OEFill_On, 1.0)
    image.DrawCircle(center, radius, pen)

    line_grey = oechem.OEColor(220, 220, 220)
    line_pen = oedepict.OEPen(line_grey, line_grey, oedepict.OEFill_On, 1.0)

    angle_inc = 360.0 / float(num_bins)

    v = oedepict.OE2DPoint(0.0, -1.0)
    for i in range(num_bins):
        end = oedepict.OELengthenVector(
            oedepict.OERotateVector(v, i * angle_inc), radius
        )
        image.DrawLine(center, center + end, line_pen)

    font_size: int = math.floor(radius * 0.1)
    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Bold,
        font_size,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )

    for i in range(4):
        angle = i * 90.0
        end = oedepict.OELengthenVector(
            oedepict.OERotateVector(v, angle), radius * 1.20
        )
        text = f"{angle:.1f}"
        dim = radius / 2.5
        text_frame = oedepict.OEImageFrame(
            image, dim, dim, center + end - oedepict.OE2DPoint(dim / 2.0, dim / 2.0)
        )
        oedepict.OEDrawTextToCenter(text_frame, text, font)

    min_radius = radius / 3.0
    white_pen = oedepict.OEPen(
        oechem.OEWhite,
        oechem.OEWhite,
        oedepict.OEFill_On,
        1.0,
        oedepict.OEStipple_NoLine,
    )
    image.DrawCircle(center, min_radius, white_pen)

    font.SetSize(int(font_size * 1.5))
    top = oedepict.OE2DPoint(image.GetWidth() / 2.0, -10.0)
    image.DrawText(top, "torsion histogram", font)

    bottom = oedepict.OE2DPoint(image.GetWidth() / 2.0, image.GetHeight() + 26.0)
    image.DrawText(bottom, f"number of conformations: {num_confs}", font)


def _get_text_angle(angle: float) -> float:
    if angle <= 180.0:  # noqa: PLR2004
        return (360 - angle + 90.0) % 360
    return (180 - angle + 90.0) % 360


def draw_dihedral_histogram(
    image: oedepict.OEImageBase,
    histogram: list[int],
    center: oedepict.OE2DPoint,
    radius: float,
    num_bins: int,
    num_confs: int,
) -> None:
    """Draw the radial histogram of a torsional angle."""
    min_radius = radius / 3.0
    max_value = max(histogram)
    radius_inc = (radius - min_radius) / max_value

    angle_inc = 360.0 / float(num_bins)

    line_grey = oechem.OEColor(220, 220, 220)
    value_pen = oedepict.OEPen(oechem.OERoyalBlue, line_grey, oedepict.OEFill_On, 0.5)

    max_value = 0
    max_value_idx = 0
    for i in range(len(histogram)):
        value = histogram[i]
        if value == 0:
            continue

        if value > max_value:
            max_value = value
            max_value_idx = i

        arc_radius = value * radius_inc + min_radius
        if arc_radius < 1.0:
            continue

        bgn_angle = i * angle_inc
        end_angle = (i + 1) * angle_inc

        image.DrawPie(center, bgn_angle, end_angle, arc_radius, value_pen)

    percent = max_value / (num_confs / 100.0)

    white_pen = oedepict.OEPen(
        oechem.OEWhite,
        oechem.OEWhite,
        oedepict.OEFill_On,
        1.0,
        oedepict.OEStipple_NoLine,
    )
    image.DrawCircle(center, min_radius, white_pen)

    font_size: int = math.floor(radius * 0.1)
    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Bold,
        font_size,
        oedepict.OEAlignment_Center,
        oechem.OEWhite,
    )
    angle = max_value_idx * angle_inc
    if angle >= 180.0:  # noqa: PLR2004
        angle += angle_inc * 0.3
    else:
        angle += angle_inc * 0.7
    text_angle = _get_text_angle(angle)
    v = oedepict.OE2DPoint(0.0, -1.0)
    pos = oedepict.OELengthenVector(oedepict.OERotateVector(v, angle), radius * 0.80)
    font.SetRotationAngle(text_angle)
    image.DrawText(center + pos, f"{percent:.1f}%", font)


def are_same_groups(
    group_one: oechem.OEGroupBase, group_two: oechem.OEGroupBase
) -> bool:
    """Determine whether the two groups identical."""
    for a, b in zip(group_one.GetAtoms(), group_two.GetAtoms(), strict=False):
        if a.GetIdx() != b.GetIdx():
            return False
    return all(
        a.GetIdx() == b.GetIdx()
        for a, b in zip(group_one.GetBonds(), group_two.GetBonds(), strict=False)
    )


def get_reference_dihedral(
    group: oechem.OEGroupBase, ref_mol: oechem.OEMCMolBase | None, tag: int
) -> oechem.OEGroupBase | None:
    """
    Return the torsion group on the reference molecule.

    The torsion group on the reference molecule that
    corresponds to the torsional group of the multi-conformer
    molecule.
    """
    if not ref_mol:
        return None

    for ref_group in ref_mol.GetGroups(oechem.OEHasGroupType(tag)):
        if are_same_groups(group, ref_group):
            return ref_group
    return None


def draw_reference_dihedral(
    image: oedepict.OEImageBase,
    group: oechem.OEGroupBase,
    tag: int,
    center: oedepict.OE2DPoint,
    radius: float,
) -> None:
    """Draw dihedral angle of the reference molecule."""
    if not group.HasData(tag):
        return
    angle = group.GetData(tag)
    v = oedepict.OE2DPoint(0.0, -1.0)
    bgn = oedepict.OELengthenVector(oedepict.OERotateVector(v, angle), radius / 6.0)
    end = oedepict.OELengthenVector(oedepict.OERotateVector(v, angle), radius / 3.0)
    red_pen = oedepict.OEPen(oechem.OERed, oechem.OERed, oedepict.OEFill_Off, 2.0)
    image.DrawLine(center + bgn, center + end, red_pen)

    font_size: int = math.floor(radius * 0.12)
    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Bold,
        font_size,
        oedepict.OEAlignment_Center,
        oechem.OERed,
    )

    dim = radius / 2.5
    text_frame = oedepict.OEImageFrame(
        image, dim, dim, center - oedepict.OE2DPoint(dim / 2.0, dim / 2.0)
    )
    oedepict.OEDrawTextToCenter(text_frame, f"{angle:.1f}", font)


def draw_highlight(
    image: oedepict.OEImageBase,
    disp: oedepict.OE2DMolDisplay,
    ab_set: oechem.OEAtomBondSet,
) -> None:
    """Highlight the atoms of the dihedral angle on the molecule display."""
    linewidth = disp.GetScale() / 2.0
    pen = oedepict.OEPen(
        oechem.OEBlueTint, oechem.OEBlueTint, oedepict.OEFill_On, linewidth
    )
    for bond in ab_set.GetBonds():
        atom_disp_bgn = disp.GetAtomDisplay(bond.GetBgn())
        atom_disp_end = disp.GetAtomDisplay(bond.GetEnd())
        image.DrawLine(atom_disp_bgn.GetCoords(), atom_disp_end.GetCoords(), pen)


def get_color_gradient(num_bins: int, flexibility: bool) -> oechem.OEColorGradientBase:
    """
    Initialize the color gradient.

    It is used to color the circle in the middle of the rotatable bond.
    """
    color_gradient = oechem.OEExponentColorGradient(0.25)

    if flexibility:
        color_gradient.AddStop(oechem.OEColorStop(1.0, oechem.OEBlack))
        color_gradient.AddStop(oechem.OEColorStop(num_bins, oechem.OERed))
    else:
        color_gradient.AddStop(oechem.OEColorStop(1.0, oechem.OEBlack))
        color_gradient.AddStop(oechem.OEColorStop(num_bins, oechem.OEBlack))

    return color_gradient


def draw_color_gradient(
    image: oedepict.OEImageBase, color_gradient: oechem.OEColorGradientBase
) -> None:
    """Draw the color gradient."""
    width, height = image.GetWidth(), image.GetHeight()
    frame = oedepict.OEImageFrame(
        image, width * 0.8, height * 0.8, oedepict.OE2DPoint(width * 0.1, height * 0.1)
    )

    opts = oegrapheme.OEColorGradientDisplayOptions()
    opts.SetColorStopPrecision(1)
    opts.SetColorStopLabelFontScale(0.5)
    opts.SetColorStopVisibility(False)

    opts.AddLabel(
        oegrapheme.OEColorGradientLabel(color_gradient.GetMinValue(), "rigid")
    )
    opts.AddLabel(
        oegrapheme.OEColorGradientLabel(color_gradient.GetMaxValue(), "flexible")
    )

    oegrapheme.OEDrawColorGradient(frame, color_gradient, opts)


def determine_flexibility(histogram: list[int]) -> float:
    """
    Determine molecule flexibility.

    Return an estimation of torsion flexibility based on the Shannon entropy
    of the dihedral angle histogram distribution. A rigid torsion concentrates
    counts in few bins (low entropy), while a flexible one spreads counts
    across many bins (high entropy). The result is scaled to [1, num_bins].
    """
    total = sum(histogram)
    if total == 0:
        return 1.0
    num_bins = len(histogram)
    entropy = 0.0
    for count in histogram:
        if count > 0:
            p = count / total
            entropy -= p * math.log2(p)
    max_entropy = math.log2(num_bins)
    if max_entropy == 0:
        return 1.0
    normalized = entropy / max_entropy
    return 1.0 + normalized * (num_bins - 1)


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