#!/usr/bin/env python
#*-----------------------------------------------------------------------*
#|                                                                       |
#|  Copyright (c) 2013 by Paul Scherrer Institute (http://www.psi.ch)    |
#|                                                                       |
#|              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
#*-----------------------------------------------------------------------*
'''
implements an image view to show a colored image of a hdf5 dataset.
'''

if __name__ == '__main__':
  #Used to guarantee to use at least Wx2.8
  import wxversion
  wxversion.ensureMinimal('2.8')
import wx
import matplotlib as mpl
if __name__ == '__main__':
  mpl.use('WXAgg')
  #or mpl.use('WX')
  #matplotlib.get_backend()

import wxutils as ut
import os
import numpy as np
from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas
import pylab as plt #used for the colormaps

#or from matplotlib.backends.backend_wx import FigureCanvasWx as FigureCanvas
#The source of the DraggableColorbar is from:
#http://www.ster.kuleuven.be/~pieterd/python/html/plotting/interactive_colorbar.html

class MPLCanvasImg(FigureCanvas):
  def __init__(self,parent,SetStatusCB=None):
    if SetStatusCB:
      self.SetStatusCB=SetStatusCB
    fig = mpl.figure.Figure()
    ax = fig.add_axes([0.075,0.075,0.85,0.85])
    FigureCanvas.__init__(self,parent, -1, fig)
    self.mpl_connect('motion_notify_event', self.OnMotion)
    self.mpl_connect('button_press_event',   self.OnBtnPress)
    self.mpl_connect('button_release_event', self.OnBtnRelease)
    self.mpl_connect('scroll_event', self.OnBtnScroll)
    self.mpl_connect('key_press_event',self.OnKeyPress)
    #self.mpl_connect('pick_event',self.OnPick) #works but locks the screed if debugging
    self.fig=fig
    self.ax=ax

  def InitChild(self,pts,rec):

    fig=self.fig
    ax=self.ax

      #res=rot.ActPos,x.ActPos,y.ActPos,rot.DesPos,x.DesPos,y.DesPos
      #idx     0        1        2          3        4        5
    #idx=np.ndarray(shape=len(pts),dtype=np.int32)
    #for i in range(len(pts)):
    #  l=rec[:,4:6]-pts[i,:]
    #  l2=l[:,0]**2+l[:,1]**2
    #  idx[i]=np.argmin(l2)

    idx=[]
    for i in range(len(pts)):
      l=rec[:,4:6]-pts[i,:]
      l2=l[:,0]**2+l[:,1]**2
      idx.extend(np.where(l2<1)[0].tolist())
    #print idx
    recPts=rec[idx,:]
    f2 = plt.figure()
    a2 = f2.add_axes([0.075,0.075,0.85,0.85])
    a2.plot(rec[:,1],'r-',label='ptsDot') #,picker=5 default value
    a2.plot(rec[:,2],'g-',label='ptsDot') #,picker=5 default value
    a2.plot(rec[:,4],'r--',label='ptsDot') #,picker=5 default value
    a2.plot(rec[:,5],'g--',label='ptsDot') #,picker=5 default value
    plt.show()


    hl=[]
    hl+=ax.plot(pts[:,0],pts[:,1],'r.',label='ptsDot') #,picker=5 default value
    hl+=ax.plot(pts[:,0],pts[:,1],'y--',label='ptsLine')


    ec = mpl.collections.EllipseCollection(1, 1, 0, units='xy', offsets=pts,transOffset=ax.transData,edgecolors='g',facecolors=(1,1,0,0.3))
    ax.add_collection(ec)


    hl+=ax.plot(rec[:, 4], rec[:, 5], 'b-',label='recDesPos')
    hl+=ax.plot(rec[:,1],rec[:,2],'g-',label='recActPos')
    hl+=ax.plot(recPts[:,1],recPts[:,2],'g.',label='recDot')
    ax.xaxis.set_label_text('x-pos um')
    ax.yaxis.set_label_text('y-pos um')

    fig.obj=self
    self.ax=ax
    self.hl=hl

  def OnPick(self,event):
    thisline = event.artist
    xdata = thisline.get_xdata()
    ydata = thisline.get_ydata()
    ind = event.ind
    points = tuple(zip(xdata[ind], ydata[ind]))
    print('onpick points:', points)
    pass


  def OnMotion(self,event):
    #print event,event.x,event.y,event.inaxes,event.xdata,event.ydata

    if event.inaxes==self.ax:

      x=int(round(event.xdata))
      y=int(round(event.ydata))
      s=''
      for h in self.ax.get_lines(): # to get also the ellipses: self.ax.get_children()
        c=h.contains(event)
        if c[0]:
          #s = str(h)+str(c[1])
          #print "over %s" % s
          if type(h)==mpl.lines.Line2D:
            lbl=h.get_label()
            arr=h.get_data()
            idx=c[1]['ind'][0]
            s += '%s[%d] = [%.2f, %.2f]'%(lbl,idx,arr[0][idx],arr[1][idx])
          else:
            s += str(h)+str(c[1])

      self.SetStatusCB(self.Parent, 0, (x, y, s))

      #try:
      #  #v=self.img.get_array()[y,x]
      #  v=0;
      #except IndexError as e:
      #  pass
      #else:
        #print x,y,v

    #if event.button==1:#move top,bottom,both
    #  pD = event.y - pS
    #  vD=(vmax-vmin)/(p1-p0)*(pS-event.y)
    #  colBar.norm.vmin = vmin+vD
    #  colBar.norm.vmax = vmax+vD
    #elif event.button==3:#scale around point
    #  scale= np.exp((pS-event.y)/100)
    #  vS=vmin+(vmax-vmin)/(p1-p0)*(pS-p0)
    #  #print scale,vS
    #  colBar.norm.vmin = vS-scale*(vS-vmin)
    #  colBar.norm.vmax = vS-scale*(vS-vmax)

  def OnBtnPress(self, event):
    """on button press we will see if the mouse is over us and store some data"""
    print dir(event.guiEvent)
    return
    if event.inaxes == self.colBar.ax:
      #if event.guiEvent.LeftDClick()==True:
      #    print dlg
      pt=self.colBar.ax.bbox.get_points()[:,1]
      nrm=self.colBar.norm
      self.colBarPressed = (nrm.vmin,nrm.vmax,pt[0],pt[1],event.y)
      #self.colBarPressed = event.x, event.y
      #print self.colBarPressed
      #self.OnMouse(event)
    pass

  def OnBtnRelease(self, event):
    """on release we reset the press data"""
    #self.OnMouse(event)
    return
    try: del self.colBarPressed
    except AttributeError: pass

  def OnBtnScroll(self, event):
    return
    #self.OnMouse(event)
    colBar=self.colBar
    if event.inaxes==colBar.ax:
      pt=colBar.ax.bbox.get_points()[:,1]
      nrm=colBar.norm
      vmin,vmax,p0,p1,pS = (nrm.vmin,nrm.vmax,pt[0],pt[1],event.y)
      if isinstance(colBar.norm,mpl.colors.LogNorm):#type(colBar.norm)==mpl.colors.LogNorm does not work...
        scale= np.exp((-event.step)/10)
        colBar.norm.vmax=vmax*scale
      else:#scale around point
        scale= np.exp((-event.step)/10)
        vS=vmin+(vmax-vmin)/(p1-p0)*(pS-p0)
        #print scale,vS
        colBar.norm.vmin = vS-scale*(vS-vmin)
        colBar.norm.vmax = vS-scale*(vS-vmax)
      self.img.set_norm(colBar.norm)#force image to redraw
      colBar.patch.figure.canvas.draw()

  def OnKeyPress(self, event):
    return
    colCycle=self.colCycle
    colBar=self.colBar
    if event.key=='down':
      self.colIndex += 1
    elif event.key=='up':
      self.colIndex -= 1
    self.colIndex%=len(colCycle)
    cmap = colCycle[self.colIndex]
    colBar.set_cmap(cmap)
    colBar.draw_all()
    self.img.set_cmap(cmap)
    self.img.get_axes().set_title(cmap)
    colBar.patch.figure.canvas.draw()

  def OnMouse(self, event):
    return
    for k in dir(event):
      if k[0]!='_':
        print k,getattr(event,k)

