#!/usr/bin/env python

""" Tool that takes python script and runs it. Returns the results
embedded in the code in a pdf """

#TODO: - Fail back to plain text if errors on the rst 
#      - Output ps, rst, html, moinmoin syntax ?
#      - Check to see if pylab is imported before trying to override the
#           "show()" command.
#      - Allow for embedding of LaTeX in the code.
#      - Long, long term: use reportlab to avoid the dependencies on
#          LaTeX

DEBUG = True

#---------------------------- Imports ----------------------------------------
# To deal with arguments
import sys
# To deal with regexp
import re
import os
import os.path

#------------------------------- Subroutines ---------------------------------
def TeX2pdf(fileName):
    """ Compiles a TeX file with pdfLaTeX and cleans up the mess afterwards """
    print >> sys.stderr, "Compiling document to pdf"
    fileName=os.path.splitext(fileName)[0]
    os.system("pdflatex --interaction scrollmode " + fileName + ".tex")
    print "Cleaning up"
    os.unlink(fileName+".tex")
    os.unlink(fileName+".log")
    os.unlink(fileName+".aux")

def epstopdf(figureList):
    """ Converts the eps files generated by the script to pdf files """
    os.environ['GS_OPTIONS'] = "-dUseFlatCompression=true -dPDFSETTINGS=/prepress -sColorImageFilter=FlateEncode -sGrayImageFilter=FlateEncode -dAutoFilterColorImages=false -dAutoFilterGrayImages=false -dEncodeColorImages=false -dEncodeGrayImages=false -dEncodeMonoImages=false"
    for file in figureList:
        os.system("epstopdf --nocompress " + file)
        os.unlink(file)

def rmFigures(figureList):
    """ Removes the pdf files of the figures generated by the script """
    for file in figureList:
        os.unlink(file[:-4]+".pdf")

def rst2latex(restString):
    """ Calls docutils' engine to convert a rst string to a LaTeX file. """
    from docutils import core as docCore
    overrides = {'output_encoding': 'latin1',
                'initial_header_level': 0}
    TeXstring = docCore.publish_string(
                source=restString, 
                writer_name='latex', settings_overrides=overrides)
    return TeXstring

def py2blocks(fileName):
    """ Returns the list of blocks to be processed of a given python file, 
    with there starting line number in the original file."""
    blockList=[["",1],]
    file=open(fileName)
    for linenum, line in enumerate(file):
        if line[0]==" " or line[0]=="\t" :
            blockList[-1][0]+=line
        elif line[0:3]=="## " :
            blockList+=[[line[3:],"commentBlock"],]
        elif line[0:2]=="##" :
            blockList+=[[line[2:],"commentBlock"],]
        else :
            blockList+=[[line,linenum],]
    return blockList

def executeBlock( namespace, block, figureList):
    """ Excute a python command block, and returns the locals variables
    created, the stderr and the stdout generated, and the list of figures
    generated."""
    blockText = block[0]
    lineNumber = block[1]
    outValue = ""
    # define a "show" command, to replace pylab's show()
    class _Show(object):
        def __init__(self, figureList ):
            self.figureList = figureList
        def __call__(self):
            figureName = 'pylab2pdf_%d.eps' % (len(self.figureList) )
            self.figureList += (figureName, )
            print "Here goes figure " + figureName
            savefig(figureName)

    show = _Show(figureList)
    namespace['show'] = show
    # to treat StdIn, StdOut as files:
    import StringIO
    
    # create file-like string to capture output
    codeOut = StringIO.StringIO()
    codeErr = StringIO.StringIO()
   
    try:
        # capture output and errors
        sys.stdout = codeOut
        sys.stderr = codeErr
        
        exec blockText in namespace
        outValue=codeOut.getvalue()
        errorValue=codeErr.getvalue()
        
        # restore stdout and stderr
        sys.stdout = sys.__stdout__
        sys.stderr = sys.__stderr__
    except Exception, inst:
        print inst 
        sys.stdout = sys.__stdout__
        sys.stderr = sys.__stderr__
        errorValue = inst.__str__()
        print "Error in executing script on line ",lineNumber ,": ", inst 
    
    if errorValue:  print errorValue
    namespace=globals()
    namespace.update(locals())
    
    codeOut.close()
    codeErr.close()
    
    if outValue:    print outValue
    return namespace, outValue, errorValue, show.figureList

