"""
This module supports embedded TeX expressions in matplotlib via dvipng
and dvips for the raster and postscript backends.  The tex and
dvipng/dvips information is cached in ~/.matplotlib/tex.cache for reuse between
sessions

Requirements:

  tex

  *Agg backends: dvipng

  PS backend: latex w/ psfrag, dvips, and Ghostscript 8.51
  (older versions do not work properly)

Backends:

  Only supported on *Agg and PS backends currently
  

For raster output, you can get RGBA numerix arrays from TeX expressions
as follows

  texmanager = TexManager()
  s = r'\TeX\ is Number $\displaystyle\sum_{n=1}^\infty\frac{-e^{i\pi}}{2^n}$!'
  Z = self.texmanager.get_rgba(s, size=12, dpi=80, rgb=(1,0,0))

To enable tex rendering of all text in your matplotlib figure, set
text.usetex in your matplotlibrc file (http://matplotlib.sf.net/matplotlibrc)
or include these two lines in your script:
from matplotlib import rc
rc('text', usetex=True)

"""

import glob, os, sys, md5, shutil
from matplotlib import get_configdir, get_home, get_data_path, \
     rcParams, verbose
from matplotlib._image import readpng
from matplotlib.numerix import ravel, where, array, \
     zeros, Float, absolute, nonzero, sqrt
     
debug = False
     
