# -*- coding: iso-8859-1 -*-
"""
    MoinMoin - Matplotlib integration

    example:   {{{#!mpl src=off, run=on, display=inline, klass=mpl, tablestyle="", rowstyle="",style=""
                ...
                ...
                }}}


    IN object contains all input files with filenames as key.
    If you use a file, use the '#! attach(filename)' command and
    reference it in the open statement with 
    f = open(IN['filename'])

    keyword parameters:

    @param display: off|link|inline (default:inline)
    @param columns: arange images in table with No columns (default:1) 
                    if more than 1 image and display is inline
    @param src: on|off (display mpl source, or not, default: off)
    @param run: on|off (if off, execution is disabled, default:on)

    @params klass: class attribute for div
    @params tablestyle: style attribute for table
    @params rowstyle: style attribute for table row
    @params style: style attribute for table cell

    @copyright: 2008 F. Zieher
    @license: GNU GPL, see COPYING for details.
"""

import os, re
import sys
import exceptions
import tempfile
import subprocess
import sha
import pickle

from path import path

from MoinMoin.Page import Page
from MoinMoin import config, wikiutil
from MoinMoin.action import AttachFile
from MoinMoin import log
loggin = log.getLogger(__name__)

# ---------------------------------------------------------------------- 

# some config settings
if 'win' in sys.platform:
    FURL = 'e:/home/zieherf/_ipython/security/ipcontroller-tc.furl'
    PYTHONCMD = 'd:/Python25/python.exe'
else:
    FURL = '/s330/moin/furls/ipcontroller-tc.furl'
    PYTHONCMD = 'python'

Dependencies = ['time']

# ---------------------------------------------------------------------- 

_mplCode = """
# -*- coding: iso-8859-1 -*-
from IPython.kernel import client
import pickle

mmPref = "%(prefix)s"
mmFmt  = "%(fmt)s"
inobj  = %(inobj)s

tc     = client.TaskClient('%(furl)s')
script = '''
%(script)s
'''

# execute script on ipengine and obtain output (pullObjs)
pushObjs = dict(IN=inobj,mmPref=mmPref,mmFmt=mmFmt)
pullObjs = ('out',)
st = client.StringTask(script,clear_before=True,pull=pullObjs,push=pushObjs)
res = tc.run(st,block=True)

if res.failure:
    f = open('%(log)s','w')
    f.write(res.failure.getTraceback())
    f.close()
else:
    f = open('%(out)s','wb')
    pickle.dump(res.results.get('out',{}),f)
    f.close()
"""

MPL_HELP = """
# ------------------------------------------------------------------------------------
# Predefined objects:
#    mpl ... matplotlib object
#    plt ... matplotlib.pyplot object
#    np  ... numpy object
#    IN  ... dictionary of absolute paths of input files keyed by filename
#    mm  ... use mm.imgName(imgno=0) and mm.nextImg() in plt.savefig
#            examples: plt.savefig(mm.imgName(1)), and for all consecutive
#                      plt.savefig(img.nextImg())
# ------------------------------------------------------------------------------------
"""

# ---------------------------------------------------------------------- 

class SandBox:
    """Implement Sandbox for executing python code.
    Sandbox has inbox/outbox directories. inbox is where ipengine
    reads all data from. outbox is where ipengine creates
    output files (i.e. image files from matplotlib).
    """

    fmode = 0666
    dmode = 0777

    def __init__(self,sbpath='',prefix=''):
        """Instantiate sandbox

        @param sbpath: if given, path to sandbox, otherwise a temporary name will be created
        @param prefix: prefix to use when snadbox name is created automatically
        """
        self.sb     = sbpath and path(sbpath) or path(tempfile.NamedTemporaryFile(prefix=prefix).name)
        self.inbox  = self.sb.joinpath('inbox')
        self.outbox = self.sb.joinpath('outbox')
        self.create()
        
    def create(self):
        """Create sandbox
        """
        for d in (self.sb,self.inbox,self.outbox):
            d.mkdir()
            d.chmod(self.dmode)

    def emptyInbox(self):
        """Empty (delete) all files in inbox
        """
        for f in self.inbox.walkfiles():
            f.remove()

    def emptyOutbox(self):
        """Empty (delete) all files in outbox
        """
        for f in self.outbox.walkfiles():
            f.remove()

    def remove(self):
        """Remove sandbox
        """
        self.sb.rmtree()

    def sendTo(self,srcname,dstname=''):
        """Copy file srcname into the inbox. 
        If dstname is not given, srcname is used.
        """
        f = path(srcname)
        if f.isfile() and f.access(os.R_OK):
            dst = dstname and self.inbox.joinpath(dstname) or self.inbox.joinpath(f.basename())
            f.copy(dst)
            dst.chmod(self.fmode)

    @property
    def outboxFiles(self):
        """Return files in the outbox
        """
        return self.outbox.files()

    @property
    def INobj(self):
        IN = {}
        for f in self.inbox.files():
            IN[f.basename()] = f.encode(config.charset).replace('\\','/')
        # skip run.py from IN object
        if 'run.py' in IN:
            del IN['run.py']
        return IN