class MAxyPlotFrame(wx.Frame):
  def __del__(self):
    self.doc.view.remove(self)

  def __init__(self, parent,doc):
    wx.Frame.__init__(self, parent, title='xy-Plot', size=wx.Size(850, 650))
    self.doc=doc;doc.view.append(self)
    imgDir=ut.Path.GetImage()
    icon = wx.Icon(os.path.join(imgDir,'PBMA.ico'), wx.BITMAP_TYPE_ICO)
    self.SetIcon(icon)

    canvas =  MPLCanvasImg(self,self.SetStatusCB)

    sizer = wx.BoxSizer(wx.VERTICAL)
    sizer.Add(canvas, 1, wx.LEFT | wx.TOP | wx.GROW)
    self.SetSizer(sizer)

    toolbar=ut.AddToolbar(canvas,sizer)

    wxAxCtrlLst=[]
    #l=len(data.shape)
    #idxXY=(l-2,l-1)
    #for idx,l in enumerate(data.shape):
    #  if idx in idxXY:
    #    continue
    #  wxAxCtrl=ut.SliderGroup(self, label='Axis:%d'%idx,range=(0,l-1))
    #  wxAxCtrl.idx=idx
    #  wxAxCtrlLst.append(wxAxCtrl)
    #  sizer.Add(wxAxCtrl.sizer, 0, wx.EXPAND | wx.ALIGN_CENTER | wx.ALL, border=5)
    #  wxAxCtrl.SetCallback(HdfImageFrame.OnSetView,wxAxCtrl)

    #sl=ut.GetSlice(idxXY,data.shape,wxAxCtrlLst)
    #wxAxCtrl=ut.SliderGroup(self, label='Axis:%d'%1,range=(0,1000))
    #wxAxCtrl.SetCallback(MAxyPlotFrame.OnSetView,wxAxCtrl)
    #sizer.Add(wxAxCtrl.sizer, 0, wx.EXPAND | wx.ALIGN_CENTER | wx.ALL, border=5)

    canvas.InitChild(doc.fh['pts'],doc.fh['rec'])

    #self.Fit()
    self.Centre()

    self.BuildMenu()
    self.canvas=canvas
    self.sizer=sizer
    self.toolbar=toolbar
    #self.data=data
    #self.idxXY=idxXY
    #self.wxAxCtrlLst=wxAxCtrlLst

  def BuildMenu(self):
    mnBar = wx.MenuBar()

    #-------- Edit Menu --------
    mn = wx.Menu()
    #mnItem=mn.Append(wx.ID_ANY, 'Setup Colormap', 'Setup the color mapping ');self.Bind(wx.EVT_MENU, self.OnColmapSetup, mnItem)
    #mnItem=mn.Append(wx.ID_ANY, 'Invert X-Axis', kind=wx.ITEM_CHECK);self.Bind(wx.EVT_MENU, self.OnInvertAxis, mnItem)
    #self.mnIDxAxis=mnItem.GetId()
    #mnItem=mn.Append(wx.ID_ANY, 'Invert Y-Axis', kind=wx.ITEM_CHECK);self.Bind(wx.EVT_MENU, self.OnInvertAxis, mnItem)
    #mnItem=mn.Append(wx.ID_ANY, 'Show Moments', 'Show image moments ', kind=wx.ITEM_CHECK);self.Bind(wx.EVT_MENU, self.OnShowMoments, mnItem)
    #self.mnItemShowMoment=mnItem
    #mnItem=mn.Append(wx.ID_ANY, 'Tomo Normalize', 'Multiplies each pixel with a normalization factor. Assumes there exist an array exchange/data_white', kind=wx.ITEM_CHECK);self.Bind(wx.EVT_MENU, self.OnTomoNormalize, mnItem)
    #self.mnItemTomoNormalize=mnItem

    mnBar.Append(mn, '&Edit')
    #mnItem=mn.Append(wx.ID_ANY, 'Animate', 'Animate the motion');self.Bind(wx.EVT_MENU, self.OnAnimate, mnItem)

    mn = wx.Menu()
    mnItem=mn.Append(wx.ID_ANY, 'Help', 'How to use the image viewer');self.Bind(wx.EVT_MENU, self.OnHelp, mnItem)
    mnBar.Append(mn, '&Help')

    self.SetMenuBar(mnBar)
    self.CreateStatusBar()

  def OnUpdate(self, msg, usrData):
    # this is the model-view-control update function
    #print 'OnUpdate',self, msg, usrData
    if msg==0:
      canvas=self.canvas
      ax=canvas.ax
      idx=usrData
      rec=self.doc.fh['rec']
      hl=canvas.hl
      hl[2].set_data(rec[:idx+1, 4], rec[:idx+1, 5])
      hl[3].set_data(rec[:idx+1, 1], rec[:idx+1, 2])
      #ax.draw_artist(hl[2])
      x=ax.get_xlim();x=(x[1]-x[0])/2;
      y=ax.get_ylim();y=(y[1]-y[0])/2;
      ax.set_xlim(rec[idx, 1]-x,rec[idx, 1]+x)
      ax.set_ylim(rec[idx, 2]-y,rec[idx, 2]+y)
      canvas.draw()

  @staticmethod
  def SetStatusCB(obj,mode,v):
    if mode==0:
      obj.SetStatusText( "x= %d y=%d val=%s"%v,0)
    else:
      raise KeyError('wrong mode')

  def OnHelp(self,event):
    msg='''to change the image selection:
use the toolbar at the bottom to pan and zoom the image
use the scrollbars at the bottom (if present) to select an other slice

to change the colorscale:
drag with left mouse button to move the colorbar up and down
drag with right mouse button to zoom in/out the colorbar at a given point
use mouse weel to zoom in/out the colorbar at a given point
double click left mouse button to set maximum and minimun colorbar values
use cursor up and down to use a different colormap'''
    dlg = wx.MessageDialog(self, msg, 'Help', wx.OK|wx.ICON_INFORMATION)
    dlg.ShowModal()
    dlg.Destroy()