class TexManager:
    """
    Convert strings to dvi files using TeX, caching the results to a
    working dir
    """
    
    oldpath = get_home()
    if oldpath is None: oldpath = get_data_path()
    oldcache = os.path.join(oldpath, '.tex.cache')

    configdir = get_configdir()
    texcache = os.path.join(configdir, 'tex.cache')

    os.environ['TEXMFOUTPUT'] = texcache
    if os.path.exists(oldcache):
        print >> sys.stderr, """\
WARNING: found a TeX cache dir in the deprecated location "%s".
  Moving it to the new default location "%s"."""%(oldcache, texcache)
        shutil.move(oldcache, texcache)

    dvipngVersion = None

    arrayd = {}
    postscriptd = {}
    pscnt = 0

    def __init__(self):
        if not os.path.isdir(self.texcache):
            os.mkdir(self.texcache)

        
    def get_prefix(self, tex):
        return md5.md5(tex).hexdigest()
        
    def get_tex_command(self, tex, fname):

        fontcmd = {'sans-serif' : r'{\sffamily %s}',
                   'monospace'  : r'{\ttfamily %s}'}.get(
            rcParams['font.family'], r'{\rmfamily %s}')
        tex = fontcmd % tex
        fh = file(fname, 'w')
        if rcParams['text.tex.engine'] == 'latex':
            s = r"""\documentclass{article}
\usepackage{%s}
\setlength{\paperwidth}{72in}
\setlength{\paperheight}{72in}
\pagestyle{empty}
\begin{document}
%s
\end{document}
""" % (rcParams['font.latex.package'], tex)
            fh.write(s)
            fh.close()
            command = 'latex -interaction=nonstopmode "%s"'%fname
        else:
            s = r"""\def\frac#1#2{ {#1 \over #2} }
\nopagenumbers
\hsize=72in
\vsize=72in
%s
\bye
""" % tex
            fh.write(s)
            fh.close()
            command = 'tex "%s"'%fname
        return command
        
    def make_dvi(self, tex, force=0):
        if debug: force = True
        
        prefix = self.get_prefix(tex)
        fname = os.path.join(self.texcache, prefix+ '.tex')
        dvibase = prefix + '.dvi'
        dvifile = os.path.join(self.texcache, dvibase)        

        if force or not os.path.exists(dvifile):
            command = self.get_tex_command(tex, fname)
            verbose.report(command, 'debug-annoying')            
            stdin, stdout, stderr = os.popen3(command)
            verbose.report(stdout.read(), 'debug-annoying')
            err = stderr.read()
            if err: verbose.report(err, 'helpful')

        # tex will put it's output in the current dir if possible, and
        # if not in TEXMFOUTPUT.  So check for existence in current
        # dir and move it if necessary and then cleanup
        if os.path.exists(dvibase):
            shutil.move(dvibase, dvifile)
            for fname in glob.glob(prefix+'*'):
                os.remove(fname)
        return dvifile
        
    def make_png(self, tex, dpi, force=0):
        if debug: force = True
        
        dvifile = self.make_dvi(tex)
        prefix = self.get_prefix(tex)
        pngfile = os.path.join(self.texcache, '%s_%d.png'% (prefix, dpi))

        self.get_dvipng_version()  # raises if dvipng is not up-to-date
        #print 'makepng', prefix, dvifile, pngfile
        #command = 'dvipng -bg Transparent -fg "rgb 0.0 0.0 0.0" -D %d -T tight -o "%s" "%s"'% (dpi, pngfile, dvifile)
        command = 'dvipng -bg Transparent -D %d -T tight -o "%s" "%s"'% (dpi, pngfile, dvifile)        

        #assume white bg
        #command = "dvipng -bg 'rgb 1.0 1.0 1.0' -fg 'rgb 0.0 0.0 0.0' -D %d -T tight -o %s %s"% (dpi, pngfile, dvifile)
        # assume gray bg
        #command = "dvipng -bg 'rgb 0.75 0.75 .75' -fg 'rgb 0.0 0.0 0.0' -D %d -T tight -o %s %s"% (dpi, pngfile, dvifile)                

        # see get_rgba for a discussion of the background
        if force or not os.path.exists(pngfile):
            verbose.report(command, 'debug-annoying')
            stdin, stdout, stderr = os.popen3(command)
            verbose.report(stdout.read(), 'debug-annoying')
            err = stderr.read()
            if err: verbose.report(err, 'helpful')
        return pngfile

    def make_ps(self, tex, dpi, force=0):
        if debug: force = True
        
        dvifile = self.make_dvi(tex)
        prefix = self.get_prefix(tex)
        psfile = os.path.join(self.texcache, '%s_%d.epsf'% (prefix, dpi))

        if not os.path.exists(psfile):
            command = 'dvips -q -E -D %d -o "%s" "%s"'% (dpi, psfile, dvifile)
            verbose.report(command, 'debug-annoying')
            stdin, stdout, stderr = os.popen3(command)
            verbose.report(stdout.read(), 'debug-annoying')
            err = stderr.read()
            if err: verbose.report(err, 'helpful')

        return psfile

    def get_ps_bbox(self, tex):
        key = tex
        val = self.postscriptd.get(key)
        if val is not None: return val
        psfile = self.make_ps(tex, dpi=72.27)
        ps = file(psfile).read()
        for line in ps.split('\n'):
            if line.startswith('%%BoundingBox:'):
                return [int(val) for val in line.split()[1:]]
        raise RuntimeError('Could not parse %s'%psfile)
        
        
    def __get_ps(self, tex, fontsize=10, dpi=80, rgb=(0,0,0)):
        """
        Return bbox, header, texps for tex string via make_ps    
        """

        # this is badly broken and safe to ignore.
        key = tex, fontsize, dpi, rgb
        val = self.postscriptd.get(key)
        if val is not None: return val
        psfile = self.make_ps(tex, dpi)
        ps = file(psfile).read()
        
        # parse the ps
        bbox = None
        header = []
        tex = []
        inheader = False
        texon = False
        fonts = []
        infont = False
        replaced = {}
        for line in ps.split('\n'):
            if line.startswith('%%EndProlog'):
                inheader = False
            if line.startswith('%%Trailer'):
                break

            if line.startswith('%%BoundingBox:'):
                bbox = [int(val) for val in line.split()[1:]]
                continue
            if line.startswith('%%BeginFont:'):
                fontname = line.split()[-1].strip()
                newfontname = fontname + str(self.pscnt)
                replaced[fontname] = newfontname
                thisfont = [line]
                infont = True
                continue
            if line.startswith('%%EndFont'):
                thisfont.append('%%EndFont\n')
                fonts.append('\n'.join(thisfont).replace(fontname, newfontname))
                thisfont = []
                infont = False                
                continue
            if infont:
                thisfont.append(line)
                continue
            if line.startswith('%%BeginProcSet:'):                
                inheader = True
            if inheader:
                header.append(line)
            if line.startswith('%%EndSetup'):
                assert(not inheader)
                texon = True
                continue



            if texon:
                line = line.replace('eop end', 'end')
                tex.append(line)

        def clean(s):
            for k,v in replaced.items():
                s = s.replace(k,v)
            return s
        
        header.append('\n')
        tex.append('\n')
        if bbox is None:
            raise RuntimeError('Failed to parse dvips file: %s' % psfile)
        
        replaced['TeXDict'] = 'TeXDict%d'%self.pscnt
        header = clean('\n'.join(header))
        tex = clean('\n'.join(tex))
        fonts = '\n'.join(fonts)
        val = bbox, header, fonts, tex
        self.postscriptd[key] = val

        self.pscnt += 1
        return val
        
    def get_rgba(self, tex, fontsize=10, dpi=80, rgb=(0,0,0)):
        """
        Return tex string as an rgba array
        """

        # dvipng assumes a constant background, whereas we want to
        # overlay these rasters with antialiasing over arbitrary
        # backgrounds that may have other figure elements under them.
        # When you set dvipng -bg Transparent, it actually makes the
        # alpha channel 1 and does the background compositing and
        # antialiasing itself and puts the blended data in the rgb
        # channels.  So what we do is extract the alpha information
        # from the red channel, which is a blend of the default dvipng
        # background (white) and foreground (black).  So the amount of
        # red (or green or blue for that matter since white and black
        # blend to a grayscale) is the alpha intensity.  Once we
        # extract the correct alpha information, we assign it to the
        # alpha channel properly and let the users pick their rgb.  In
        # this way, we can overlay tex strings on arbitrary
        # backgrounds with antialiasing
        #
        # red = alpha*red_foreground + (1-alpha)*red_background

        # Since the foreground is black (0) and the background is
        # white (1) this reduces to red = 1-alpha or alpha = 1-red
        
        # assuming standard 10pt design size space
        dpi = fontsize/10.0 * dpi
        
        r,g,b = rgb
        key = tex, dpi, tuple(rgb)
        Z = self.arrayd.get(key)
        
        if Z is None:
            # force=True to skip cacheing while debugging
            pngfile = self.make_png(tex, dpi, force=False) 
            X = readpng(pngfile)
            vers = self.get_dvipng_version()
            #print 'dvipng version', vers
            if vers<'1.6':
                # hack the alpha channel as described in comment above
                alpha = sqrt(1-X[:,:,0])
            else:
                # dvipng 1.6 and above handles the alpha channel
                # properly [JDH: for some reason I had square root in
                # here which isn't correct
                #alpha = sqrt(X[:,:,-1])
                alpha = X[:,:,-1]
            

            Z = zeros(X.shape, Float)
            Z[:,:,0] = r
            Z[:,:,1] = g
            Z[:,:,2] = b
            Z[:,:,3] = alpha
               
            self.arrayd[key] = Z

        return Z

    def get_dvipng_version(self):
        if self.dvipngVersion is not None: return self.dvipngVersion
        sin, sout = os.popen2('dvipng --version')
        for line in sout.readlines():
            if line.startswith('dvipng '):
                self.dvipngVersion = line.split()[-1]
                return self.dvipngVersion
        raise RuntimeError('Could not obtain dvipng version')
            
