# License: Creative Commons Zero (almost public domain) http://scpyce.org/cc0

"""
Example usage of matplotlibs widgets: Build a small 2d Data viewer.

Shows 2d-data as a pcolormesh, a click on the image shows the crossection
(x or y, depending on the mouse button) and draws a corresponding line in
the image, showing the location of the crossections. A reset button deletes all
crossections plots.

Works with matplotlib 1.0.1.
"""
from __future__ import division
from lyse import *
from pylab import *
from analysislib.common.fit_gaussian_2d import fit_2d
from analysislib.spinor.aliases import *
from time import time
from scipy.ndimage import *
from scipy.optimize import *
import os
import pandas as pd
import numpy as np
import numexpr as ne

from matplotlib.widgets import Cursor, Button
import matplotlib.pyplot as plt

plt.rcParams['font.size']=8

class viewer_2d(object):
    def __init__(self,z,x=None, y=None):
        """
        Shows a given array in a 2d-viewer.
        Input: z, an 2d array.
        x,y coordinters are optional.
        """
        if x==None:
            self.x=np.arange(z.shape[0])
        else:
            self.x=x
        if y==None:
            self.y=np.arange(z.shape[1])
        else:
            self.y=y
        self.z=z
        self.fig=plt.figure()
        #Doing some layout with subplots:
        self.fig.subplots_adjust(0.05,0.05,0.98,0.98,0.1)
        self.overview=plt.subplot2grid((8,4),(0,0),rowspan=7,colspan=2)
        self.overview.pcolormesh(self.z, cmap='hot')
        self.overview.autoscale(1,'both',1)
        self.x_subplot=plt.subplot2grid((8,4),(0,2),rowspan=4,colspan=2)
        self.y_subplot=plt.subplot2grid((8,4),(4,2),rowspan=4,colspan=2)

        #Adding widgets, to not be gc'ed, they are put in a list:
        
        cursor=Cursor(self.overview, useblit=True, color='black', linewidth=2 )
        but_ax=plt.subplot2grid((8,4),(7,0),colspan=1)
        reset_button=Button(but_ax,'Reset')
        but_ax2=plt.subplot2grid((8,4),(7,1),colspan=1)
        legend_button=Button(but_ax2,'Legend')
        self._widgets=[cursor,reset_button,legend_button]
        #connect events
        reset_button.on_clicked(self.clear_xy_subplots)
        legend_button.on_clicked(self.show_legend)
        self.fig.canvas.mpl_connect('button_press_event',self.click)
    
    def show_legend(self, event):        
        """Shows legend for the plots"""
        for pl in [self.x_subplot,self.y_subplot]:
            if len(pl.lines)>0:
                pl.legend()
        plt.draw()

    def clear_xy_subplots(self,event):
        """Clears the subplots."""
        for j in [self.overview,self.x_subplot,self.y_subplot]:
            j.lines=[]
            j.legend_ = None
        plt.draw()


    def click(self,event):
        """
        What to do, if a click on the figure happens:
            1. Check which axis
            2. Get data coord's.
            3. Plot resulting data.
            4. Update Figure
        """
        if event.inaxes==self.overview:
            #Get nearest data
            xpos=np.argmin(np.abs(event.xdata-self.x))
            ypos=np.argmin(np.abs(event.ydata-self.y))
            
            #Check which mouse button:
            if event.button==1:
                #Plot it                
                c,=self.y_subplot.plot(self.y, self.z[ypos,:],label=str(self.x[xpos]))
                self.overview.axvline(self.x[xpos],color=c.get_color(),lw=2)

            elif event.button==3:
                #Plot it                
                c,=self.x_subplot.plot(self.x, self.z[:,xpos],label=str(self.y[ypos]))
                self.overview.axhline(self.y[ypos],color=c.get_color(),lw=2)

        if event.inaxes==self.y_subplot:
            ypos=np.argmin(np.abs(event.xdata-self.y))
            c=self.x_subplot.plot(self.x, self.z[ypos,:],label=str(self.y[ypos]))
            self.overview.axhline(self.y[ypos],color=c.get_color(),lw=2)

        if event.inaxes==self.x_subplot:
            xpos=np.argmin(np.abs(event.xdata-self.x))
            c,=self.y_subplot.plot(self.y, self.z[:,xpos],label=str(self.x[xpos]))
            self.overview.axvline(self.x[xpos],color=c.get_color(),lw=2)
        #Show it
        plt.draw()
        
if not spinning_top:
# if the script isn't being run from the lyse GUI, 
# load the DataFrame from lyse and use the first shot
    try:
        df = data()
        path = df['filepath'][0]
    except:
        # if lyse isn't running, use an explicitly specified shot file path
        try:
            path = sys.argv[1]
        except:
            path = '20150325T151931_test_BEC_0_rep00028.h5'

image_group = 'data'
dataset = data(path)
run = Run(path)
    
def raw_to_OD(name):
    with h5py.File(path) as h5_file:
        abs = ma.masked_invalid(array(h5_file['data']['imagesYZ_1_Flea3'][name])[0])
        probe = ma.masked_invalid(array(h5_file['data']['imagesYZ_1_Flea3'][name])[1])
        bckg = ma.masked_invalid(array(h5_file['data']['imagesYZ_1_Flea3'][name])[2])
        return (-log((abs - bckg)/(probe - bckg)))
		
A = raw_to_OD('Raw')
B = np.array(A)
print B
# x=np.linspace(-3,3,300)
xx= np.linspace(0, 487, 488)
# y=np.linspace(-4,4,400)
yy = np.linspace(0, 647, 648)
# X,Y=np.meshgrid(x,y)
# z=np.sqrt(X**2+Y**2)+np.sin(X**2+Y**2)
#Put it in the viewer
# fig_v=viewer_2d(z,x,y)
fig_v2=viewer_2d(B, xx, yy)

#Show it
plt.show()