# ---------------------------------------------------------------------- 

class MplClient:
    """Client interface to ipengines.

    1) Create sandbox
    2) Copy all required input data into inbox
    3) Supply python script file to run --> copy into inbox as run.py
    4) Run run.py with TaskClient and pull output out (dictionary)
    5) Get png files from outbox and attach to page name
    6) Cleanup
    """

    furl = FURL

    def __init__(self,request,pagename,script,fmt='png'):

        self.request    = request
        self.pagename   = pagename
        self.fmt        = fmt
        self.attach_dir = path(AttachFile.getAttachDir(self.request,self.pagename))
        self.sb     = None
        self.script = ''
        self.inbox  = []
        self.imgdata = []
        self.logdata = None
        self.outdata = {}

        # base script   
        self.script = script

        # treat attachments referenced in script and create input list
        self.parseAttachments()
        self.parseIncludes()

    @property
    def imgPrefix(self):
        raw = self.script
        for f in self.inbox:
            raw += str(f.stat().st_mtime)
        self.imgPrefix = 'mpl_' + sha.new(raw).hexdigest() + '_chart'
        return self.imgPrefix

    def addInput(self,infile):
        """Add infile to the set of input files.
        """
        f = path(infile)
        if f.isfile() and f.access(os.R_OK):
            self.inbox.append(f)

    def parseAttachments(self):
        """Collect input files from "#! attach(file-attachment-path)" in self.inbox

        The path can reference an attachment on a differenet page than the
        current. It supports relative path syntax, i.e. "../file.dat".
        The file part can include widcard character that are used for
        globing more than one file.
        
        Don't use quotation marks on the file-path.
        """
        rec = re.compile(r"""(#! *attach\()(?P<fname>[A-Za-z0-9 ._=+*/-]+)(\).*)""")
        for line in self.script.splitlines():
            m = rec.match(line)
            if m:
                # find page and attach_dir where attachments are stored
                aName = m.group('fname')

                # get pageName and attchDir
                fparts = aName.split('/')
                pageName = '/'.join(fparts[0:-1])

                if not pageName:
                    pageName = self.pagename
                else:
                    if pageName.endswith('..'):
                        pageName += '/'
                    pageName = wikiutil.AbsPageName(self.pagename,pageName)
                filePat  = fparts[-1]

                attachDir = AttachFile.getAttachDir(self.request,pageName)

                # attachment found, could be pattern
                for fname in path(attachDir).files(filePat):
                    self.inbox.append(fname)

    def parseIncludes(self):
        """Replace #! include(pagename) with the python src from page.
        """
        rec = re.compile(r"""(#! *include\()(?P<pname>[A-Za-z0-9 ._=+*/-]+)(\).*)""")
        script = self.script.splitlines()

        for lno,line in enumerate(script):
            m = rec.match(line)
            if m:
                # include detected
                pageName = m.group('pname')
                if pageName.endswith('..'):
                    pageName += '/'
                pageName  = wikiutil.AbsPageName(self.pagename,pageName)
                pageTxt = Page(self.request, pageName).get_raw_body()
                pageTxt = '\n'.join(pageTxt.splitlines()[1:])
                script[lno] = pageTxt 

        self.script = '\n'.join(script)

    def run(self):
        """Execute script in sandbox by ipengine
        """

        # make sandbox
        self.sb = SandBox(prefix='mpl_')  
        self.prefix  = str(self.sb.outbox.joinpath('mplplot')).replace('\\','/')
        self.logfile = self.sb.outbox.joinpath('error.log')
        self.outfile = self.sb.outbox.joinpath('mpl.out')

        # make mpl script
        self.mplScript = MplScript(self.script).mplScript        

        # send files (data files and python files)
        for f in self.inbox:
            self.sb.sendTo(f)

        # run
        self.runProcess()

        # read image data
        self.imgdata = []
        for no,imgfile in enumerate(self.sb.outbox.files('mplplot-*.%s'%self.fmt)):
            self.imgdata.append(file(imgfile,'rb').read())

        # read pickled output object
        if self.outfile.exists():
            f = open(self.outfile,'rb')
            self.outdata = pickle.load(f)
            f.close()

        # read logfile
        if self.logfile.exists():
            self.logdata = file(self.logfile,'r').read()

        # finally cleanup
        self.cleanup()

    def runProcess(self):
        """Run python/mpl script on an external ipengine
        Direct use of TaskClient is not possible because of 
        twisted web server used for moin application.
        """

        # build run script in sandbox
        self.pyScript = self.sb.inbox.joinpath('run.py')
        logfile = str(self.logfile).replace('\\','/')
        outfile = str(self.outfile).replace('\\','/')
        f = open(self.pyScript,'w')
        params = dict(furl=self.furl,script=self.mplScript,prefix=self.prefix,
                fmt=self.fmt,log=logfile,out=outfile,inobj=self.sb.INobj)
        f.write( _mplCode % params)
        f.close()

        # execute sandbox script 
        self.failure = None
        p = subprocess.Popen([PYTHONCMD,self.pyScript])
        try:        
            p.wait()
        except exceptions.EnvironmentError, ex:
            os.kill(p.pid, 9)
            self.failure = ex

    def cleanup(self):
        """Cleanup sandbox
        """
        self.sb.remove()

