#!/usr/bin/env python

import pylab as p
import matplotlib.patches as mpa

DEBUG = True

class GrenzeRechteck:
    "wrapper around matplotlib.patches.Rectangle making the rectangle draggable"
    lock = None
    def __init__(self, axes, parent, xy, width, height):
        "initialize a Rectangle instance"
        self.parent = parent
        #self.rect = rect
        self.press = None
        self.press1 = None
        self.background = None
        self.motion = None # Valid Values: "Move", "Resize"

        self.rect = mpa.Rectangle(xy, width, height, alpha=.2)
        axes.add_patch(self.rect)

    def connect(self):
        'connect to all the events we need'
        self.cidpress = self.rect.figure.canvas.mpl_connect(
            'button_press_event', self.on_press)
        self.cidrelease = self.rect.figure.canvas.mpl_connect(
            'button_release_event', self.on_release)
        self.cidmotion = self.rect.figure.canvas.mpl_connect(
            'motion_notify_event', self.on_motion)
        self.cidresize = self.rect.figure.canvas.mpl_connect(
            'motion_notify_event', self.on_resize)

    def on_press(self, event):
        'on button press we will see if the mouse is over us and store some data'
        if event.inaxes != self.rect.axes: return
        if GrenzeRechteck.lock is not None: return
        contains, attrd = self.rect.contains(event)
        if not contains: return
        if event.button == 1:
            self.motion = 'Move'
        elif event.button == 3:
            self.motion = 'Resize'
        else:
            return
        x0, y0 = self.rect.xy
        w, h = self.rect.get_width(), self.rect.get_height()
        self.press = x0, y0, event.xdata, event.ydata
        self.press1 = x0, y0, event.xdata, event.ydata, w, h
        GrenzeRechteck.lock = self

        # draw everything but the selected rectangle and store the pixel buffer
        canvas = self.rect.figure.canvas
        axes = self.rect.axes
        self.rect.set_animated(True)
        canvas.draw()
        self.background = canvas.copy_from_bbox(self.rect.axes.bbox)

        # now redraw just the rectangle
        axes.draw_artist(self.rect)

        # and blit just the redrawn area
        canvas.blit(axes.bbox)
    
    def on_motion(self, event):
        'on motion we will move the rect if the mouse is over us'
        if not self.motion == 'Move': return
        if GrenzeRechteck.lock is not self:
            return
        if event.inaxes != self.rect.axes: return
        x0, y0, xpress, ypress = self.press
        dx = event.xdata - xpress

        self.rect.set_x(x0 + dx)

        canvas = self.rect.figure.canvas
        axes = self.rect.axes
        canvas.restore_region(self.background)

        # redraw just the current rectangle
        axes.draw_artist(self.rect)

        # blit just the redrawn area
        canvas.blit(axes.bbox)

    def on_resize(self, event):
        'on motion we will resize the rect if the mouse is over us'
        if not self.motion == 'Resize': return
        if GrenzeRechteck.lock is not self:
            return
        if event.inaxes != self.rect.axes: return
        x0, y0, xpress, ypress, w, h = self.press1
        dx = event.xdata - xpress
        self.rect.set_width(dx + w)

        canvas = self.rect.figure.canvas
        axes = self.rect.axes
        # restore the background region
        canvas.restore_region(self.background)

        # redraw just the current rectangle
        axes.draw_artist(self.rect)

        # blit just the redrawn area
        canvas.blit(axes.bbox)

    def on_release(self, event):
        'on release we reset the press data'
        if GrenzeRechteck.lock is not self:
            return

        w = self.rect.get_width()
        xy = self.rect.get_xy()
        lower_limit = p.argmin(abs(self.parent.x-xy[0]))
        upper_limit = p.argmin(abs(self.parent.x-xy[0]-w))
        if DEBUG:
            print xy[0], w
            print lower_limit, upper_limit

        self.press = None
        self.press1 = None
        self.motion = None
        GrenzeRechteck.lock = None

        self.rect.axes.cla()    
        self.rect.axes.plot(self.parent.x,self.parent.y)
        dy = self.parent.get_integral(lower_limit,upper_limit)
        self.rect.axes.plot(self.parent.x[lower_limit:upper_limit],dy)
        self.rect.axes.add_patch(self.rect)
        
        # turn off the rect animation property and reset the background
        self.rect.set_animated(False)
        self.background = None

        # redraw the full figure
        self.rect.figure.canvas.draw()      

    def disconnect(self):
        'disconnect all the stored connection ids'
        self.rect.figure.canvas.mpl_disconnect(self.cidpress)
        self.rect.figure.canvas.mpl_disconnect(self.cidrelease)
        self.rect.figure.canvas.mpl_disconnect(self.cidmotion)
        self.rect.figure.canvas.mpl_disconnect(self.cidresize)
        if self.cidhit:
            self.rect.figure.canvas.mpl_disconnect(self.cidhit)

