# Copyright 2009 Robert Schroll, rschroll@gmail.com
# http://jfi.uchicago.edu/~rschroll/refigure/
#
# This file is distributed under the terms of the BSD license, available
# at http://www.opensource.org/licenses/bsd-license.php
#
# Thanks to Owen Taylor for replot:
# http://git.fishsoup.net/cgit/reinteract/tree/lib/replot.py
#
# Thanks to Nicolas Rougier for GTK Python console:
# http://www.loria.fr/~rougier/pycons.html
#
# Thanks to Kai Willadsen for rematplotlib:
# http://ramshacklecode.googlepages.com/#rematplotlib
#
# Thanks to Sage for the syntax:
# http://www.sagemath.org/

"""refigure is an extension for Reinteract that embeds
matplotlib figures in worksheets.

"""
__version__ = "0.3"

import gtk
import copy
import gc, sys

# Monkey patch gtk to keep matplotlib from setting the window icon.
# This can't be an intelligent thing to do....
gtk.window_set_default_icon_from_file = lambda x: None

import matplotlib.pyplot as _p
from matplotlib.figure import Figure
import reinteract.custom_result as custom_result

_p.rcParams.update({'figure.figsize': [6.0, 4.5],
                   'figure.subplot.bottom': 0.12,
                   })
_backend = _p.get_backend()
if _backend[:3] == 'GTK':
    _gui_elements = ['FigureCanvas'+_backend, 'NavigationToolbar2'+_backend]
    if _backend == 'GTKCairo': _gui_elements[1] = 'NavigationToolbar2GTK'
    _temp = __import__('matplotlib.backends.backend_' + _backend.lower(),
                        globals(), locals(), _gui_elements)
    FigureCanvas = getattr(_temp, _gui_elements[0])
    NavigationToolbar = getattr(_temp, _gui_elements[1])
    
    def embedFigure(f):
        c = f.canvas
        box = gtk.VBox()
        box.pack_start(c, True, True)
        toolbar = NavigationToolbar(c, None) # Last is supposed to be window?
        e = gtk.EventBox() # For setting cursor
        e.add(toolbar)
        box.pack_end(e, False, False)
        c.set_size_request(*map(int, f.get_size_inches()*f.get_dpi()))
        box.show_all()
        toolbar.connect("realize", lambda widget:
            widget.window.set_cursor(gtk.gdk.Cursor(gtk.gdk.LEFT_PTR)))
        # Note: backend_gtk.py adds a function notify_axes_change, but this
        # doesn't seem to be needed for everything to work.
        gc.collect()  # This may help, but doesn't eliminate leak
        return box
    
    del(_gui_elements, _temp) #, FigureCanvas, NavigationToolbar)
        
else:
    raise NotImplementedError, """
You must use a GTK-based backend with refigure.  Adjust
your matplotlibrc file, or before importing refigure run
    >>> from matplotlib import use
    >>> use( < 'GTK' | 'GTKAgg' | 'GTKCairo' > )
"""

del(_backend)

class PlotError(Exception):
    pass
    
# CloseError inherits from PlotError so that internals pass it on as if it
# were a PlotError, instead of wrapping in its own PlotError, as with general
# exceptions.
class CloseError(PlotError):
    """Thrown to stop evaluation of a PlotResult object."""
    pass

class setOnceDict(dict):
    """A Dictionary where only the first setting of a key sticks."""
    def __setitem__(self, k, v):
        if not self.has_key(k):
            dict.__setitem__(self, k, v)

