#!/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.
'''
from __future__ import print_function
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 MouseData: #data structure for mouse events on image
  usageStr='unknown mode. available m:move t: transform'
  def __init__(self):
    self.mode=None

  def addTrfPt(self,p1):
    try:
      p0=self.btnPress
      del self.btnPress
    except AttributeError: return
    try:
      trf=self.trfPts
    except AttributeError:
      trf=self.trfPts=[]

    p=p0+p1
    print('add trf: {}'.format(p))
    trf.append(p)


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])
    ax.invert_xaxis();ax.invert_yaxis()
    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.ax=ax
    self.mouseData=MouseData()

  def InitChild(self,doc):

    fig=self.figure
    ax=self.ax
    pts=doc.fh['pts']
    rec=doc.fh['rec']
    idxTrigger=doc.idxTrigger

    # idx=[]
    # for i in range(len(pts)):
    #   l=rec[:,(3,2)]-pts[i,:]
    #   l2=l[:,0]**2+l[:,1]**2
    #   idx.extend(np.where(l2<1)[0].tolist())
    # recPts=rec[idx,:]

    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[:, 3], rec[:, 2], 'b-',label='recDesPos')
    hl+=ax.plot(rec[:,1],rec[:,0],'g-',label='recActPos')
    #hl+=ax.plot(recPts[:,1],recPts[:,0],'g.',label='recDot')
    hl+=ax.plot(rec[-1:,1],rec[-1:,0],'r.--',label='recNxt')
    #hh=ax.plot(rec[idxTrigger,1],rec[idxTrigger,0],'rx',label='trig')
    #hh[0].get_markeredgewidth();hh[0].get_markersize()
    hl+=ax.plot(rec[idxTrigger,1],rec[idxTrigger,0],'rx',label='trig',markeredgewidth=1.,markersize=6.)

    #hl+=ax.plot(recPts[:,1],recPts[:,0],'g.',label='recDot')
    fn=doc.fh.fid.name[:-4]+'.jpg'
    try:
      img=mpl.image.imread(fn)
      try:
        self.imgTrf=trf=doc.fh['imgTrf']
      except KeyError:
        mn=pts.min(0)
        mx=pts.max(0)
        ext=(mn[0], mx[0], mn[1], mx[1])
        skew=None
      else:
        xmin=trf[0, 2]
        ymin=trf[1, 2]
        xmax=trf[0, :].sum()
        ymax=trf[1, :].sum()
        xskew=trf[0, 2]+trf[0, 0]
        yskew=trf[1, 2]+trf[1, 0]
        ext=(xmin, xmax, ymin, ymax)
        skew=(xskew, yskew)

      #interpolation must be none to allow skew image
      himg=ax.imshow(img,extent=ext,interpolation='none')
      himg._image_skew_coordinate=skew

      pass
    except IOError as e:
      print(e)
    ax.xaxis.set_label_text('x-pos um')
    ax.yaxis.set_label_text('y-pos um')
    ax.axis('equal')
    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:
      md=self.mouseData
      if md.mode==2 and event.button==1 and self.toolbar.mode=='': # move
        try:
          himg=self.ax.get_images()[0]
        except IndexError:
          print('no image to transform')
          return
        trf=self.imgTrf
        xOfs=event.xdata-md.btnPress[0]
        yOfs=event.ydata-md.btnPress[1]
        xmin=trf[0, 2]+xOfs
        ymin=trf[1, 2]+yOfs
        xmax=trf[0, :].sum()+xOfs
        ymax=trf[1, :].sum()+yOfs
        xskew=trf[0, 2]+trf[0, 0]+xOfs
        yskew=trf[1, 2]+trf[1, 0]+yOfs
        himg.set_extent((xmin, xmax, ymin, ymax))
        himg._image_skew_coordinate=(xskew, yskew)
        self.figure.canvas.draw()

      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"""
    if self.toolbar.mode!='':
      print(self.toolbar.mode)
      return
    if event.inaxes==self.ax:
      md=self.mouseData
      md.mode
      if md.mode==1:#transform
        if event.guiEvent.RightDown()==True:
          self.doTransform()
        if event.guiEvent.LeftDown()==True:
          p=(event.xdata,event.ydata)
          md.btnPress=p
      elif md.mode==2:#move
        if event.guiEvent.LeftDown()==True:
          p=(event.xdata,event.ydata)
          md.btnPress=p
      else:
        print(md.usageStr)

          #print('self.btnPress1', p)
    #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"""
    print('OnBtnRelease')
    md=self.mouseData
    if md.mode==1:  #transform
      p1=(event.xdata, event.ydata)
      md.addTrfPt(p1)
    elif md.mode==2:#move
      trf=self.imgTrf
      xOfs=event.xdata-md.btnPress[0]
      yOfs=event.ydata-md.btnPress[1]
      #trf[0,2]+=xOfs
      #trf[1,2]+=yOfs
      #trf[1,2]+=yOfs
    return

  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._goImg.set_norm(colBar.norm)#force image to redraw
      colBar.patch.figure.canvas.draw()

  def OnKeyPress(self, event):
    #self.trfPts=[(-1831.2063808574276, 1349.6011964107677, -1966.6001994017945, 1433.3499501495512), (-1830.4810087443084, 917.64916303133487, -1958.3629895701151, 870.93212944552363), (-1561.5041487047886, 917.17727380319548, -1445.8912878106089, 873.29157558622114)]
    #self.trfPts=[(-1960., 872., -1960., 872.), (-1450., 872, -1450., 872), (-1960., 1432.,-1960., 1432.)]
    #self.trfPts=[(-1960., 872., -1760., 872.), (-1450., 872, -1250., 872), (-1960., 1432.,-1760., 1432.)]
    #self.doTransform()
    md=self.mouseData
    if event.key=='t':
      md.mode=1;print('transform image mode (drag drop to add point, right mopuse to execute)')
    elif event.key=='x':
      self.doTransform()
    elif event.key=='m':
      md.mode=2;print('move image mode')
    elif event.key=='s':
      self.saveTransform()
    else:
      try:
        print('unknown mode. available m:move t: transform')
        md.mode=None
      except AttributeError:
        pass
    return


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

  def doTransform(self):
    #print('doTransform')
    try:
      trfPts=self.mouseData.trfPts
      del self.mouseData.trfPts
    except AttributeError:
      print('no transformation points')
      return
    print(trfPts)
    trfPts=np.array(trfPts)

    try:
      himg=self.ax.get_images()[0]
    except IndexError:
      print('no image to transform')
      return
    xmin,xmax,ymin,ymax=himg.get_extent()
    #himg.set_extent((xmin,xmax+10,ymin,ymax))
    skew=himg._image_skew_coordinate
    if skew is None:
      xskew=xmax
      yskew=ymin
    else:
      (xskew, yskew)=skew

    trf0=np.array([[xskew-xmin,xmax-xskew,xmin],
                  [yskew-ymin,ymax-yskew,ymin],
                  [0,0,1]])


    #calculate least square transformation matrix of the transformation points
    n=trfPts.shape[0]
    if n<3:
      print('tot enough points (needs at least 3)')
      return

    M=np.mat(trf0).I
    trfPtsImgI=M*np.vstack((trfPts[:,(0,1)].T,np.ones((1,n))))
    trfPtsImgO=M*np.vstack((trfPts[:,(2,3)].T,np.ones((1,n))))
    trfPtsImg=np.hstack((trfPtsImgI[0:2,:].T,trfPtsImgO[0:2,:].T))

    A=np.zeros((n*2, 6))
    A[0:n, (0, 1)]=trfPtsImg[:,(0,1)]
    A[0:n, 2]=1
    A[n:n*2, (3, 4)]=trfPtsImg[:,(0,1)]
    A[n:n*2, 5]=1
    A=np.mat(A)
    B=trfPtsImg[:,(2,3)].T.reshape(-1, 1)

    # (A'*A)^-1*A'*B
    r=(A.T*A).I*A.T*B
    trf1=np.hstack((r.T, [[0, 0, 1]])).reshape(3, -1)

    trf=np.mat(trf0)*trf1
    xmin=trf[0, 2]
    ymin=trf[1, 2]
    xmax=trf[0, :].sum()
    ymax=trf[1, :].sum()
    xskew=trf[0, 2]+trf[0, 0]
    yskew=trf[1, 2]+trf[1, 0]
    self.imgTrf=trf
    print(trf)
    himg.set_extent((xmin,xmax,ymin,ymax))
    himg._image_skew_coordinate=(xskew, yskew)
    self.figure.canvas.draw()

  def saveTransform(self):
    fhr=self.Parent.doc.fh
    fn_npz=fhr.fid.name
    fn_npzOrig=fn_npz+'.orig'
    d=dict(fhr.items())
    fhr.close()
    if not os.path.exists(fn_npzOrig):
      os.rename(fn_npz, fn_npzOrig)
    d['imgTrf']=self.imgTrf
    np.savez_compressed(fn_npz,**d)
    self.Parent.doc.fh=np.load(fn_npz)

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)

    #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, 3], rec[:idx+1, 2])
      hl[3].set_data(rec[:idx+1, 1], rec[:idx+1, 0])
      hl[4].set_data(rec[idx:idx+10, 1], rec[idx:idx+10, 0])
      #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, 0]-y,rec[idx, 0]+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()