#!/usr/bin/env python3
# (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 OpenEye 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 OpenEye offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  CADENCE 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.

"""Converts a CSV or SDF file into PDF with molecular depictions."""

import argparse
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Convert CSV or SDF files into PDF reports with molecular depictions"
__SCRIPT_TOOLKITS__ = ["oechem", "oedepict", "depiction"]
__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]",
    )

    io_group = parser.add_argument_group("Input options")
    io_group.add_argument(
        "--mol",
        "--mol-file",
        metavar="MOL-FILE",
        type=str,
        required=True,
        help="input molecule file (.csv or .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(1, 4),
        metavar="N",
        help="number of rows 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 (default: %(default)s)",
    )

    parser.add_argument("--help-image", action=HelpPreviewAction)
    parser.add_argument(
        "--save-console-svg",
        default=False,
        action="store_true",
        help=f"run command and capture console output in {__SCRIPT_NAME__}.svg file",
    )
    return parser.parse_args()


def main() -> int:
    """Convert molecule file to PDF report."""
    args = parse_options()

    _check_report_file(args)

    mol_list: list[oechem.OEMolBase] = read_molecules(args.mol)
    tags: list[str] = collect_data_tags(mol_list)

    rows, cols = args.rows, 2
    report_opts: oedepict.OEReportOptions = oedepict.OEReportOptions(rows, cols)
    report_opts.SetHeaderHeight(25)
    report_opts.SetFooterHeight(25)
    report_opts.SetCellGap(2)
    report_opts.SetPageMargins(10)
    report: oedepict.OEReport = oedepict.OEReport(report_opts)

    cell_width, cell_height = report.GetCellWidth(), report.GetCellHeight()
    opts: oedepict.OE2DMolDisplayOptions = oedepict.OE2DMolDisplayOptions(
        cell_width, cell_height, oedepict.OEScale_AutoScale
    )

    depict_molecules_with_data(
        report, mol_list, pathlib.Path(args.mol).name, tags, opts
    )

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

    return os.EX_OK


def read_molecules(mol_filename: str) -> list[oechem.OEMolBase]:
    """Read molecules from a CSV or SDF file and return a list of OEMolBase objects."""
    mol_path = pathlib.Path(mol_filename)
    if not mol_path.exists():
        oechem.OEThrow.Fatal(f"Cannot open input file '{mol_path.name}'!")

    ifs = oechem.oemolistream()
    if not ifs.open(str(mol_path)):
        oechem.OEThrow.Fatal(f"Cannot open input file '{mol_path.name}'!")

    if ifs.GetFormat() not in [oechem.OEFormat_CSV, oechem.OEFormat_SDF]:
        oechem.OEThrow.Fatal("Input must be a CSV or SDF file!")
    mol_list: list[oechem.OEMolBase] = [
        oechem.OEGraphMol(m) for m in ifs.GetOEGraphMols()
    ]
    return mol_list


def collect_data_tags(mol_list: list[oechem.OEMolBase]) -> list[str]:
    """Collect all unique SD data tags from molecules."""
    tags: list[str] = []
    for mol in mol_list:
        for dp in oechem.OEGetSDDataIter(mol):
            tag: str = dp.GetTag()
            if tag not in tags:
                tags.append(tag)
    return tags


def depict_molecules_with_data(
    report: oedepict.OEReport,
    mol_list: list[oechem.OEMolBase],
    input_name: str,
    tags: list[str],
    opts: oedepict.OE2DMolDisplayOptions,
) -> None:
    """Depict molecules with their associated data in report."""
    for mol in mol_list:
        cell: oedepict.OEImageBase = report.NewCell()
        oedepict.OEPrepareDepiction(mol)
        disp: oedepict.OE2DMolDisplay = oedepict.OE2DMolDisplay(mol, opts)
        oedepict.OERenderMolecule(cell, disp)
        oedepict.OEDrawCurvedBorder(cell, oedepict.OELightGreyPen, 10.0)

        # Render corresponding data
        cell = report.NewCell()
        render_data(cell, mol, tags)

    # Add input filename to headers
    header_font: oedepict.OEFont = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Default,
        12,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )
    header_pos: oedepict.OE2DPoint = oedepict.OE2DPoint(
        report.GetHeaderWidth() / 2.0, report.GetHeaderHeight() / 2.0
    )

    for header in report.GetHeaders():
        header.DrawText(header_pos, input_name, header_font)

    # Add page number to footers
    footer_font: oedepict.OEFont = oedepict.OEFont(
        oedepict.OEFontFamily_Default,
        oedepict.OEFontStyle_Default,
        12,
        oedepict.OEAlignment_Center,
        oechem.OEBlack,
    )
    footer_pos: oedepict.OE2DPoint = oedepict.OE2DPoint(
        report.GetFooterWidth() / 2.0, report.GetFooterHeight() / 2.0
    )

    for page_idx, footer in enumerate(report.GetFooters()):
        footer.DrawText(footer_pos, f"- {page_idx + 1} -", footer_font)


def render_data(
    image: oedepict.OEImageBase, mol: oechem.OEMolBase, tags: list[str]
) -> None:
    """Render SD data as a table in the image."""
    data: list[tuple[str, str]] = []
    for tag in tags:
        value: str = "N/A"
        if oechem.OEHasSDData(mol, tag):
            value = oechem.OEGetSDData(mol, tag)
        data.append((tag, value))

    nr_data: int = len(data)

    table_opts: oedepict.OEImageTableOptions = oedepict.OEImageTableOptions(
        nr_data, 2, oedepict.OEImageTableStyle_LightBlue
    )
    table_opts.SetColumnWidths([10, 20])
    table_opts.SetMargins(2.0)
    table_opts.SetHeader(False)
    table_opts.SetStubColumn(True)
    table: oedepict.OEImageTable = oedepict.OEImageTable(image, table_opts)

    for row, (tag, value) in enumerate(data):
        cell: oedepict.OEImageBase = table.GetCell(row + 1, 1)
        table.DrawText(cell, f"{tag}:")
        cell = table.GetBodyCell(row + 1, 1)
        table.DrawText(cell, value)


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


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