class PlotResult(custom_result.CustomResult):

    def __init__(self, command, args=None, kw=None, attrlist=None, figkw=None):
        if isinstance(command, list):
            # Be sure to make a copy!
            self.command_queue = []
            self.command_queue.extend(command)
        elif isinstance(command, SingleCommand):
            self.command_queue = [command]
        else:
            self.command_queue = [SingleCommand(command, args, kw, attrlist)]
        self.figkw = {}
        if figkw is not None:
            self.figkw.update(figkw)
        self.figure = None
    
    def __getattr__(self, name):
        if len(self.command_queue) == 1:
            return PlotResult(getattr(self.command_queue[0], name), figkw=self.figkw)
        else:
            raise NotImplementedError, "Cannot access member of multiple commands."
    
    def __getitem__(self, name):
        if len(self.command_queue) == 1:
            return PlotResult(self.command_queue[0][name], figkw=self.figkw)
        else:
            raise NotImplementedError, "Cannot subscript multiple commands."
    
    def __call__(self, *args, **kw):
        if len(self.command_queue) == 1:
            return PlotResult(self.command_queue[0](*args, **kw), figkw=self.figkw)
        else:
            raise NotImplementedError, "Cannot call multiple commands."
    
    def __add__(self, other):
        if isinstance(other, PlotResult):
            figkw = {}
            figkw.update(self.figkw)
            figkw.update(other.figkw)
            return PlotResult(self.command_queue + other.command_queue, figkw=figkw)
        return NotImplemented
    
    def __iadd__(self, other):
        if isinstance(other, PlotResult):
            self.command_queue.extend(other.command_queue)
            self.figkw.update(other.figkw)
            return self
        return NotImplemented

# This may not be fully kosher, but it should be perfectly safe.  It ought
# to save some copying, as well.
    def __copy__(self):
        return PlotResult(self.command_queue, figkw=self.figkw)
    
    def evaluate(self):
        global _prev_rc
        _prev_rc = setOnceDict()
        f = None
        try:
            for command in self.command_queue:
                f = command.evaluate()
        finally: # We don't trap errors here; we just reset rcParams
            _p.rcParams.update(_prev_rc)
        return f
                
    
    def create_widget(self):
        f = Figure(**self.figkw)
        c = FigureCanvas(f)  # savefig, eg, requires a canvas
        global _current_fig
        _current_fig = f
        
        try:
            self.evaluate()
        except CloseError: # Must be first, else caught by PlotError
#            return None # Doesn't work -> calls None.show() in on_add_custom_result()
            l = gtk.Label("") #<i>Closed figure.</i>")
            #l.set_use_markup(True)
            l.set_size_request(0,0)
            return l
        except PlotError, e:
            return gtk.Label(e)

        return embedFigure(f)
    
    def __repr__(self):
        l = len(self.command_queue)  
        if l == 0:
            return "<PlotResult>: empty"
        elif l == 1:
            return repr(self.command_queue[0])
        else:
            return "<PlotResult>: " + ("\n" + " "*14).join([repr(c) for c in self.command_queue])
   
   
_current_fig = None 
_p.gcf = lambda: _current_fig
def close():
    """Close the current figure and ignore all further commands.  No plot will
    be embedded in the notebook.  You probably want to precede this with a call
    to savefig, if you want to see the resultant figure."""
    raise CloseError
_p.close = close

_prev_rc = None
def rclocal(group, **kwargs):
    """Adjust the rcParams object for only this plot.  Takes arguments either in
    the style of rc() (a group string followed by keyword pairs) or of 
    rcParams.update() (a dictionary)."""
    if isinstance(group, basestring):
        kw = {}
        for k,v in kwargs.items():
            kw[group + '.' + k] = v
    elif isinstance(group, dict):
        kw = group
    else:
        raise TypeError, "The arguments of rclocal must be either a string followed by keyword pairs or a dictionary."
    for k in kw.keys():
        _prev_rc[k] = _p.rcParams[k]
    _p.rcParams.update(kw)
_p.rclocal = rclocal

class SingleCommand(object):
    
    def __init__(self, command, args=None, kw=None, queue=None):
        self.command = command
        self.args = args or []
        self.kw = kw or {}
        self.queue = []
        if queue is not None:
            self.queue.extend(queue)
    
    def __getattr__(self, name):
        return SingleCommand(self.command, self.args, self.kw, self.queue + [name])
    
    def __getitem__(self, name):
        return SingleCommand(self.command, self.args, self.kw, self.queue + [[name]])
    
    def __call__(self, *args, **kw):
        return SingleCommand(self.command, self.args, self.kw, self.queue + [(args, kw)])
    
    def __add__(self, other):
        if isinstance(other, PlotResult):
            return PlotResult(self) + other
        return NotImplemented
    
    def __radd__(self, other):
        if isinstance(other, PlotResult):
            return other + PlotResult(self)
        return NotImplemented

    def __copy__(self):
        return SingleCommand(copy.copy(self.command),
                             copy.copy(self.args),
                             copy.copy(self.kw),
                             copy.copy(self.queue))
    
    def evaluate(self):
        args = [evaluateCommand(a) for a in self.args]
        kw = dict([(k, evaluateCommand(v)) for k,v in self.kw.items()])
        try:
            f = eval('_p.' + self.command + '(*args, **kw)')
            if self.queue is not None:
                for attr in self.queue:
                    if isinstance(attr, tuple):
                        args = [evaluateCommand(a) for a in attr[0]]
                        kw = dict([(k, evaluateCommand(v)) for k,v in attr[1].items()])
                        f = f(*args, **kw)
                    elif isinstance(attr, list):
                        f = f[attr[0]]
                    else:
                        f = getattr(f, attr)
        except PlotError, e:
            raise
        except Exception, e:
            raise PlotError, '%s in `%s`: %s'%(repr(type(e))[18:-2], repr(self), e)
        return f
    
    def __repr__(self):
        s = self.command + functionCallAsString(self.args, self.kw)
        for attr in self.queue:
            if isinstance(attr, tuple):
                s += functionCallAsString(*attr)
            elif isinstance(attr, list):
                s += '[%s]'%repr(attr[0])
            else:
                s += '.' + attr
        return s