# ---------------------------------------------------------------------- 

class MplScript:

    pre = """
# --- PRE CODE -------------------------------------------
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

class MoinMpl:
    def __init__(self,prefix,fmt='png'):
        self.prefix = prefix
        self.imgno  = 0
        self.fmt    = fmt
    def imgName(self,imgno=0):
        self.imgno = int(imgno)
        return '%s-%d.%s'% (self.prefix,self.imgno,self.fmt)
    def nextImg(self):
        return self.imgName(self.imgno+1)
# ------------------------------------------

mpl.rcdefaults()
mm = MoinMpl(mmPref,mmFmt)
out = {} # output will be handed to caller
# --- END PRE CODE ---------------------------------------

"""
    post = """

# --- POST CODE ------------------------------------------
plt.close()
# --- END POST CODE --------------------------------------
"""

    def __init__(self,pyscript=''):
        self.script = pyscript

    @property
    def mplScript(self):
        return self.pre+self.script+self.post

# ---------------------------------------------------------------------- 

def mpl_settings(run='on',src='off',mxsrc=9999,fmt='png',display='inline',columns=1,klass="mpl",tablestyle="",rowstyle="",style=""):
    """
    Initialize default parameters. The parameters are checked for wrong input.
    @param run: don't execute script of run == 'off' (default on)
    @param src: show src, default=off
    @param mxsrc: number of lines to show, default=9999  all lines basically 
    """
    return locals()

# ---------------------------------------------------------------------- 

