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


"""Calculates XLogP of set of molecules and visualizes the atom contributions using property map."""

import argparse
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Depict XLogP of set of molecules."
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "oegrapheme", "oemolprop", "oequacpac"]
__SCRIPT_CATEGORIES__ = ["depiction"]


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_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(
        "--rows",
        type=int,
        default=3,
        choices=range(2, 6),
        metavar="N",
        help="number of rows per page (default: %(default)s)",
    )
    report_group.add_argument(
        "--cols",
        type=int,
        default=2,
        choices=range(1, 3),
        metavar="N",
        help="number of columns per page (default: %(default)s)",
    )
    report_group.add_argument(
        "--page-by-page",
        action="store_true",
        help="write pages of report to separate numbered image files",
    )
    return parser.parse_args()


def main() -> int:
    """Visualizes XLogP of set of molecules."""
    args = parse_options()

    _check_report_file(args)

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

    # initialize multi-page report

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

    # setup depiction options

    width, height = report.GetCellWidth(), report.GetCellHeight()
    opts = oedepict.OE2DMolDisplayOptions(width, height, oedepict.OEScale_AutoScale)
    opts.SetAtomColorStyle(oedepict.OEAtomColorStyle_WhiteMonochrome)

    # depict molecule with XLogP atom contributions
    depict_molecules_xlogp(report, mols, opts)

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

    return os.EX_OK


def set_atom_properties(
    mol: oechem.OEMolBase, tag: int, min_value: float, max_value: float
) -> tuple[float, float]:
    """Attache the XLogP atom contribution to each atom with the given tag."""
    oequacpac.OERemoveFormalCharge(mol)

    atom_values = oechem.OEFloatArray(mol.GetMaxAtomIdx())
    xlogp = oemolprop.OEGetXLogP(mol, atom_values)

    mol.SetTitle(f"{mol.GetTitle()} -- OEXLogP = {xlogp:.2f}")

    for atom in mol.GetAtoms():
        val = atom_values[atom.GetIdx()]
        atom.SetData(tag, val)
        min_value = min(min_value, val)
        max_value = max(max_value, val)

    return min_value, max_value


def depict_molecules_xlogp(
    report: oedepict.OEReport,
    mols: list[oechem.OEMolBase],
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """Generate a report of molecules depicting the atom contribution of XLogP."""
    str_tag = "XLogP"
    int_tag = oechem.OEGetTag(str_tag)

    min_value, max_value = float("inf"), float("-inf")
    for mol in mols:
        min_value, max_value = set_atom_properties(mol, int_tag, min_value, max_value)

    mol_scale = float("inf")
    for mol in mols:
        oedepict.OEPrepareDepiction(mol)
        mol_scale = min(mol_scale, oegrapheme.OEGetMoleculeSurfaceScale(mol, opts))
    opts.SetScale(mol_scale)

    prop_map = oegrapheme.OE2DPropMap(opts.GetBackgroundColor())
    prop_map.SetNegativeColor(oechem.OEDarkGreen)
    prop_map.SetPositiveColor(oechem.OEDarkPurple)
    prop_map.SetLegendLocation(oegrapheme.OELegendLocation_Left)
    prop_map.SetMinValue(min_value)
    prop_map.SetMaxValue(max_value)

    for mol in mols:
        disp = oedepict.OE2DMolDisplay(mol, opts)
        prop_map.Render(disp, str_tag)
        cell = report.NewCell()
        oedepict.OERenderMolecule(cell, disp)


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 outout 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.OEMolBase]:
    ifs = oechem.oemolistream()
    if not ifs.open(filename):
        oechem.OEThrow.Fatal(f"Cannot open {filename} input file!")

    mols: list[oechem.OEMolBase] = [oechem.OEGraphMol(m) for m in ifs.GetOEGraphMols()]
    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())