"""
Utilities for creating MatSci reports.
Copyright Schrodinger, LLC. All rights reserved.
"""
import os
from collections import namedtuple
from matplotlib.backends import backend_pdf
from reportlab.lib.pagesizes import letter
import schrodinger.application.desmond.report_helper as rhelper
from schrodinger.application.matsci import jobutils
from schrodinger.ui.qt import filedialog
from schrodinger.ui.qt import swidgets
from schrodinger.utils import fileutils
DEFAULT_LEFT_MARGIN = 3 / 8. * rhelper.inch
DEFAULT_RIGHT_MARGIN = rhelper.inch / 4.
DEFAULT_TOP_MARGIN = 1 * rhelper.inch
DEFAULT_BOTTOM_MARGIN = 1 * rhelper.inch
EXPORT_PDF = "PDF Report"
EXPORT_IMAGES = "Plots"
EXPORT_DATA = "Data"
MATSCI_LOGO_PATH = os.path.join(fileutils.get_mmshare_scripts_dir(),
                                "event_analysis_dir", "schrodinger_ms_logo.png")
Footer = namedtuple("Footer", ["text", "height"])
[docs]class ReportInfo:
    """
    Manages information about the report request. Contains user inputs as well
    as panel data.
    """
[docs]    def __init__(self, panel, output_dir, base_name, outputs):
        """
        Create an instance and store arguments
        :param `af2.App` panel: The panel for which the report is requested
        :param str output_dir: The directory to write the report files in
        :param str base_name: The base name for the exported files and folders
        :param list outputs: The outputs that the user has requested for the report
        """
        self.panel = panel
        self.output_dir = output_dir
        self.base_name = base_name
        self.outputs = outputs
        self.figures = {} 
[docs]    def pdfRequested(self):
        """
        Determine whether a pdf was requested for the report
        :rtype: bool
        :return: Whether a pdf was requested for the report
        """
        return EXPORT_PDF in self.outputs 
[docs]    def imagesRequested(self):
        """
        Determine whether images were requested for the report
        :rtype: bool
        :return: Whether images were requested for the report
        """
        return EXPORT_IMAGES in self.outputs 
[docs]    def dataRequested(self):
        """
        Determine whether raw data was requested for the report
        :rtype: bool
        :return: Whether raw data was requested for the report
        """
        return EXPORT_DATA in self.outputs 
[docs]    def getPDFPath(self):
        """
        Get the path to the report pdf file
        :rtype: str
        :return: path to the report pdf file
        """
        return os.path.join(self.output_dir, self.base_name + "_report.pdf") 
[docs]    def getCsvFilePath(self):
        """
        Get the path to the data csv file
        :rtype: str
        :return: path to the data csv file
        """
        return os.path.join(self.output_dir, self.base_name + "_data.csv") 
[docs]    def getImagesDirPath(self):
        """
        Get the path to the images directory
        :rtype: str
        :return: path to the images directory
        """
        return os.path.join(self.output_dir, self.base_name + "_images") 
[docs]    def okToWrite(self):
        """
        Ensure that target files and folders don't already exist, or if they do,
        they can be overwritten
        :rtype: bool
        :return: Whether we are clear to write
        """
        paths_to_check = []
        if self.pdfRequested():
            paths_to_check.append(self.getPDFPath())
        if self.imagesRequested():
            paths_to_check.append(self.getImagesDirPath())
        if self.dataRequested():
            paths_to_check.append(self.getCsvFilePath())
        existing_paths = []
        for path in paths_to_check:
            if os.path.exists(path):
                existing_paths.append(os.path.basename(path))
        if existing_paths:
            question = ("The following files or folders already exist in the"
                        " selected directory and will be overwritten, continue?"
                        "\n\n" + "\n".join(existing_paths))
            overwrite = self.panel.question(question,
                                            button1='Yes',
                                            button2='No')
            if not overwrite:
                return False
        return True  
[docs]class PDFBuilder:
    """
    Contains features shared by all MatSci PDF builder classes.
    """
