#!/usr/bin/env python
# (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.


"""Visualizes the variability in torsional angles sampled across the dataset."""

import argparse
import os
import pathlib
import statistics
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Visualizes the variability in torsional angles sampled."
__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_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)",
    )
    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(
        "--page-by-page",
        action="store_true",
        help="write pages of report to separate numbered image files",
    )
    parser.add_argument("--help-image", action=HelpPreviewAction)
    return parser.parse_args()


def main() -> int:
    """Visualizes variability in torsions sampled in dataset."""
    args = parse_options()

    _check_report_file(args)

    mols: list[oechem.OEMCMolBase] = _read_molecules(args.mol)
    console = rich.console.Console()
    console.print(f"Imported {len(mols)} molecules from {args.mol}")

    bond_torsions: dict[oechem.OEBondBase, list[float]] = get_unique_torsions(mols)

    normalized_torsion_variances: dict[oechem.OEBondBase, float] = (
        calculate_torsion_variances(bond_torsions)
    )

    # transfer normalized variance score to bonds for depiction
    tag = oechem.OEGetTag("torsion variability")
    for bond, score in normalized_torsion_variances.items():
        bond.SetData(tag, score)

    report_opts = oedepict.OEReportOptions()
    report_opts.SetHeaderHeight(35)
    report_opts.SetFooterHeight(45)
    report_opts.SetPageMargins(10)
    report_opts.SetCellGap(5)
    report = oedepict.OEReport(report_opts)

    depict_molecules(report, mols, tag)

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

    return os.EX_OK


def get_unique_torsions(
    mols: list[oechem.OEMCMolBase],
) -> dict[oechem.OEBondBase, list[float]]:
    """
    Retrieve unique torsions for molecules.

    Identify rotatable bonds/torsion associated with each molecule and return them along with
    list of the corresponding dihedral angels (in radian) in each conformations in a dictionary.
    """
    bond_torsions: dict[oechem.OEBondBase, list[float]] = {}
    for mol in mols:
        unique_bonds: set[oechem.OEBondBase] = set()
        for torsion in oechem.OEGetTorsions(mol, oechem.OEIsRotor()):
            if any(
                x.IsHydrogen() for x in [torsion.a, torsion.b, torsion.c, torsion.d]
            ):
                continue
            if (bond := torsion.b.GetBond(torsion.c)) and bond not in unique_bonds:
                unique_bonds.add(bond)
                bond_torsion_angels: list[float] = [
                    oechem.OEGetAbsTorsion(
                        conf, torsion.a, torsion.b, torsion.c, torsion.d
                    )
                    for conf in mol.GetConfs()
                ]
                bond_torsions[bond] = bond_torsion_angels

    return bond_torsions


def calculate_torsion_variances(
    bond_torsions: dict[oechem.OEBondBase, list[float]],
) -> dict[oechem.OEBondBase, float]:
    """Calculate sample variance of of torsion angels and normalize them to range [0.0, 1.0]."""
    torsion_variances: dict[oechem.OEBondBase, float] = {}

    for bond, torsions in bond_torsions.items():
        torsion_variances[bond] = statistics.variance(torsions)

    min_variance = min(i for i in torsion_variances.values())
    max_variance = max(i for i in torsion_variances.values())

    normalized_variances: dict[oechem.OEBondBase, float] = {}
    for bond, value in torsion_variances.items():
        normalized_variances[bond] = (value - min_variance) / (
            max_variance - min_variance
        )

    return normalized_variances


def depict_molecules(
    report: oedepict.OEReport,
    mols: list[oechem.OEMCMolBase],
    tag: int,
) -> None:
    """Depict molecules."""
    color_gradient = oechem.OELinearColorGradient()
    color_gradient.AddStop(oechem.OEColorStop(0.0, oechem.OEWhite))
    color_gradient.AddStop(oechem.OEColorStop(1.0, oechem.OERedOrange))

    bond_glyph = ColorBondByScore(color_gradient, tag)

    disp_options = oedepict.OE2DMolDisplayOptions(
        report.GetCellWidth(), report.GetCellHeight(), oedepict.OEScale_AutoScale
    )
    disp_options.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)

    font = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Bold,
        12,
        oedepict.OEAlignment_Left,
        oechem.OELightGrey,
    )
    for mol in mols:
        oegrapheme.OEPrepareDepictionFrom3D(mol)
        disp = oedepict.OE2DMolDisplay(mol, disp_options)
        oegrapheme.OEAddGlyph(disp, bond_glyph, oechem.IsTrueBond())

        cell = report.NewCell()

        layer: oedepict.OEImageBase = disp.GetLayer(oedepict.OELayerPosition_Below)
        layer.DrawText(
            oedepict.OE2DPoint(10.0, cell.GetHeight() - 10),
            f"confs={mol.NumConfs()}",
            font,
        )

        oedepict.OERenderMolecule(cell, disp)
        oedepict.OEDrawCurvedBorder(cell, oedepict.OEBlackPen, 20)

    header_text = "Normalized Variability in Torsions Sampled"
    font.SetAlignment(oedepict.OEAlignment_Center)
    font.SetColor(oechem.OEBlack)
    for header in report.GetHeaders():
        oedepict.OEDrawTextToCenter(header, header_text, font)

    color_options = oegrapheme.OEColorGradientDisplayOptions()
    color_options.SetColorStopVisibility(False)
    color_options.AddLabel(oegrapheme.OEColorGradientLabel(0.0, "low variance"))
    color_options.AddLabel(oegrapheme.OEColorGradientLabel(1.0, "high variance"))
    for footer in report.GetFooters():
        oegrapheme.OEDrawColorGradient(footer, color_gradient, color_options)


class ColorBondByScore(oegrapheme.OEBondGlyphBase):
    """Bond coloring class."""

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

    def RenderGlyph(  # noqa: N802
        self, disp: oedepict.OE2DMolDisplay, bond: oechem.OEBondBase
    ) -> bool:
        """Highlight bond."""
        bond_disp = disp.GetBondDisplay(bond)
        if not bond_disp or not bond_disp.IsVisible():
            return False

        if not bond.HasData(self._tag):
            return False

        line_width = disp.GetScale() / 3.0
        color = self._color_gradient.GetColorAt(bond.GetData(self._tag))
        pen = oedepict.OEPen(color, color, oedepict.OEFill_Off, line_width)

        atom_disp_bgn: oedepict.OE2DAtomDisplay = disp.GetAtomDisplay(bond.GetBgn())
        atom_disp_end: oedepict.OE2DAtomDisplay = disp.GetAtomDisplay(bond.GetEnd())

        layer = disp.GetLayer(oedepict.OELayerPosition_Below)
        layer.DrawLine(atom_disp_bgn.GetCoords(), atom_disp_end.GetCoords(), pen)

        return True

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


def _check_report_file(args: argparse.Namespace) -> bool:
    ext = pathlib.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


def _read_molecules(filename: str) -> list[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:
        # enable to read multi-conformer molecules to non OEB file
        compare_title = True
        conf_test = oechem.OEIsomericConfTest(not compare_title)
        ifs.SetConfTest(conf_test)

    mols: list[oechem.OEMCMolBase] = []
    for mol in ifs.GetOEMols():
        if mol.GetDimension() != 3:  # noqa: PLR2004
            oechem.OEThrow.Warning("Ignore molecule without 3D conformations")
            continue
        mols.append(oechem.OEMol(mol))

    if not mols:
        oechem.OEThrow.Fatal(f"No molecules could be read from {filename}")
    return mols


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