class Integral:
    x = None
    dx = None
    y = None

    def __init__(self,**kwargs):
        self.x = kwargs.get('x')
        self.y = kwargs.get('y')
        self.dx = p.zeros(len(self.x),dtype=self.x.dtype)
        for i in range(1,len(self.x)):
            self.dx[i] = self.x[i]-self.x[i-1]

    def get_integral(self,lower_limit=None,upper_limit=None):
        if not lower_limit: lower_limit = 0
        if not upper_limit: upper_limit = len(self.y)

        tmp = self.y*self.dx
        Y = p.zeros(len(tmp),dtype=tmp.dtype)
        for i in range(lower_limit+1,upper_limit):
            Y[i] = p.sum(tmp[lower_limit:i])
        return Y[lower_limit:upper_limit]

class Ableitung:
    x = None
    y = None
    z = None

    def __init__(self,**kwargs):
        self.x = kwargs.get('x')
        self.y = kwargs.get('y')
        self.z = kwargs.get('z')

    def get_ableitung(self):
        if not (self.x.ndim == 1 and self.y.ndim == 1):
            raise ValueError, 'This function is for 1 dimensional data only'
        l = len(self.y)
        dy = p.zeros(l,dtype=self.y.dtype)
        for i in range(l):
            a,b = (i+1)%l, (i-1)%l
            if i > 0 and i < l-1:
                dx = self.x[a]-self.x[b]
            elif i == 0:
                dx = 2*(self.x[1]-self.x[0])
            elif i == l-1:
                dx = 2*(self.x[l-1]-self.x[l-2])
            else: raise IndexError, 'Index out of range: %i' % i

            dy[i] = (self.y[a]-self.y[b])/dx
        return dy

    def get_gradient(self):
        if not (self.x.ndim == 2 and self.y.ndim == 2):
            raise ValueError, 'This function is for 2 dimensional data only'
        L = self.z.shape
        dz = p.zeros(L+(2,),dtype=self.y.dtype)
        for l in range(L[0]):
            a,b = (l-1)%L[0],(l+1)%L[0]
            for i in range(L[1]):
                if l > 0:
                    dx = self.y[b,i]-self.y[a,i]
                elif l == 0:
                    dx = 2*(self.y[1,i]-self.y[0,i])
                elif l == L[0]-1:
                    dx = 2*(self.y[L[0]-1,i]-self.y[L[0]-2,i])
                else:
                    raise IndexError, 'Index out of range: %i' % l
                c,d = (i-1)%L[1],(i+1)%L[1]
                if i > 0 and i < L[1]-1:
                    dy = self.x[l,d]-self.x[l,c]
                elif i == 0:
                    dy = 2*(self.x[l,1]-self.x[l,0])
                elif i == L[1]-1:
                    dy = 2*(self.x[l,L[1]-1]-self.x[l,L[1]-2])
                else: raise IndexError, 'Index out of range: %i' % i
                
                dz[l,i,1] = (self.z[b,i]-self.z[a,i])/dx
                dz[l,i,0] = (self.z[l,d]-self.z[l,c])/dy
        return dz

if __name__ == '__main__':
    x = p.arange(-2*p.pi,2*p.pi,.1)
    y = p.sin(x)

    I = Integral(x=x,y=y)
    A = Ableitung(x=x,y=y)
    dy = A.get_ableitung()
    F = p.figure()
    p.plot(x,y,x,dy)
    l = GrenzeRechteck(
        F.gca(), I, (x[10],y.min()), x[-10]-x[10], y.max()-y.min()
    )
    l.connect()

    X,Y = p.meshgrid(x,x)
    Z = p.cos(X+.2*Y**2)
    A = Ableitung(x=X,y=Y,z=Z)
    z = A.get_gradient()
    p.figure()
    p.pcolormesh(X,Y,Z)
    p.colorbar()
    p.quiver(X[::3,::3],Y[::3,::3],
        z[::3,::3,0],z[::3,::3,1],pivot='mid',color='r')
    p.show()