[docs]    def __init__(self,
                 report_info,
                 left_margin=DEFAULT_LEFT_MARGIN,
                 top_margin=DEFAULT_TOP_MARGIN,
                 right_margin=DEFAULT_RIGHT_MARGIN,
                 bottom_margin=DEFAULT_BOTTOM_MARGIN):
        """
        Create a PDFBuilder instance. Store and initialize variables.
        :param `ReportInfo` report_info: The report information object
        :param int left_margin: The left margin for the pdf
        :param int top_margin: The top margin for the pdf
        :param int right_margin: The right margin for the pdf
        :param int bottom_margin: The bottom margin for the pdf
        """
        self.report_info = report_info
        self.panel = report_info.panel
        self.schrodinger_temp_dir = fileutils.get_directory_path(fileutils.TEMP)
        self.figures = list(report_info.figures.values())
        report_path = report_info.getPDFPath()
        self.doc = rhelper.platypus.SimpleDocTemplate(report_path,
                                                      pagesize=letter)
        self.doc.leftMargin = left_margin
        self.doc.rightMargin = right_margin
        self.doc.topMargin = top_margin
        self.doc.bottomMargin = bottom_margin
        self.files_to_cleanup = []
        self.footers = {}  # Dictionary mapping page number to Footer
        self.elements = [] 
[docs]    def build(self):
        """
        Build the report pdf. Calls methods that should be defined in derived
        classes to add flowables and footers.
        """
        self.addFlowables()
        self.addFooters()
        # Pass footers to the canvas to be written as the document is saved.
        # TODO: Find a cleaner way, such as passing a custom function to doc
        CanvasWithHeaderAndFooter.FOOTERS = self.footers
        self.doc.build(self.elements, canvasmaker=CanvasWithHeaderAndFooter)
        self.cleanupTempFiles() 
[docs]    def addFlowables(self):
        """
        Add flowables to the pdf, such as headers, titles, paragraphs, tables
        and images. Should be defined in derived classes.
        """
        raise NotImplementedError("The addFlowables method should be"
                                  " defined in all derived classes.") 
[docs]    def writeParagraphs(self, paragraphs):
        """
        Add the paragraphs to the pdf as flowables.
        :param list paragraphs: A list of strings, each of which should be a
            paragraph
        """
        for paragraph in paragraphs:
            rhelper.pargph(self.elements, paragraph, leading=15)
            self.addSpacer() 
[docs]    def getTempImageFilename(self):
        """
        Get a temporary image file path and add it to the list of files to be
        cleaned up.
        :rtype: str
        :return: A temporary image file path
        """
        temp_filename = fileutils.get_next_filename(
            os.path.join(self.schrodinger_temp_dir, "temp_image.png"), "")
        self.files_to_cleanup.append(temp_filename)
        return temp_filename 
[docs]    def cleanupTempFiles(self):
        """
        Clean up all the temporary files
        """
        for temp_file in self.files_to_cleanup:
            fileutils.force_remove(temp_file) 
[docs]    def getTableStyle(self,
                      column_headers=True,
                      row_headers=False,
                      align_center=True):
        """
        Create a platypus table style based on desired headers and alignment
        :param bool column_headers: Whether the table columns have headers
        :param bool row_headers: Whether the table rows have headers
        :param bool align_center: Whether align center should be used for all cells
        :rtype: list of tuples
        :return: A platypus table style based on desired headers and alignment
        """
        # Format: (0, 1), (-1, -1): From column 0, row 1, to last column and row
        table_style = [('NOSPLIT', (0, 0), (-1, -1)),
                       ('BOTTOMPADDING', (0, 0), (-1, -1), 1),
                       ('TOPPADDING', (0, 0), (-1, -1), 1)] # yapf:disable
        if column_headers:
            table_style.extend([
                ('BOTTOMPADDING', (0, 0), (-1, 0), 2),
                ('TEXTCOLOR', (0, 0), (-1, 0), rhelper.gray)
            ]) # yapf:disable
        if row_headers:
            table_style.append(('TEXTCOLOR', (0, 0), (0, -1), rhelper.gray))
        if align_center:
            table_style.append(('ALIGN', (0, 0), (-1, -1), 'CENTER'))
        return table_style 
