import math, os, subprocess, sys

class LatexBuilder(object):
    """
    A simple class for automatically generating LaTeX and compiling reports in python
    """

    LATEX_TEMPLATE_FILE = '../templates/report.tex'
    MAKEFILE_TEMPLATE_FILE = '../templates/latex_makefile'
    SUBFIGURE_RELATIVE_WIDTH = 0.65

    def __init__(self, author, title, latex_template_file = LATEX_TEMPLATE_FILE, makefile_template_file = MAKEFILE_TEMPLATE_FILE, graphicspath = None):
        """

        :param author: Author of the report
        :param title: Title of the report
        :param latex_template_file: Latex template file to use
        :param makefile_template_file: makefile template file to use
        :param graphicspath: Path to figures
        :return:
        """
        self.latex_template_file = latex_template_file
        self.makefile_template_file = makefile_template_file
        self.graphicspath = graphicspath

        self.report_author = author
        self.report_title = title

        self.pages = []

    def addPage(self, text):
        """
        Add page to the report
        :param text:
        :return:
        """
        self.pages.append(text)

    def addSection(self, section_name, level = 0, clearPage = False):
        """
        Add new section to the report
        :param section_name: Name of the section
        :param level: Level of the section
        :param clearPage: Whether to start section from new page
        :return:
        """
        self.addPage("%s\%ssection{%s}" % ('\clearpage' if clearPage else '', 'sub' * level,section_name.replace('_','-')))

    def addGroupFigure(self, title, image_list = None, column_count = None, image_caption_list = None, row_image_list = None, row_caption_list = None, image_label_suffix='', caption_suffix = '', image_angle = 0, max_images_per_figure = None):
        """
        Add group figure
        :param title: Title of the main figure
        :param image_list: List of subfigures
        :param column_count: Number of columns in the figure
        :param image_caption_list: List of captions for subfigures
        :param row_image_list:
        :param row_caption_list: List of captions for rows
        :param image_label_suffix: Suffix for image labels
        :param caption_suffix: Suffix for figure caption
        :param image_angle: Rotate images
        :param max_images_per_figure: If number of subfigures exceeds this number, split the figure into several parts
        :return:
        """
        if max_images_per_figure is None or len(image_caption_list)<max_images_per_figure:
            self.addPage(self.getGroupFigureLatexString(title, image_list, column_count, image_caption_list, row_image_list, row_caption_list, image_label_suffix, caption_suffix, image_angle))
        else:#split into several figures
            remaining_figures = len(image_list)
            i_start = 0
            num_of_figures = int(math.ceil(remaining_figures / float(max_images_per_figure)))
            while remaining_figures>0:
                self.addPage(self.getGroupFigureLatexString(title + ' - Part %s/%s' % (int(i_start / max_images_per_figure)+1,num_of_figures),
                                                            image_list[i_start:(i_start + min(remaining_figures,max_images_per_figure))],
                                                            column_count,
                                                            image_caption_list[i_start:(i_start + min(remaining_figures,max_images_per_figure))],
                                                            row_image_list, row_caption_list, image_label_suffix, caption_suffix, image_angle))

                i_start += min(remaining_figures,max_images_per_figure)
                remaining_figures -= min(remaining_figures,max_images_per_figure)


    def save(self, output_file, create_makefile = False, makefile_pdf_quality = 'prepress'):
        """
        Save text file
        :param output_file:
        :param create_makefile:
        :param makefile_pdf_quality:
        :return:
        """

        if not output_file.find('.tex') == len(output_file) - 4:
            output_file += '.tex'

        if not os.path.exists(os.path.dirname(output_file)):
            os.makedirs(os.path.dirname(output_file))

        if isinstance(self.graphicspath, basestring):
            self.graphicspath = [self.graphicspath]

        template_values = {
            'author': self.report_author,
            'title': self.report_title,
            'graphicspath': ''.join(['{%s}' % path for  path in self.graphicspath]),
            'body': '\n\n'.join(self.pages)
        }

        template = file(self.latex_template_file, 'r').read()

        open(output_file, mode='wb').write(template % template_values)


        if create_makefile:

            template_values =  {
                'tex_filename': os.path.splitext(os.path.basename(output_file))[0],
                'pdf_quality': makefile_pdf_quality
            }

            template_makefile = file(self.makefile_template_file, 'r').read()

            open(os.path.dirname(output_file) + '/makefile', mode='w').write(
                template_makefile % template_values
            )



    def getGroupFigureLatexString(self, title, image_list = None, column_count = None, image_caption_list = None, row_image_list = None, row_caption_list = None, image_label_suffix='', caption_suffix = '', image_angle = 0):
        """
        Get latex string for a figure containing sub figures
        :param image_list: List of all images to be included
        :param title: Title of the figure
        :param column_count: Number of columns
        :param image_caption_list:
        :param row_caption_list: List of all images to be included groupped by row
        :param image_label_suffix:
        :return:
        """
        latex_string = '''
\\begin{figure}[H]
\centering\n'''


        #Either all images in one list or images groupped by row
        if image_list is None and row_image_list is None:
            return ''

        if len(image_list) != len(image_caption_list):
            raise Exception('There should be the same number of plots and captions: %s != %s',len(image_list),len(image_caption_list))

        #If grouped by row, determine column count (max row length)
        if column_count is None:
            if image_list is None:
                column_count = 0
                for image_list_for_the_row in row_image_list:
                    if len(image_list_for_the_row)>column_count:
                        column_count = len(image_list_for_the_row)
            else:
                column_count = int(math.ceil(math.sqrt(len(image_list))))



        # if column_count > 9:
        #     image_width = str(0.75/(column_count))[:5]
        # elif column_count > 4:
        #     image_width = str(0.85/(column_count))[:5]
        # else:
        #     image_width = str(0.65/(column_count))[:5]

        image_width = str(LatexBuilder.SUBFIGURE_RELATIVE_WIDTH/(column_count))[:5]

        latex_string += '\\begin{tabular}{%s}\n' % ('c' * column_count)

        caption = '\caption{\\textbf{%s.} ' % title

        #If all images in one list, group them by row based on column count
        if row_image_list is None:
            row_image_list = []
            temp_list = []
            i=0
            for image in image_list:
                i+=1
                temp_list.append(image)
                if i%column_count == 0:
                    row_image_list.append(temp_list)
                    temp_list = []

            if len(temp_list)>0:
                row_image_list.append(temp_list)

        first_in_the_row_image_name = None #used for references in row captions
        i=0 #used for row captions
        k=0 #used for image captions
        for row in row_image_list:
            j=0 #used to determine the first image in the row
            for image_file in row:
                head, tail = os.path.split(image_file)
                fileName, fileExtension = os.path.splitext(tail)
                image_name = fileName#image name is used for labels

                if j==0: #first image in the row
                    first_in_the_row_image_name = image_name
                j+=1

                latex_string += '\subfloat[]{\label{fig:%s}\includegraphics[width=%s\linewidth, angle=%s]{%s}} &\n' % (image_name + image_label_suffix, image_width, image_angle, image_file)

                #if needed, add image caption
                if image_caption_list is not None:
                    caption += ' \protect\subref{fig:%s} -  ' % (image_name + image_label_suffix)  + image_caption_list[k] + ','
                k+=1

            latex_string = latex_string[:-3] + ' \\'+'\\' + '\n'

            #if needed, add row caption
            if row_caption_list is not None:
                caption += ' \protect\subref{fig:%s}-\protect\subref{fig:%s} -  ' % (first_in_the_row_image_name + image_label_suffix,image_name + image_label_suffix)\
                           + row_caption_list[i] + ','
            i+=1

        caption = caption[:-1] + '.' + caption_suffix + '}'

        latex_string += '\end{tabular}%s\end{figure}\n' % caption

        return latex_string





    def get_latex_figure_string(self, image_name, caption):
        """
        Get figure containing one image
        """
        return '''
\\begin{figure}[H]
\centering
\includegraphics[width=0.9\\textwidth]{''' + image_name + '''}
\caption{''' +  caption + '''}
\end{figure}
'''


    def compile(self, report_directory, report_name, pdflatex = False):
        """
        Compile document
        :param report_directory: Outpu directory
        :param report_name: Name of the report
        :param pdflatex: Use pdflatex
        :return:
        """
        if not os.path.exists(report_directory):
            os.makedirs(report_directory)

        self.save(report_directory + report_name + '/' + report_name + '.tex',create_makefile=True, makefile_pdf_quality='screen')

        print 'Compiling document...'
        p = subprocess.Popen(['make','pdflatex' if pdflatex else 'pdf'], stdout=subprocess.PIPE, cwd = report_directory + report_name)
        for line in iter(p.stdout.readline, ''):
            sys.stdout.write(line)
        p.stdout.close()