def condenseOutputList(outputList):
    """ Takes the "outputList", made of block of different type and merges 
    successiv blocks of the same type."""
    outList=[['commentBlock',''],]
    for block in outputList:
        if block[0]==outList[-1][0]:
            outList[-1][1]+=block[1]
            if block[0]=='outputBlock':
                outList[-1][2]+=block[2]
                outList[-1][1]=re.sub(r"(\n)+",r"\n",outList[-1][1])
        else:
            outList+=[block]
    return outList

def addIndent(lines):
    """ Add an indentation (4 spaces) to all the lines"""
    output=re.sub(r"\n","\n    ",lines)
    return output

def protect(string):
    """ Protects all the "\" in a string by adding a second one before """
    return re.sub(r"\\",r"\\\\",string)

def RstString(outputList):
    """ Giving a list of blocks being either input blocks, output blocks or 
    error blocks, returns a rst string """
    emptyListing=re.compile(r"""\\begin\{lstlisting\}\{\}\s*\\end\{lstlisting\}""",re.DOTALL)
    inputBlocktpl = r"""
    
.. raw:: LaTeX

    {\inputBlocksize
    \lstset{backgroundcolor=\color{lightblue},fillcolor=\color{lightblue},numbers=left,name=pythoncode,firstnumber=%(linenumber)d,xleftmargin=0pt,fillcolor=\color{white},frame=single,fillcolor=\color{lightblue},rulecolor=\color{lightgrey}}
    \begin{lstlisting}{}
    %(textBlock)s
    \end{lstlisting}
    }


"""
    errorBlocktpl = r"""

.. raw:: LaTeX


    {\color{red}{\bfseries Error: }
    \begin{verbatim}%s\end{verbatim}}
    
"""
    outputBlocktpl = r"""
.. raw:: LaTeX

    \vskip -1ex
    \lstset{backgroundcolor=,numbers=none,name=answer,xleftmargin=3ex,frame=none}
    \begin{lstlisting}{}
    %s
    \end{lstlisting}
    \vskip -1ex
    
    """
    figuretpl = r'''
\end{lstlisting}
\\vskip -0.5em
\\vskip 0.3cm
\\centerline{\includegraphics[scale=0.5]{%s}}
\\begin{lstlisting}{}'''
    
    outString=""
    for block in outputList:
        if block[0]=="inputBlock":
            data = {'linenumber' : block[2],
                    'textBlock' : addIndent(block[1]),
                    }
            TeXtext = inputBlocktpl % data
            TeXtext=re.sub(emptyListing,"",TeXtext)
            outString += TeXtext
        elif block[0]=="errorBlock":
            outString += errorBlocktpl % (addIndent(block[1]))
        elif block[0]=="commentBlock":
            outString += "\n" + block[1] + "\n" 
        elif block[0]=="outputBlock":
            if not len(block[2])==0:
                for file in block[2]:
                    block[1]=re.sub("Here goes figure "+ re.escape(file),
                            (figuretpl % ( os.path.splitext(file)[0] )) ,
                            block[1])
            TeXtext = outputBlocktpl % (addIndent(block[1]))
            TeXtext = re.sub(emptyListing,"",TeXtext)
            outString += TeXtext
    outString += "\n"
    return outString

#------------------------------- Entry point ---------------------------------
def commandlineCall():

    from optparse import OptionParser
    usage = """usage: %prog [options] pythonfile
    
    Processes a python script and pretty prints the results using LateX. If 
    the script uses "show()" commands (from pylab) they are caught by 
    pylab2pdf and the resulting graphs are inserted in the output pdf.
    Comments lines starting with two "##" are interprated as rst lines
    and pretty printed accordingly in the pdf.
    """
    parser = OptionParser(usage=usage)
    parser.add_option("-o", "--outfile", dest="outfile",
                    help="write report to FILE", metavar="FILE")
    parser.add_option("-d", "--double",
                    dest="double", action="store_true", default=False,
                    help="compile to two columns per page")
    parser.add_option("-n", "--nocode",
                    dest="nocode", action="store_true", default=False,
                    help="do not display the source code")
    parser.add_option("-t", "--type",
                    action="store", type="string", dest="outputtype",
                    default="pdf",
                    help="output a tex file instead of a pdf")
    parser.add_option("-q", "--quiet",
                    action="store_false", dest="verbose", default=True,
                    help="don't print status messages to stdout")

    (options, args) = parser.parse_args()


    main(args[0], outputtype=options.outputtype, 
                double=options.double,
                nocode=options.nocode,
                outfile=options.outfile
                )