if __name__ == '__main__':
  import os,sys,argparse #since python 2.7
  def GetParser(required=True):
    fnHDF='/scratch/detectorData/e14472_00033.hdf5'
    #lbl='mcs'
    lbl='pilatus_1'
    #lbl='spec'
    elem='/entry/dataScan00033/'+lbl
    exampleCmd='--hdfFile='+fnHDF+' --elem='+elem
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=__doc__,
                                     epilog='Example:\n'+os.path.basename(sys.argv[0])+' '+exampleCmd+'\n ')
    parser.add_argument('--hdfFile', required=required, default=fnHDF, help='the hdf5 to show')
    parser.add_argument('--elem', required=required, default=elem, help='the path to the element in the hdf5 file')
    return parser
    args = parser.parse_args()
    return args

  class App(wx.App):
    def OnInit(self):
      parser=GetParser()
      #parser=GetParser(False) # debug with exampleCmd
      args = parser.parse_args()
      try:
        self.fid=fid=h5py.h5f.open(args.hdfFile)
      except IOError as e:
        sys.stderr.write('Unable to open File: '+args.hdfFile+'\n')
        parser.print_usage(sys.stderr)
        return True
      try:
        hid = h5py.h5o.open(fid,args.elem)
      except KeyError as e:
        sys.stderr.write('Unable to open Object: '+args.elem+'\n')
        parser.print_usage(sys.stderr)
        return True
      frame = HdfImageFrame(None,args.elem,hid)
      frame.Show()
      self.SetTopWindow(frame)
      return True

    def OnExit(self):
      self.fid.close()

  ut.StopWatch.Start()
  app = App()
  app.MainLoop()
