#!/usr/bin/env python
#*-----------------------------------------------------------------------*
#|                                                                       |
#|  Copyright (c) 2018 by Paul Scherrer Institute (http://www.psi.ch)    |
#|                                                                       |
#|              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
#*-----------------------------------------------------------------------*
'''
implements an image alalyser for ESB MX
'''

from scipy import fftpack, ndimage
import scipy.ndimage as ndi
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
#plt.ion()

def ffttest():
  for phi in np.arange(0.,180.,10.):
    ffttest2(phi=phi)
    plt.show()
    pass

def ffttest2(phi=45.,frq=4.2,amp=1.,n=256.):
  #find the main frequency and phase in 1-D
  #plt.ion()
  plt.figure(1)
  d2r=np.pi/180.
  x=np.arange(n)
  phi=phi*d2r
  y=amp*np.cos(frq*x/n*2.*np.pi-phi)
  plt.plot(x,y,'y')
  #y[np.where(y<-.5)]=0
  #y*=np.hamming(n)
  w=np.hanning(n)
  plt.plot(x,w,'y')
  y*=w
  #y*=1.-np.cos(x/(n-1.)*2.*np.pi)

  #y=[1,-1,1,-1,1,-1,1,-1,1,-1,1,-1,1,-1,1,-1]
  plt.stem(x,y)

  fy=np.fft.fft(y)
  fya=np.abs(fy)
  plt.figure(2)
  plt.subplot(211)
  plt.stem(x,fya)
  plt.subplot(212)
  plt.stem(x,np.angle(fy)/d2r)
  print(np.angle(fy[frq])/d2r)

  fya[0]=0
  i=(fya.reshape(2,-1)[0,:]).argmax()
  (vn,v0,vp)=fya[i-1:i+2]
  frq_=i+(vn-vp)/(2.*(vp+vn-2*v0))
  print('freq: %g phase %g %g %g'%((frq_,)+tuple(np.angle(fy[i-1:i+2])/d2r)))

  #PHASE CALCULATION
  plt.figure(1)
  y=np.zeros(x.shape)
  z=np.zeros(x.shape)
  for ii in (i,i-1,i+1):
    y+=np.abs(fy[ii])/n*np.cos(ii*x/n*2.*np.pi+np.angle(fy[ii]))
    z+=np.abs(fy[ii])/n*np.sin(ii*x/n*2.*np.pi+np.angle(fy[ii]))
  y*=2 #double because of conjugate part
  z*=2 #double because of conjugate part
  plt.plot(x,y,'r')
  amp_=fya[i-1:i+2].sum()/n*2.
  t=int(n/2)-1
  #->phase: find maximum or where the sin is 0
  w=np.arccos(y[t]/amp_)
  if(z[t]<0): w=-w
  print('amplitude %g, value at middle (%d) cos->%g sin->%g -> acos %g deg'%(amp_,t,y[t],z[t],w/d2r))
  #rot angle at middle x=127 =frq_*t/n*2.*np.pi-phi_=w '%(amp_,t,y[t],phi_/d2r))
  phi_=frq_*t/n*2.*np.pi-w
  y=amp_*np.cos(frq_*x/n*2.*np.pi-phi_)
  print('y[%d] %g'%(t,y[t]))

  plt.plot(x,y,'g')

  pass

