#!/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
'''
import os
from scipy import fftpack, ndimage
import scipy.ndimage as ndi
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import imgStack

def findGrid(image,numPeak=2,minFrq=2,maxFrq=None,debug=0):
  d2r=np.pi/180.
  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
  if debug&1:
    plt.figure(num='hamming window*img')
    plt.imshow(image, interpolation="nearest")

  fft2 = np.fft.fft2(image)
  fft2=np.fft.fftshift(fft2)
  fa =abs(fft2)
  if debug&1:
    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]-minFrq+1:ctr[0]+minFrq, ctr[1]-minFrq+1:ctr[1]+minFrq]=0 # set dc to 0
  # limit to maximal frequency
  if maxFrq is not None:
    fa[:ctr[0]-maxFrq,:]=0;fa[ctr[0]+maxFrq:,:]=0
    fa[:,:ctr[1]-maxFrq]=0;fa[:,ctr[1]+maxFrq:]=0
  x=np.arange(s[1])/float(s[1])*2.*np.pi
  y=np.arange(s[0])/float(s[0])*2.*np.pi
  if debug&1:
    hi.set_data(fa)
    gen = np.zeros(fft2.shape)
    #x=np.linspace(0,2*np.pi,s[1],endpoint=False)
    #y=np.linspace(0,2*np.pi,s[0],endpoint=False)
    xx, yy = np.meshgrid(x, y)

  my=int(s[0]/2)
  mx=int(s[1]/2)

  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]
    if debug&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+ix] + fy*y[my+iy] + ang)
        sumSin+=amp*np.sin(fx*x[mx+ix] + fy*y[my+iy] + ang)
        if debug&1:
          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)
    #have main frequency positive and phase positive (for convinient)
    if (freq_x<0 and abs(freq_x)> abs(freq_y)) or \
       (freq_y<0 and abs(freq_y)> abs(freq_x)):
      freq_x = -freq_x; freq_y = -freq_y; phi_=-phi_
    if phi_<0: phi_+=2*np.pi

    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
    if debug&1:
      hi.set_data(fa)

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

  if debug&1:
    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')
  return res

def plotGrid(grid,shape):
  #x=np.linspace(0,2*np.pi,shape[1],endpoint=False)
  #y=np.linspace(0,2*np.pi,shape[0],endpoint=False)
  #for (freq_x, freq_y, phase) in grid:
  #find points were: np.cos(freq_x*xx + freq_y*yy - phase) is max
  #freq_x*xx + freq_y*yy - phase = 0 2pi, 4pi
  # grid should have 2 entries
  # -> 2 equations to solve
  # points for:
  # entry1 entry2 = (0 0), (0, 2*pi), (2*pi 0)
  (fx0,fy0,p0)=grid[0]
  (fx1,fy1,p1)=grid[1]
  A=np.array([[fx0,fy0],[fx1,fy1]])
  A*=np.array([2*np.pi/shape[1],2*np.pi/shape[0]])
  Ai=np.asmatrix(A).I


  na=int(max(abs(fx0),abs(fy0)))+1
  nb=int(max(abs(fx1),abs(fy1)))+1
  p=np.ndarray((na*nb,2))
  #p=np.ndarray((3,2))
  #i=0
  #for x,y in ((0,0),(1,0),(0,1)):
  #  v = np.array([p0+x*2.*np.pi, p1+y*2.*np.pi]).reshape(-1, 1)
  #  p[i,:]=(Ai*v).reshape(-1)
  #  i+=1

  i=0
  for b in range(nb):
    for a in range(na):
      v = np.array([p0+a*2.*np.pi, p1+b*2.*np.pi]).reshape(-1, 1)
      p[i,:]=(Ai*v).reshape(-1)
      i+=1




  #plt.plot([400,500],[400,500],'r+',markeredgewidth=2, markersize=10)
  plt.plot(p[:,0],p[:,1],'r+',markeredgewidth=2, markersize=10)
  plt.axis('image')
  pass

def ShowImage(img,title=None,cmap='gray',vmin=None, vmax=None):
  plt.figure(title)
  plt.imshow(img, interpolation="nearest", cmap=cmap,vmin=vmin,vmax=vmax)  # ,vmin=m-3*s, vmax=m+3*s)
  plt.colorbar()

def imgEqualize(img,num_bins=256):
    # get image histogram
    hist, bins = np.histogram(img.flatten(), num_bins, normed=True)
    cdf = hist.cumsum() # cumulative distribution function
    cdf = 255 * cdf / cdf[-1] # normalize

    # use linear interpolation of cdf to find new pixel values
    imgEqu=np.interp(img.flatten(), bins[:-1], cdf)
    imgEqu=np.uint8(imgEqu.reshape(img.shape))

    plt.figure();plt.plot(hist)
    h2,b2=np.histogram(imgEqu.flatten(), num_bins, normed=True)
    plt.figure();plt.plot(h2)
    return imgEqu.reshape(img.shape)#, cdf


def findObj(image,objSize=150,tol=0,debug=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()
  image=image[500:2500,1000:2500]
  s=image.shape
  f=np.array([0.9595264, 0.9600567, 0.9608751, 0.9620137, 0.9634765, 0.9652363, 0.9672352, 0.9693891, 0.9715959, 0.9737464, 0.9757344, 0.9774676, 0.9788761, 0.9799176, 0.9805792, 0.9808776, 0.9808528, 0.9805624, 0.9800734, 0.9794550])
  f=f*1000
  f=f-f.mean()

  m=image.mean();s=image.std()
  ShowImage(image,vmin=-.85,vmax=-.99)
  ShowImage(image,vmin=m-3*s,vmax=m+3*s)
  ieq=imgEqualize(image)
  ShowImage(ieq)
  m=image.mean();s=image.std()
  i2=(image-m)*(256./(3*s))+128.
  i2[i2>255.]=255.
  i2[i2<0.]=0.
  i2=np.uint8(i2)
  ShowImage(i2)
  image=ndi.filters.gaussian_filter1d(image,sigma=5./3,truncate=3.,axis=0)
  image=ndi.filters.gaussian_filter1d(image,sigma=5./3,truncate=3.,axis=1)
  ShowImage(image,vmin=-.85,vmax=-.99)

  img2=ndi.filters.convolve1d(image,f,0)
  ShowImage(img2,vmin=-3,vmax=3)

  img3=ndi.filters.convolve1d(image,f,1)
  ShowImage(img3,vmin=-3,vmax=3)
  img4=np.maximum(abs(img2),abs(img3))
  ShowImage(img4,vmin=0,vmax=3)

  m=image.mean();s=image.std()
  img5=np.zeros(image.shape)
  w=np.where(img4>1.5)
  img5[w]=1
  ShowImage(img5)
  #imgStack.Run([image, img2,img3,img4,img5])
  #w=np.where(img2>image*1.01)
  #img2[:]=0
  #img2[w]=1
  l=int(objSize/30);l=5
  img6=ndi.binary_fill_holes(img5, structure=np.ones((l,l)))

  #img6 = ndi.binary_dilation(img5, iterations=2)
  #img6 = ndi.binary_erosion(img5, iterations=2)
  ShowImage(img6)

  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 debug&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 debug&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 debug&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)
  return image

def phase_retrieval_intensity_tranport(image, mu, delta, I_in, M, pixel_size, R2):
  '''calculates the projected thickness as in eq. 12 of Paganin et al, 2002, J. Microscopy, 
  from an image, the absoption coefficient mu (um-1), the real part of the deviation of the refractive index 
  from unity delta, the uniform intensity of the incident radiation I_in (ph/s/um2), the magnification of the 
  image from the point source illumination M, the pixel size (um), the propagation distance R2 (um).''' 
  F = np.fft.rfft2(image)
  k_x = 1/(pixel_size*image.shape[0]) # is there a 2pi factor?
  k_y = 1/(pixel_size*image.shape[1]) # is there a 2pi factor?
  #print k_x,k_y
  k_array = np.zeros((F.shape[0],F.shape[1]))
  A = np.zeros((F.shape[0],F.shape[1]), dtype=np.complex128)
  for i_k in range(k_array.shape[0]):
    for j_k in range(k_array.shape[1]):
      k_array[i_k,j_k] = np.sqrt((i_k * k_x)**2 + (j_k * k_y)**2)
      A[i_k,j_k] = mu * F[i_k,j_k] / (I_in * (R2 * delta * k_array[i_k,j_k]**2 / M + mu))
  Fm1 = np.fft.irfft2(A)
  T = np.multiply(-1/mu, np.log(Fm1))
  return T

if __name__ == '__main__':
  def testfftLoop():
    plt.ioff()
    for phi in np.arange(0.,180.,10.):
      testfft(phi=phi)
      plt.show()
      pass

  def testfft(phi=45.,frq=4.2,amp=1.,n=256.):
    #find the main frequency and phase in 1-D
    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 testFindGrid():
    plt.ioff()
    imggrid=(
    ('grid_20180409_115332_45deg.png', 2,None),
    ('grid_20180409_115332.png',       2,None),
    ('honeycomb.png',                  2,25),
    ('MS01_20180411_110544.png',       10,None),
    ('MS02_20180411_120354.png',       10,None),
    ('MS03_20180411_135524.png',       10,None),
    ('MS04_20180411_143045.png',       10,None),
    ('MS04_20180411_144239.png',       10,None),
    )
    for (file,minFrq,maxFrq) in imggrid:
      image = ndimage.imread(os.path.join(basePath,file))
      image=-image
      grid=findGrid(image,minFrq=minFrq,maxFrq=maxFrq,numPeak=2)
      plt.figure('grid');plt.imshow(image, interpolation="nearest", cmap='gray')
      plotGrid(grid, image.shape)
      plt.show()

  def testFindObj():
    plt.ioff()
    imggrid=(
    ('grid_20180409_115332_45deg.png', 2,None),
    )
    for (file,p0,p1) in imggrid:
      image = ndimage.imread(os.path.join(basePath,file))
      image=-image
      objPos=findObj(image,debug=0)
      plt.figure('findObj');plt.imshow(image, interpolation="nearest", cmap='gray')
      plt.plot(objPos[:,1],objPos[:,0],'r+',markeredgewidth=2, markersize=10)
      plt.axis('image')
      plt.show()

  def testPhaseEnhance():

    mu=0.001  # um-1, optimal value found is 0.001
    phase_shift=-0.4  # rad/um, optimal value found is -1
    E=18  # keV
    wavelength=12.4/E/10000  # um
    delta=-wavelength/(2*np.pi)*phase_shift
    delta=4.38560287631e-06
    M=1
    R2=19*1e4  # um
    pixel_size=0.325  # um
    flux=1e12  # ph/s
    for fn in ('lyso1_scan_18keV_190mm_MosaicJ_noblend_crop.tif',
               'PepT2_scan_18keV_190mm_MosaicJ_noblend_crop_FFT40.tif',
               'SiN_lysoS1_scan_18keV_190mm_MosaicJ_noblend_crop_FFT40.tif',
               'SiN_PepT2_scan_18keV_190mm_MosaicJ_noblend_crop.tif',):
      img = ndimage.imread(os.path.join(basePath,fn))
      FOV_pix = img.shape
      FOV_um = (FOV_pix[0] * pixel_size, FOV_pix[1] * pixel_size)
      I_in = flux / (FOV_um[0] * FOV_um[1]) # ph/s/um2
      phase_image = phase_retrieval_intensity_tranport(img, mu, delta, I_in, M, pixel_size, R2)
      phase_image_flipped = phase_retrieval_intensity_tranport(np.fliplr(img), mu, delta, I_in, M, pixel_size, R2)
      added_image = np.add(phase_image,np.fliplr(phase_image_flipped))


      m=img.mean();s=img.std()
      ShowImage(img, title=fn+' raw', vmin=m-3*s, vmax=m+3*s)
      m=added_image.mean();s=added_image.std()
      ShowImage(added_image, title=fn+' phase retrival', vmin=m-3*s, vmax=m+3*s)
      plt.show()

  basePath='/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/images/'
  #testfftLoop()
  #testFindGrid()
  testFindObj()
  #testPhaseEnhance()
  exit(0)

  plt.ion()
  #image = ndimage.imread(os.path.join(basePath, 'lyso1_scan_18keV_190mm_MosaicJ_noblend_crop.tif'))
  image = ndimage.imread(os.path.join(basePath, 'lyso1_scan_18keV_190mm_MosaicJ_noblend_crop.tif'))
  #image = ndimage.imread(os.path.join(basePath, 'grid_20180409_115332.png'))
  image = -image
  #plt.figure('input');
  #plt.imshow(image, interpolation="nearest", cmap='gray')
  #sbl=ndi.sobel(image)
  #plt.figure('sobel');
  #plt.imshow(sbl, interpolation="nearest", cmap='gray')
  objPos = findObj(image, objSize=50,debug=255)
  plt.figure('findObj');
  plt.imshow(image, interpolation="nearest", cmap='gray')
  plt.plot(objPos[:, 1], objPos[:, 0], 'r+', markeredgewidth=2, markersize=10)
  plt.axis('image')
  plt.show()

  #d2r=np.pi/180.
  #image=genImg((600,800),(4.,1.0,10.*d2r))
  #image=genImg((600,800),(4.5,-3.2,70.*d2r))
  #image=genImg((600,800),(-4.5,3.2,290.*d2r)) #same image
  #findGrid(image,numPeak=1)
  #findGrid(image,numPeak=1,debug=2)

  #for v in np.arange(0,2,.3):
  #  image=genImg((600,800),(8,.2+v,40.*d2r),(.4,5.2,30.*d2r))
  #  #image = genImg((600, 800), (.4,5.2,30.))
  #  plt.figure(10);plt.cla()
  #  plt.imshow(image)
  #  grid=findGrid(image,numPeak=2,debug=255)
  #  plt.figure(1);plt.cla()
  #  plt.imshow(image, interpolation="nearest", cmap='gray')
  #  plotGrid(grid,image.shape)
  #  plt.show()
  #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))