def main(pyFileName,outputtype="pdf", double=False, nocode=False,outfile=None):

    if not outfile:
        outfile = os.path.splitext(pyFileName)[0]+".pdf"
    texFileName=os.path.splitext(outfile)[0]+".tex"

    global DEBUG

    LaTeXpreambule = r"""
    \usepackage{listings}
    \usepackage{color}
    \usepackage{graphicx}

    \definecolor{darkgreen}{cmyk}{0.7, 0, 1, 0.5}
    \definecolor{darkblue}{cmyk}{1, 0.8, 0, 0}
    \definecolor{lightblue}{cmyk}{0.05,0,0,0.05}
    \definecolor{grey}{cmyk}{0.1,0.1,0.1,1}
    \definecolor{lightgrey}{cmyk}{0,0,0,0.5}
    \definecolor{purple}{cmyk}{0.8,1,0,0}

    \makeatletter
        \let\@oddfoot\@empty\let\@evenfoot\@empty
        \def\@evenhead{\thepage\hfil\slshape\leftmark
                        {\rule[-0.11cm]{-\textwidth}{0.03cm}
                        \rule[-0.11cm]{\textwidth}{0.03cm}}}
        \def\@oddhead{{\slshape\rightmark}\hfil\thepage
                        {\rule[-0.11cm]{-\textwidth}{0.03cm}
                        \rule[-0.11cm]{\textwidth}{0.03cm}}}
        \let\@mkboth\markboth
        \markright{{\bf %s }\hskip 3em  \today}
        \def\maketitle{
            \centerline{\Large\bfseries\@title}
            \bigskip
        }
    \makeatother

    \lstset{language=python,
            extendedchars=true,
            basicstyle=\ttfamily,
            keywordstyle=\sffamily\bfseries,
            identifierstyle=\sffamily,
            commentstyle=\slshape\color{darkgreen},
            stringstyle=\rmfamily\color{blue},
            showstringspaces=false,
            tabsize=4,
            breaklines=true,
            numberstyle=\footnotesize\color{grey},
            classoffset=1,
            morekeywords={zeros,zeros_like,ones,ones_like,array,rand,indentity,mat,vander},keywordstyle=\color{darkblue},
            classoffset=2,
            otherkeywords={[,],=,:},keywordstyle=\color{purple}\bfseries,
            classoffset=0
            }
    """ % ( re.sub( "_", r'\\_', pyFileName) )

    if nocode:
        LaTeXcolumnsep=r"""
    \setlength\columnseprule{0.4pt}
    """
    else:
        LaTeXcolumnsep=""


    LaTeXdoublepage = r"""
    \usepackage[landscape,left=1.5cm,right=1.1cm,top=1.8cm,bottom=1.2cm]{geometry}
    \usepackage{multicol}
    \def\inputBlocksize{\normalsize}
    \makeatletter
        \renewcommand\normalsize{%
        \@setfontsize\normalsize\@ixpt\@xipt%
        \abovedisplayskip 8\p@ \@plus4\p@ \@minus4\p@
        \abovedisplayshortskip \z@ \@plus3\p@
        \belowdisplayshortskip 5\p@ \@plus3\p@ \@minus3\p@
        \belowdisplayskip \abovedisplayskip
        \let\@listi\@listI}
        \normalsize
        \renewcommand\small{%
        \@setfontsize\small\@viiipt\@ixpt%
        \abovedisplayskip 5\p@ \@plus2\p@ \@minus2\p@
        \abovedisplayshortskip \z@ \@plus1\p@
        \belowdisplayshortskip 3\p@ \@plus\p@ \@minus2\p@
        \def\@listi{\leftmargin\leftmargini
                    \topsep 3\p@ \@plus\p@ \@minus\p@
                    \parsep 2\p@ \@plus\p@ \@minus\p@
                    \itemsep \parsep}%
        \belowdisplayskip \abovedisplayskip
        }
        \renewcommand\footnotesize{%
        \@setfontsize\footnotesize\@viipt\@viiipt
        \abovedisplayskip 4\p@ \@plus2\p@ \@minus2\p@
        \abovedisplayshortskip \z@ \@plus1\p@
        \belowdisplayshortskip 2.5\p@ \@plus\p@ \@minus\p@
        \def\@listi{\leftmargin\leftmargini
                    \topsep 3\p@ \@plus\p@ \@minus\p@
                    \parsep 2\p@ \@plus\p@ \@minus\p@
                    \itemsep \parsep}%
        \belowdisplayskip \abovedisplayskip
        }
        \renewcommand\scriptsize{\@setfontsize\scriptsize\@vipt\@viipt}
        \renewcommand\tiny{\@setfontsize\tiny\@vpt\@vipt}
        \renewcommand\large{\@setfontsize\large\@xpt\@xiipt}
        \renewcommand\Large{\@setfontsize\Large\@xipt{13}}
        \renewcommand\LARGE{\@setfontsize\LARGE\@xiipt{14}}
        \renewcommand\huge{\@setfontsize\huge\@xivpt{18}}
        \renewcommand\Huge{\@setfontsize\Huge\@xviipt{22}}
        \setlength\parindent{14pt}
        \setlength\smallskipamount{3\p@ \@plus 1\p@ \@minus 1\p@}
        \setlength\medskipamount{6\p@ \@plus 2\p@ \@minus 2\p@}
        \setlength\bigskipamount{12\p@ \@plus 4\p@ \@minus 4\p@}
        \setlength\headheight{12\p@}
        \setlength\headsep   {25\p@}
        \setlength\topskip   {9\p@}
        \setlength\footskip{30\p@}
        \setlength\maxdepth{.5\topskip}
    \makeatother

    \AtBeginDocument{
    \setlength\columnsep{1.1cm}
    """ + LaTeXcolumnsep + r"""
    \begin{multicols}{2}
    \small}
    \AtEndDocument{\end{multicols}}
    """

    if double:
        LaTeXpreambule += LaTeXdoublepage
    else:
        LaTeXpreambule += """\usepackage[top=2.1cm,bottom=2.1cm,left=2cm,right=2cm]{geometry}
    \def\inputBlocksize{\small}
        """

    # Slice the input file into blocks
    blockList=py2blocks(pyFileName)

    # Process the blocks
    # Each element of "outputList" is made of a string, describing the
    # type of the entry, either "inputBlock", or "outputBlock", or
    # "errorBlock", and the content of the entry.
    namespace = {}
    outputList = []
    figureList = ()
    for block in blockList:
        if block[1]=="commentBlock":
            outputList+=[['commentBlock',block[0]],]
        else:
            [ namespace, outValue, errorValue, figureList ] = executeBlock(namespace, block, figureList)
            if not nocode:
                outputList+=[['inputBlock', block[0], block[1]], ]
            if errorValue:
                outputList+=[['errorBlock', errorValue], ]
            if outValue:
                    outputList+=[['outputBlock', outValue, figureList], ]

    if DEBUG:
        f = open('debug','w')
        f.write(outputList.__repr__())
        f.close()

    # Clean up the output list
    outputList=condenseOutputList(outputList)

    if DEBUG:
        f = open('debug2','w')
        f.write(outputList.__repr__())
        f.close()

    restString=RstString(outputList)

    if DEBUG:
        restFile = open(texFileName[:-4]+".rst","w")
        restFile.write(restString)
        restFile.close()

    texString = rst2latex(restString)

    texString=re.sub(r"\\begin{document}",  protect(LaTeXpreambule) + r"\\begin{document}",texString)

    # Write the tex file
    texFile = open(texFileName,"w")
    texFile.write(texString)
    texFile.close()

    # Compile it
    epstopdf(figureList)

    if outputtype=="pdf":
        TeX2pdf(texFileName)
        rmFigures(figureList)


    if DEBUG:
        texFile = open(texFileName,"w")
        texFile.write(texString)
        texFile.close()


if __name__ == '__main__':
    commandlineCall()