class Parser:
    """
        Sends plot images generated by matplotlib
    """
    
    extensions = []
    Dependencies = Dependencies

    def __init__(self, raw, request, **kw):
        self.raw = raw
        self.request = request

        args = kw.get('format_args', '')
        # we use a macro definition to initialize the default init parameters
        # if a user enters a wrong parameter the failure is shown by the exception
        try:
            pyplot = mpl_settings()
            for k,v in pyplot.iteritems():
                setattr(self,k,v)
            settings = wikiutil.invoke_extension_function(request, mpl_settings, args)
            for k, v in settings.iteritems():
                if v: 
                    pyplot[k]=v
                    setattr(self,k,v)

        except ValueError, err:
            msg = u"matplotlib: %s" % err.args[0]
            request.write(self.request.formatter.text(msg))

    def format(self, formatter):
        """ Send the text. """

        self.request.flush() # to identify error text
        self.formatter = formatter

        if self.src.lower() in ('on','1','true'):
            # use highlight parser to show python code
            from highlight import Parser as HLParser
            #pre = MplScript().pre
            #post= MplScript().post
            #hlp = HLParser(pre+self.raw+post,self.request,format_args='python')
            hlp = HLParser(MPL_HELP+self.raw,self.request,format_args='python')
            hlp.format(formatter)

        if self.run.lower() in ('off','0','false'):
            self.request.write(formatter.preformatted(1)+'Execution disabled, set run=on.'+formatter.preformatted(0))
            return

        self.pagename = formatter.page.page_name
        self.attach_dir=AttachFile.getAttachDir(self.request,self.pagename,create=1)

        # -------------------- 

        # mplclient is created here, need to include data files in hexdigest
        mpl = MplClient(self.request,self.pagename,self.raw,self.fmt)
        imgPrefix = mpl.imgPrefix
        update, charts = self._updateImgs(formatter,imgPrefix)

        if update:

            mpl.run()

            if mpl.logdata:
                # error occured
                self.request.write(formatter.preformatted(1)+mpl.logdata+formatter.preformatted(0))
                return

            # attach generated image(s)
            for no,imgdata in enumerate(mpl.imgdata):
                imgName = "%s-%d.%s" % (imgPrefix,no,self.fmt)
                attached_file = file(self.attach_dir + "/" + imgName, 'wb')
                attached_file.write(imgdata)
                attached_file.close()
                charts.append(imgName)

        self.renderCharts(charts)

    def renderCharts(self,charts):
        """Render charts according to settings
        """
    
        if self.display.lower() in ('off','0','false') or not charts:
            return

        fmt = self.formatter

        html = []
        html.append(fmt.div(1,attr={'class':self.klass}))

        if self.display.lower() == 'link':
            # link display
            html.append(fmt.bullet_list(1))
            for chart in charts:
                url = AttachFile.getAttachUrl(self.pagename, chart, self.request)
                html.append(fmt.listitem(1))
                html.append(fmt.url(1,url)+chart+fmt.url(0))
                html.append(fmt.listitem(0))
            html.append(fmt.bullet_list(0))
        
        else:
            # inline display in table form
            noImgs = len(charts)
            rows  = noImgs / self.columns
            rows += noImgs%self.columns and 1 or 0

            T = fmt.table
            R = fmt.table_row
            C = fmt.table_cell

            html.append(T(1,style=self.tablestyle))
            id = 0
            for row in range(rows):
                html.append(R(1,style=self.rowstyle))
                for col in range(self.columns):
                    chart = charts[id]
                    url = AttachFile.getAttachUrl(self.pagename, chart, self.request)
                    html.append(C(1,style=self.style)+fmt.url(1,url)+fmt.image(src="%s" % url, alt=chart)+fmt.url(0)+C(0))
                    id += 1
                    if id >= noImgs:
                        break
                html.append(R(0))
                if id >= noImgs:
                    break
            html.append(T(0))

        html.append(fmt.div(0))
        self.request.write('\n'.join(html))

    def _updateImgs(self, formatter,imgPrefix):
        """Delete outdated charts
        @param formatter: formatter object
        """

        # use delete.me.to.regenerate.images trick from dot.py
        dm2ri = self.attach_dir + '/' + "delete.me.to.regenerate.images"
        charts = []

        updateImgs = True
        deleteImgs = not path(dm2ri).exists()
        if deleteImgs:
            open(dm2ri,'w').close()

        # delete.me. ... exists, if with pref exists, then continue and do not 
        # rerun mpl execution
        attach_files = AttachFile._get_files(self.request, self.pagename)
        reChart = re.compile(r"mpl_.*_chart-[0-9]+.%s"%self.fmt)
        for chart in attach_files:
            if reChart.match(chart) and deleteImgs:
                fullpath = os.path.join(self.attach_dir, chart).encode(config.charset)
                os.remove(fullpath)

            if not deleteImgs and imgPrefix in chart:
                charts.append(chart)
                updateImgs = False

        return updateImgs,charts