def findGrid(image,numPeak=2):
  d2r=np.pi/180.
  #image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/grid_20180409_115332.png', flatten=True)     # flatten=True gives a greyscale image
  s=image.shape
  w1=np.hamming(s[0]).reshape((-1,1))
  w2=np.hamming(s[1]).reshape((1,-1))
  wnd=w1*w2
  #plt.figure(num='hamming window')
  #plt.imshow(wnd, interpolation="nearest")
  image=wnd*image
  #plt.figure(num='hamming window*img')
  #plt.imshow(image, interpolation="nearest")

  fft2 = np.fft.fft2(image)
  fft2=np.fft.fftshift(fft2)
  fa =abs(fft2)
  plt.figure(num='log of fft: hamming wnd*image')
  hi=plt.imshow(fa, interpolation="nearest",norm=mpl.colors.LogNorm(vmin=.1, vmax=fa.max()))
  #plt.xlim(s[1]/2-50, s[1]/2+50);plt.ylim(s[0]/2-50, s[0]/2+50)

  ctr=np.array(image.shape,dtype=np.int16)/2
  fa[ctr[0] - 1:ctr[0] + 2, ctr[1] - 1:ctr[1] + 2]=0 # set dc to 0
  hi.set_data(fa)
  gen = np.zeros(fft2.shape)
  x=np.arange(s[1])/float(s[1])*2.*np.pi
  y=np.arange(s[0])/float(s[0])*2.*np.pi
  #x=np.linspace(0,2*np.pi,s[1],endpoint=False)
  #y=np.linspace(0,2*np.pi,s[0],endpoint=False)
  my=int(s[0]/2)
  mx=int(s[1]/2)
  xx, yy = np.meshgrid(x, y)

  res=[] #list of tuples (freq_x,freq_y, phase)
  for i in range(numPeak):
    maxAmpIdx=fa.argmax()
    maxAmpPos=np.array(divmod(maxAmpIdx,fa.shape[1]),dtype=np.int16)
    peakPos=maxAmpPos-ctr
    peak=fft2[maxAmpPos[0] - 1:maxAmpPos[0] + 2, maxAmpPos[1] - 1:maxAmpPos[1] + 2]
    print(peakPos)
    print(abs(peak))
    (vn, v0, vp)=np.log(fa[maxAmpPos[0], maxAmpPos[1] - 1:maxAmpPos[1] + 2]) #using log for interpolation is more precise
    freq_x=peakPos[1]+(vn-vp)/(2.*(vp+vn-2*v0))
    (vn, v0, vp)=np.log(fa[maxAmpPos[0]-1:maxAmpPos[0]+2,maxAmpPos[1]])
    freq_y=peakPos[0]+(vn-vp)/(2.*(vp+vn-2*v0))
    #calculate phase
    sumCos=0.
    sumSin=0.
    sumAmp=0.
    n=np.prod(s)
    for iy in (-1,0,1):#(0,):
      for ix in (-1,0,1):
        v=peak[iy+1,ix+1]
        fx=peakPos[1]+ix
        fy=peakPos[0]+iy
        amp=np.abs(v)/n; ang=np.angle(v)
        sumAmp+=amp
        sumCos+=amp*np.cos(fx*x[mx] + fy*y[my] + ang)
        sumSin+=amp*np.sin(fx*x[mx] + fy*y[my] + ang)
        gen+=amp*np.cos(fx*xx + fy*yy + ang)
    sumAmp*=2. #double because of conjugate part
    sumCos*=2.
    sumSin*=2.
    #sumAmp=np.abs(peak).sum()/n*2 #double because of conjugate part
    w=np.arccos(sumCos/sumAmp)
    if sumSin<0: w=-w
    phi_= freq_x*x[mx]+freq_y*y[my]-w
    phi_%=(np.pi*2)
    res.append((freq_x,freq_y,phi_))
    fa[maxAmpPos[0]-1:maxAmpPos[0]+2,maxAmpPos[1]-1:maxAmpPos[1]+2]=0 # clear peak
    maxAmpPos_=2*ctr-maxAmpPos
    fa[maxAmpPos_[0]-1:maxAmpPos_[0]+2,maxAmpPos_[1]-1:maxAmpPos_[1]+2]=0 # clear conjugated peak
    hi.set_data(fa)

  gen*=2. # double because of conjugate part
  for fx,fy,phase in res:
    print('fx: %g fy: %g phase: %g deg'%(fx,fy,phase/d2r))


  plt.xlim(s[1]/2-50, s[1]/2+50)
  plt.ylim(s[0]/2-50, s[0]/2+50)
  plt.figure('image*wnd')
  plt.imshow(image,interpolation="nearest")
  plt.figure('reconstruct')
  plt.imshow(gen,interpolation="nearest")

  plt.figure()
  x=range(s[1])
  y=int(s[0]/2)-1
  plt.plot(x,image[y,:],'r')
  plt.plot(x,gen[y,:],'g')
  pass

