# -*- coding: utf-8 -*-
"""
Created on Mon Jun 03 15:50:30 2013

@author: invernizzi
"""

from PyQt4 import QtGui
import sys
from Ui_MainWindow import Ui_MainWindow
import numpy as np
import libImage
from matplotlib.backends.backend_qt4agg import NavigationToolbar2QTAgg as NavigationToolbar
from matplotlib.widgets import  RectangleSelector
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.animation as animation
info=np.dtype([('xdim','int'),('ydim','int'),('zdim','int'),('t','int'),('xp','float64'),('yp','float64'),('zp','float64'),('f','float64'),('venc','float64'),('slope','float64')])

class MainWindow(QtGui.QMainWindow):
    def __init__(self,parent=None):
        QtGui.QMainWindow.__init__(self)
        #Setup ui
        self.ui=Ui_MainWindow()
        self.ui.setupUi(self)
        self.ui.sBox.setMinimum(0)
        self.ui.bRect.setCheckable(True)
        self.ui.bPixel.setCheckable(True)
        self.ui.bTresh.setCheckable(True)
        self.ui.mplwidgetDown.setHidden(True)
        self.ui.sX.setReadOnly(True)
        self.ui.sY.setReadOnly(True)
        self.ui.sTresh.setReadOnly(True)
        #self.ui.bStop.setHidden(True)
        #Signal and Slot
        self.Connection()
        #class Attribute
        self.index=0
        self.cc=0
        self.f=self.ui.mplWidgetUp.figure
        self.f2=self.ui.mplwidgetDown.figure
        self.ax = self.f.add_subplot(111)
        self.ax2=self.f2.add_subplot(111)
        axcolor = 'lightgoldenrodyellow'
        self.axslide = self.f2.add_axes([0.25, 0.02, 0.35, 0.03], axisbg=axcolor)
        self.f.subplots_adjust(bottom=0.2) 
        self.ui.tab.setEnabled(False)
        self.pause=False
    def Connection(self):
        self.ui.aOpen.triggered.connect(self.FileOpen)
        self.ui.bNext.clicked.connect(self.Next)
        self.ui.bBefore.clicked.connect(self.Before)
        self.ui.sBox.valueChanged.connect(self.GoToFrame)
        self.ui.bRect.clicked.connect(self.Rectangular)
        self.ui.bPixel.clicked.connect(self.Picker)
        self.ui.bTresh.clicked.connect(self.TreshHold)
        self.ui.bShow.clicked.connect(self.Show)
        self.ui.bPlay.clicked.connect(self.Play)
        self.ui.bStop.clicked.connect(self.Stop)

    def Stop(self):
        self.pause=True
    def Play(self):
        self.pause=False
        if(not self.pause): 
            print 'Play tool active'
            self.f2=self.ui.mplwidgetDown.figure
            #self.f2.clear()   
            self.ax2 = self.f2.add_subplot(1,1,1,projection='3d')
            x,y = np.mgrid[:self.matrix_morpho.shape[0],:self.matrix_morpho.shape[1]]
            pl=self.ax2.plot_surface(x,y,self.matrix_morpho[:,:,0],cmap=plt.cm.jet,rstride=1,cstride=1,linewidth=0.,antialiased=False)       
            ani = animation.FuncAnimation(self.f2, self.updateAnim, 
                    frames=xrange(0,12), 
                    fargs=(x,y), interval=30)
            self.ax2.set_title('3D Animation frame=0')
            self.ax2.set_zlim3d(0,100)
            self.ui.mplwidgetDown.draw() 

    def updateAnim(self,i,x2,y2):
         
            if(not self.pause):
                self.ax2.cla()
                wframe = self.ax2.plot_surface(x2, y2,self.matrix_morpho[:,:,i], cmap=plt.cm.jet,rstride=1,cstride=1,linewidth=0.,antialiased=False)
                self.ax2.set_zlim3d(0,100)
                self.ax2.set_title('3D Animation frame=%d'%i)
                return wframe,

              
    def Show(self):
        print 'Show morphed image'
        self.axesClear()
        if(self.ui.bTresh.isChecked()):
            self.ui.bTresh.setChecked(False)
        if(self.ui.bPixel.isChecked):
            self.ui.bPixel.setChecked(False)
        if(self.ui.bRect.isChecked):
            self.ui.bRect.setChecked(False)
        self.matrix_masked=libImage.mask(self.sub_matrix,self.ui.sTresh.value(),self.cc)
        self.matrix_morpho=libImage.morpho(self.matrix_masked)
        xp=1.0;yp=1.0
        x1=np.arange(0,xp*(self.sub_matrix.shape[1]),xp)
        y1=np.arange(0,yp*(self.sub_matrix.shape[0]),yp)
        ### ContourF plot base
        self.f.clear()
        self.ax1=self.f.add_subplot(211)
        cf1=self.ax1.contourf(x1,y1,self.sub_matrix[:,:,self.index])
        self.f.colorbar(cf1)
        self.ax1.grid()
        self.ax1.set_title('matrix velocity filled contours at TS='+str(self.index))
        ### ContourF plot of morph-matrix 
        self.ax2=self.f.add_subplot(212)
        cf3=self.ax2.contourf(x1,y1,self.matrix_morpho[:,:,self.index])
        self.f.colorbar(cf3)
        self.ax2.grid()
        self.ax2.set_title('morph-matrix velocity filled contours masked at TS='+str(self.index))
        self.ui.mplWidgetUp.draw()
        self.f2.clear()
        self.ui.mplwidgetDown.draw()
        
        
    def TreshHold(self):
        print 'TreshHold tool active'
        self.axesClear()
        if(self.ui.bRect.isChecked()):
            self.ui.bRect.setChecked(False)
        if(self.ui.bPixel.isChecked()):
            self.ui.bPixel.setChecked(False)
        c=[]
        self.inc=0.01
        value=np.arange(0,1,0.01)
        self.f.clear()
        self.cc=libImage.cross_corr(self.sub_matrix,self.sub_matrix[self.ui.sX.value(),self.sub_matrix.shape[1]-self.ui.sY.value(),:])
        x=np.arange(len(self.sub_matrix[self.ui.sX.value(),self.ui.sY.value(),:]))
        self.ax=self.f.add_subplot(1,1,1)
        self.ax.plot(x,self.sub_matrix[self.ui.sX.value(),self.sub_matrix.shape[1]-self.ui.sY.value(),:],color='r')
        self.ax.set_title('Reference pixel velocity value across time v(t)')
        self.ui.mplWidgetUp.draw()
        self.f2.clear()
        for v in value:
	    count=0
 	    for i in xrange(self.cc.shape[0]):
			for j in xrange(self.cc.shape[1]):
				if(self.cc[i][j]>=v):
					count+=1
	    c.append(count)
       
        self.ax2=self.f2.add_subplot(1,1,1)
        self.bar=self.ax2.bar(value,c, width=0.02, color='r')
        self.ax2.set_title('Number of pixel vs Thresh_hold',size='12')
        axcolor = 'lightgoldenrodyellow'
        self.axslide = self.f2.add_axes([0.25, 0.02, 0.35, 0.03], axisbg=axcolor)
        self.sslide= libImage.DiscreteSlider(self.axslide, 'TreshHold', 0.0, 1.0,increment=self.inc, valinit=self.inc)
        self.ui.mplwidgetDown.draw()
        self.sslide.on_changed(self.updateBar)
    
    def updateBar(self,val):
            self.ax.hold(True)
            for el in self.bar:
                if el.get_x()==val:
                    el.set_facecolor('g')
                else:
                    el.set_facecolor('r')
            self.ui.sTresh.setValue(val)

    def Picker(self):
        print 'Picker tool active'
        self.axesClear()
        if (not self.ui.mplwidgetDown.isHidden()):
            self.ui.mplwidgetDown.setHidden(True)
        if(self.ui.bRect.isChecked):
            self.ui.bRect.setChecked(False)
        if(self.ui.bTresh.isChecked()):
            self.ui.bTresh.setChecked(False)
        
        self.ax=self.f.add_subplot(1,1,1)
        self.ax.set_picker(5)
        self.ax.set_title("Choose a Central Reference Pixel")
        self.ui.mplWidgetUp.mpl_connect('pick_event', self.onpick)
        self.sub_matrix=self.matrix[self.coord[0]:self.coord[1],self.coord[2]:self.coord[3],:]
        im=self.ax.matshow(np.transpose(self.sub_matrix[:,:,self.index]),picker=True,animated=True)
        self.ax.set_title("Choose a Central Reference Pixel")
        self.ui.mplwidgetDown.setHidden(False)
    
    def onpick(self,event):
            artist = event.artist
            if isinstance(artist, matplotlib.axes.Axes):
                x=event.mouseevent.xdata
                y=event.mouseevent.ydata
                self.ui.sX.setValue(int(y))
                self.ui.sY.setValue(int(x))
                self.update(y,x)
                self.ui.mplWidgetUp.draw()                
    def update(self,x,y):
            x=int(x)
            y=int(y)
           # y=self.sub_matrix.shape[1]-y
            n=1
            self.f2.clear()
            for i in [1,0,-1]:
                for j in (-1,0,1):
                    self.ax2=self.f2.add_subplot(3,3,n)
                    l=np.arange(len(self.sub_matrix[x+j,y+i,:]))
                    self.ax2.plot(l,self.sub_matrix[x+j,y+i,:],color='r')
                    xc=x+j
                    yc=y+i
                    stringa="Pixel "+str(xc)+ " " +str(yc) 
                    xx=self.ax2.set_xlabel('Time Istant')
                    t=self.ax2.set_title(stringa)
                    t.set_size('7')
                    xx.set_size('7')
                    n+=1
                    self.ui.mplwidgetDown.draw()
                    
    def Rectangular(self):
        print 'Rectangular selection tool active'
        self.axesClear()
        if (not self.ui.mplwidgetDown.isHidden()):
            self.ui.mplwidgetDown.setHidden(True)   
        if(self.ui.bTresh.isChecked()):
            self.ui.bTresh.setChecked(False)
        if(self.ui.bPixel.isChecked):
            self.ui.bPixel.setChecked(False)
        self.coord=[0,self.dataInfo['xdim'],0,self.dataInfo['ydim']]
        self.ax=self.f.add_subplot(1,1,1)
        self.ax.matshow(np.transpose(self.matrix[:,:,self.index]),animated=True)
        self.ax.grid()	  
        self.ax.set_title("Choose Rectangular Selection for selected frame "+str(self.index))
        self.ui.mplWidgetUp.draw()
        self.RS = RectangleSelector(self.ax, self.onselect,
                                           drawtype='box', useblit=True,
                                           button=[1,3], # don't use middle button
                                           minspanx=5, minspany=5,
                                           spancoords='pixels')
    def onselect(self,eclick, erelease):
        'eclick and erelease are the press and release events'
        x1, y1 = eclick.xdata, eclick.ydata
        x2, y2 = erelease.xdata, erelease.ydata
        if y1>y2:
            y1,y2=y2,y1
        if x1>x2:
            x1,x2=x2,x1
        self.ax.set_ylim(y2,y1)
        self.ax.set_xlim(x1,x2)
        self.ui.mplWidgetUp.draw()
        self.coord[0]=int(x1)
        self.coord[1]=int(x2)
        self.coord[2]=int(y1) #int(y2)
        self.coord[3]=int(y2) #int(y1)
     
    def GoToFrame(self,i):
        self.axesClear()
        if(not self.ui.mplwidgetDown.isHidden()):
            self.ui.mplwidgetDown.setHidden(True)
        if(self.ui.bTresh.isChecked()):
            self.ui.bTresh.setChecked(False)
        if(self.ui.bPixel.isChecked):
            self.ui.bPixel.setChecked(False)
        if(self.ui.bRect.isChecked):
            self.ui.bRect.setChecked(False)
        self.axesClear()     
        self.index=i
        self.ax=self.f.add_subplot(111)
        self.img=self.ax.imshow(np.transpose(self.matrix[:,:,self.index]))
        self.ax.set_title('Instant %d'%(self.index))
        print 'Frame number %d'%self.index
        self.ui.mplWidgetUp.draw()        
    def Next(self):
        self.axesClear()
        if(not self.ui.mplwidgetDown.isHidden()):
            self.ui.mplwidgetDown.setHidden(True)
        if(self.ui.bTresh.isChecked()):
            self.ui.bTresh.setChecked(False)
        if(self.ui.bPixel.isChecked):
            self.ui.bPixel.setChecked(False)
        if(self.ui.bRect.isChecked):
            self.ui.bRect.setChecked(False)
        
        if(self.index<self.dataInfo['t']-1):
                    self.index += 1
                    self.ax=self.f.add_subplot(111)
                    self.img=self.ax.imshow(np.transpose(self.matrix[:,:,self.index]))
                    self.ax.set_title('Instant %d'%(self.index))
                    self.ui.mplWidgetUp.draw()
                    self.ui.sBox.setValue(self.index)
                    print 'Frame number %d'%self.index
    def Before(self):
        self.axesClear()
        if(not self.ui.mplwidgetDown.isHidden()):
            self.ui.mplwidgetDown.setHidden(True)
        if(self.ui.bTresh.isChecked()):
            self.ui.bTresh.setChecked(False)
        if(self.ui.bPixel.isChecked):
            self.ui.bPixel.setChecked(False)
        if(self.ui.bRect.isChecked):
            self.ui.bRect.setChecked(False)
        if(self.index>0):
                    self.index -= 1
                    self.ax=self.f.add_subplot(111)
                    self.img=self.ax.imshow(np.transpose(self.matrix[:,:,self.index]))
                    self.ax.set_title('Instant %d'%(self.index))
                    self.ui.mplWidgetUp.draw()
                    self.ui.sBox.setValue(self.index)
                    print 'Frame number %d'%self.index
        
    def FileOpen(self):
        print 'Open Input File'
        filename=QtGui.QFileDialog.getOpenFileName(self,'Open Data file','.',"Document files (*.new);;All files (*.*)")
        self.matrix, self.dataInfo=libImage.load_data('data.txt',str(filename))
        self.matrix=self.matrix[:,:,0,:]
        self.matrix=-libImage.sscale(self.matrix,self.dataInfo['slope'],self.dataInfo['venc'] )
        self.ui.sBox.setMaximum(self.dataInfo['t']-1)
        self.ui.tab.setEnabled(True)
        self.GoToFrame(0)
        
    def axesClear(self):
        self.ax.clear()
        self.ax2.clear()
        self.axslide.clear()
        self.ui.mplWidgetUp.draw()
        self.ui.mplwidgetDown.draw()
        self.f.clear()
        self.f2.clear()
        self.f2=self.ui.mplwidgetDown.figure
        self.f=self.ui.mplWidgetUp.figure
        
class OutLog:
    def __init__(self, edit, out=None, color=None):
        
        self.edit = edit
        self.out = None
        self.color = color

    def write(self, m):
        if self.color:
            tc = self.edit.textColor()
            self.edit.setTextColor(self.color)
        self.edit.moveCursor(QtGui.QTextCursor.End)
        self.edit.insertPlainText( m )
        if self.color:
            self.edit.setTextColor(tc)
        if self.out:
            self.out.write(m)
            
if __name__=='__main__':
    app=QtGui.QApplication(sys.argv)
    m=MainWindow()
    m.show()
    sys.stdout = OutLog( m.ui.bOut, sys.stdout)
    app.exec_()