[docs]    def addSpacer(self):
        """
        Add a space between the last and next flowables
        """
        rhelper.add_spacer(self.elements)  
[docs]class ReportOutputsDialog(swidgets.SDialog):
    """
    Dialog for allowing the user to specify requested report outputs.
    Calls the report generation method of the panel when the user accepts.
    """
    DIR_SELECTOR_ID = "REPORT_UTILS_DIR_SELECTOR_ID"
[docs]    def __init__(self,
                 *args,
                 default_base_name="",
                 report_btn=True,
                 images_btn=True,
                 data_btn=True,
                 **kwargs):
        """
        Create an instance.
        :param bool report_btn: Whether PDF Report checkbox should be in the dialog
        :param bool images_btn: Whether Plots checkbox should be in the dialog
        :param bool data_btn: Whether Data checkbox should be in the dialog
        """
        self.default_base_name = default_base_name
        self.report_btn = report_btn
        self.images_btn = images_btn
        self.data_btn = data_btn
        kwargs['title'] = kwargs.get('title', 'Export Options')
        super().__init__(*args, **kwargs) 
[docs]    def layOut(self):
        """
        Lay out the widgets.
        """
        layout = self.mylayout
        dator = swidgets.FileBaseNameValidator()
        self.base_name_le = swidgets.SLabeledEdit(
            "File base name:",
            edit_text=self.default_base_name,
            validator=dator,
            always_valid=True,
            layout=layout,
            stretch=False)
        self.base_name_le.setMinimumWidth(170)
        self.checkboxes = {}
        if self.report_btn:
            self.checkboxes[EXPORT_PDF] = swidgets.SCheckBox(EXPORT_PDF,
                                                             checked=True,
                                                             layout=layout)
        if self.images_btn:
            self.checkboxes[EXPORT_IMAGES] = swidgets.SCheckBox(EXPORT_IMAGES,
                                                                checked=True,
                                                                layout=layout)
        if self.data_btn:
            self.checkboxes[EXPORT_DATA] = swidgets.SCheckBox(EXPORT_DATA,
                                                              checked=True,
                                                              layout=layout) 
[docs]    def accept(self):
        """
        Get user inputs and call the report generation method of the panel.
        """
        outputs = []
        for label, checkbox in self.checkboxes.items():
            if checkbox.isChecked():
                outputs.append(label)
        if not outputs:
            self.error("At least one option needs to be selected.")
            return
        output_dir = filedialog.get_existing_directory(
            self, caption="Choose Export Directory", id=self.DIR_SELECTOR_ID)
        if not output_dir:
            return
        report_info = ReportInfo(self.master, output_dir,
                                 self.base_name_le.text(), outputs)
        if not report_info.okToWrite():
            return
        self.user_accept_function(report_info)
        self.info("The requested items have been exported.")
        return super().accept()  
[docs]def save_images(report_info):
    """
    Save all images in the report information object to the output directory
    :param `ReportInfo` report_info: The report information object containing
        the figures
    """
    image_dir = report_info.getImagesDirPath()
    fileutils.force_rmtree(image_dir)
    fileutils.mkdir_p(image_dir)
    for title, figure in report_info.figures.items():
        file_path = os.path.join(image_dir, title + ".png")
        save_figure(figure, file_path, report_info.panel) 
[docs]def sub_script(text, size=10):
    """
    :type text: str or int or float
    :param text: The text to write as a subscript
    :param int size: The font size of the subscript
    :rtype: str
    :return: The text formatted to appear as a subscript in the report
    """
    return f"<sub><font size={size}>{text}</font></sub>" 
[docs]def super_script(text, size=10):
    """
    :type text: str or int or float
    :param text: The text to write as a superscript
    :param int size: The font size of the superscript
    :rtype: str
    :return: The text formatted to appear as a superscript in the report
    """
    return f"<super><font size={size}>{text}</font></super>"