def findObj(image,objSize=150,tol=0,viz=0):

  #objSiz is the rough diameter of the searched features in pixels
  #tol      = tolerance in object size (not yet implemented)
  #tolShape = roudness tolerance in object roundness  (not yet implemented)

  from scipy.signal import convolve2d
  #plt.ion()
  s=image.shape
  #box=np.ones((1,objSize*3),dtype=np.float32)/500.

  #img2=ndi.filters.convolve1d(image,box.reshape(-1),0)
  #img2=ndi.filters.convolve1d(img2,box.reshape(-1),1)
  img2=ndi.filters.uniform_filter(np.float32(image),objSize*2)

  if viz&32:
    plt.imshow(image, interpolation="nearest", cmap='gray')
    plt.figure()
    plt.imshow(img2, interpolation="nearest", cmap='gray')
    #plt.show()
  w=np.where(img2>image)
  img2[w]=image[w]
  img3=image-img2
  if viz&16:
    plt.figure()
    plt.imshow(img3, interpolation="nearest", cmap='gray')
    #plt.show()
  l=int(objSize/30)
  if l>0:
    img4=ndi.binary_fill_holes(img3, structure=np.ones((l,l)))
    if viz&8:
      plt.figure()
      plt.imshow(img4, interpolation="nearest", cmap='gray')
      #plt.show()
  else:
    img4=img3

  l=int(objSize/5)#=int(objSize/10)
  if l>=3:
    #img5=ndi.binary_opening(img4, structure=np.ones((l,l)))
    #img5=ndi.binary_erosion(img4, structure=np.ones((l,l)))
    img5=ndi.binary_erosion(img4, iterations=l)
    if viz&4:
      plt.figure()
      plt.imshow(img5, interpolation="nearest", cmap='gray')
      #plt.show()
  else:
    img5=img4

  import cv2
  #cvi1=np.zeros(shape=img5.shape+(3,), dtype=np.uint8)
  #cvi1[:,:,0]=image
  #image*np.array([1,1,1]).reshape(-1,1,1)
  s=image.shape+(1,)
  cvi1=image.reshape(s)*np.ones((1,1,3),dtype=np.uint8)
  #cvi1=np.ones((3,1,1),dtype=np.uint8)image
  contours, hierarchy=cv2.findContours(np.uint8(img5),cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
  #contours,hierarchy=cv2.findContours(np.uint8(img5),1,2)
  cv2.drawContours(cvi1, contours, -1, (0,255,0), 3)
  plt.figure()
  plt.imshow(cvi1 , interpolation="nearest", cmap='gray')
  m=cv2.moments(contours[0])

  lbl = ndi.label(img5)
  if viz&2:
    plt.figure()
    plt.imshow(lbl[0], interpolation="nearest")
    #plt.show()
  ctr=ndi.measurements.center_of_mass(image, lbl[0],range(lbl[1]))
  ctr=np.array(ctr)
  ctr2=np.ndarray(shape=(len(contours),2),dtype=np.uint16)
  i=0
  for c in contours:
    m=cv2.moments(c)
    try:
      m00=m['m00']
      m10=m['m10']
      m01=m['m01']
      print m00
      if m00>1000 and m00<7000:
        ctr2[i,:]=(m10/m00,m01/m00)
        i+=1
    except ZeroDivisionError:
      pass
       #ctr2[i, :]=c[0,0]

  if viz&1:
    plt.figure()
    plt.imshow(image, interpolation="nearest", cmap='gray')
    plt.plot(ctr[:,1],ctr[:,0],'or',markeredgewidth=2, markersize=10)
    plt.plot(ctr2[:i,0],ctr2[:i,1],'+y',markeredgewidth=2, markersize=10)
    plt.show()
  return ctr

def genImg(shape,*args):
  '''args is a list of tuples (freq_x,freq_y, phase) multiple args can be added'''
  image=np.ndarray(shape)#,dtype=np.uint8)
  x=np.linspace(0,2*np.pi,shape[1],endpoint=False)
  y=np.linspace(0,2*np.pi,shape[0],endpoint=False)
  #xx, yy = np.meshgrid(x, y, sparse=True)
  xx, yy = np.meshgrid(x, y)
  dc=0
  for i,f in enumerate(args):
    (freq_x, freq_y, phase)=f
    if i==0:
      image = dc+np.cos(freq_x*xx + freq_y*yy - phase)
    else:
      image += np.cos(freq_x * xx + freq_y * yy - phase)

  plt.imshow(image, interpolation="nearest")
  return image

if __name__ == '__main__':
  #plt.ion()
  #ffttest()
  image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/grid_20180409_115332.png')
  #image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/grid_20180409_115332_45deg.png')
  image=-image
  #image = ndimage.imread('/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/honeycomb.png')

  d2r=np.pi/180.
  #for phi in np.arange(0.,180.,10.):
  #  image = genImg((600, 800), (4.5, .2, phi*d2r))
  #  plt.show()
  #image=genImg((600,800),(4.,1.0,10.*d2r))
  #image=genImg((600,800),(4.5,.2,20.*d2r))
  #findGrid(image,numPeak=1)

  #image=genImg((600,800),(9.5,.2,20.*d2r),(.4,5.2,60.*d2r))
  #image=genImg((600,800),(3.5,0.,0.*d2r),(0,5.,0.*d2r))
  findGrid(image,numPeak=2)

  #image=genImg((600,800),(9.5,.2,0),(.4,5.2,0),(4,8,0))
  #findObj(image,viz=1)
  #findObj(image,viz=255)
  #print(findObj(image))
  plt.show()