def evaluateCommand(obj):
    if isinstance(obj, SingleCommand) or isinstance(obj, PlotResult):
        return obj.evaluate()
    else:
        return obj

def functionCallAsString(args, kw):
    return '(' + ', '.join([repr(a) for a in (args or [])] 
                 + [k + '=' + repr(v) for k,v in (kw or {}).items()]) + ')'

_PassThrough = ['cm',                   # Contains all colormaps
                'colorbar_doc',         # Docstring?
                'colormaps',            # Do-nothing function with doc
                'colors',               # Do-nothing function with doc
                'get_backend',          # Returns string of current backend
                'get_cmap',             # Returns a colormap instance
                'get_plot_commands',    # Returns tuple of plotting commands
                'get_scale_docs',       # Return info on scaling options.
                'get_scale_names',      #  Probably don't need, but doesn't hurt
                'hold',                 # Since these affect global state,
                'ioff',                 #  having them be PlotResults would be
                'ion',                  #  confusing.
                'imread',               # Read in images
                'isinteractive',        # Report interactive state
                'is_numlike',           # Might be useful?
                'is_string_like',       # Ditto
                'normalize',            # Could come in handy, I guess
                'plotting',             # Do-nothing function with doc
                'rc',                   # All the rc stuff should work fine
                'rcdefaults',           #  without modification
                'rcParams',             #  |
                'rcParamsDefault',      #  |
               ]
_FigureCommands = ['figure']
_ExcludedCommands = ['dedent',              # Don't really need this
                     'draw_if_interactive', # Called automatically - don't need
                     'get_current_fig_manager', # Bound to cause problems
                     'ginput',              # Doesn't work, might be fixable
                     'matplotlib',          # Unnecessary (I hope)
                     'matshow',             # Broken for me
                     'mlab',    # Math functions - import explicitly if needed
                     'new_figure_manager',  # Also bound to cause problems
                     'over',                # Doesn't work, prob. not needed
                     'pylab_setup',         # Should be unecessary
                     'silent_list',         # No need for users to have
                     'sys',                 # Import it yourself!
                     'waitforbuttonpress',  # Why would you want this?
                     'x',                   # What is this!?!
                    ] + _FigureCommands + _PassThrough

for command in dir(_p):
    if (command[0].islower()    # Only commands that start with lowercase
            and command not in _ExcludedCommands):
        exec("def %s(*args, **kw):\n '''%s'''\n return PlotResult('%s', args, kw)"%(command, getattr(_p, command).__doc__, command))

for command in _PassThrough:
    exec("%s = _p.%s"%(command, command))

del(_PassThrough, _FigureCommands, _ExcludedCommands)

def figure(**kw):
    """Returns a PlotResult object with no commands.  Any keywords will be
    passed to the matplotlib figure when it is created.
    
    Optional keyword arguments:"""
    return PlotResult([], figkw=kw)
figure.__doc__ += _p.figure.__doc__.split("Optional keyword arguments:", 1)[-1]



# Bug - Hovering over a PlotResult with a single command produces problems.
# This seems to be because pydoc uses inspect.isclass(), which thinks custom
# objects with a __getattr__ are classes (http://bugs.python.org/issue1225107).
# PlotResult has a working __getattr__ only when it has a single command.