#!/usr/bin/env python
# coding: utf-8

# ### Comparison with vectorized and original functions
# #### Edited by Erich Zimmer
# #### Created at 20210817, 2109 CTZ

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
from openpiv.pyprocess import find_first_peak,\
    vectorized_correlation_to_displacements
from openpiv.tools import imread
from numpy import log
from glob import glob


# ### Vectorized solution for subpixel estimation

# In[2]:


N = 64

corr = np.zeros((N,N))

corr[2:5,2:5] = 1
corr[3,3] = 2
corr[3,4] = 3
corr[3,5] = 1
corr


# In[3]:


pos,height = find_first_peak(corr)


# In[4]:


pos,height


# In[5]:


from openpiv.pyprocess import find_subpixel_peak_position


# In[6]:


find_subpixel_peak_position(corr) 


# In[7]:


np.flip(vectorized_correlation_to_displacements(corr[np.newaxis, :, :]) + np.floor(corr.shape[0] / 2))


# ## let's find some corner cases

# In[8]:


# peak on the border 
corr = np.zeros((N,N))

corr[:3,:3] = 1
corr[0,0] = 2
corr[0,2] = 3
corr[0,3] = 1
corr


# ## Corner case 1: peak on the border
# 
# it is disregarded in our function because we cannot define well the subpixel
# position. Or do we? 

# In[9]:


find_subpixel_peak_position(corr)


# In[10]:


np.flip(vectorized_correlation_to_displacements(corr[np.newaxis, :, :]) + np.floor(corr.shape[0] / 2))


# In[11]:


# peak on the border 
corr = np.flipud(corr)
corr


# In[12]:


find_subpixel_peak_position(corr)


# In[13]:


np.flip(vectorized_correlation_to_displacements(corr[np.newaxis, :, :]) + np.floor(corr.shape[0] / 2))


# In[14]:


corr = np.fliplr(corr)
corr[-2,-1]=5
corr


# In[15]:


find_subpixel_peak_position(corr)


# In[16]:


np.flip(vectorized_correlation_to_displacements(corr[np.newaxis, :, :]) + np.floor(corr.shape[0] / 2))


# In[17]:


## Corner case 2: negative value next to peak - the log(n<0) fails


# In[18]:


corr = np.zeros((N,N))

corr[2:5,2:5] = 1
corr[3,3] = 2
corr[3,4] = 3
# corr[3,5] = 1
corr -= 0.5
corr


# In[19]:


find_subpixel_peak_position(corr) # automatically uses parabolic method


# In[20]:


np.flip(vectorized_correlation_to_displacements(corr[np.newaxis, :, :]) + np.floor(corr.shape[0] / 2))


# In[21]:


## Corner case 3: zero next to the peak - the log(0) fails


# In[22]:


corr = np.zeros((N,N))

corr[2:5,2:5] = 1
corr[3,3] = 2
corr[3,4] = 3
# corr[3,5] = 1
corr


# In[23]:


find_subpixel_peak_position(corr)


# In[24]:


np.flip(vectorized_correlation_to_displacements(corr[np.newaxis, :, :]) + np.floor(corr.shape[0] / 2))


# In[25]:


eps = 1e-7
for method in ['gaussian','parabolic','centroid']:
    i,j = find_subpixel_peak_position(corr,method)
    print(i,j)
    i,j = find_subpixel_peak_position(corr+eps,method)
    print(i,j)


# In[26]:


for method in ['gaussian','parabolic','centroid']:
    j, i = vectorized_correlation_to_displacements(corr[np.newaxis, :, :], subpixel_method = method) + np.floor(corr.shape[0] / 2)
    print(i, j)


# ### Speed increase demonstration

# In[27]:


import pylab


# In[28]:


frame_a = imread('../test11/A001_1.tif')
frame_b = imread('../test11/A001_2.tif')
pylab.imshow(np.c_[frame_a,np.ones((frame_a.shape[0],20)),frame_b],
             cmap=pylab.cm.gray)


# In[29]:


window_size = 32
overlap = 16
from openpiv.pyprocess import sliding_window_array, get_field_shape,\
    get_coordinates, fft_correlate_images, correlation_to_displacement
n_rows, n_cols = get_field_shape(
            frame_a.shape, 
            window_size, 
            overlap
        )
x, y = get_coordinates(frame_a.shape, window_size, overlap)

aa = sliding_window_array(
            frame_a, 
            window_size, 
            overlap
)

bb = sliding_window_array(
            frame_b, 
            window_size,
            overlap
)
corr = fft_correlate_images(
            aa, bb,
            correlation_method='circular',
            normalized_correlation = True
)


# In[30]:


from openpiv.pyprocess import find_all_first_peaks, find_all_second_peaks, find_second_peak
peaks_v = find_all_second_peaks(corr)[0]
peaks_o = []
for i in range(len(corr)):
    (k, m), _ = find_second_peak(corr[i,:,:])
    peaks_o.append([i, k, m])
print(['original', 'vectorized'])
for i in range(len(peaks_v)):
    #print(peaks_o[i], peaks_v[i])
    if peaks_v[i][1] != peaks_o[i][1] or peaks_v[i][2] != peaks_o[i][2]:
        print(False)


# In[31]:


get_ipython().run_cell_magic('time', '', "u_o, v_o = correlation_to_displacement(\n            corr, \n            n_rows,\n            n_cols,\n            subpixel_method='gaussian'\n        )\n")


# In[32]:


get_ipython().run_cell_magic('time', '', "u_v, v_v = vectorized_correlation_to_displacements(\n    corr,\n    n_rows,\n    n_cols,\n    subpixel_method='gaussian',\n    #eps = 1e-7\n)\n")


# In[33]:


# slight descrepancies possibly caused by setting eps to 1e-10
print('[u original, u vectorized]')
print(np.stack((u_o[0, 0:12], u_v[0, 0:12])).T)
print((np.nanmean(u_o), np.nanmean(u_v)))


# ### Vectorized solution for signal-to-noise calculation

# In[34]:


from openpiv.pyprocess import vectorized_sig2noise_ratio,\
    sig2noise_ratio


# In[35]:


get_ipython().run_cell_magic('time', '', "peak2peak_o = sig2noise_ratio(corr, 'peak2peak')\n")


# In[36]:


get_ipython().run_cell_magic('time', '', "peak2peak_v = vectorized_sig2noise_ratio(corr, 'peak2peak')\n")


# In[37]:


get_ipython().run_cell_magic('time', '', "peak2mean_o = sig2noise_ratio(corr, 'peak2mean')\n")


# In[38]:


get_ipython().run_cell_magic('time', '', "peak2mean_v = vectorized_sig2noise_ratio(corr, 'peak2mean')\n")


# In[39]:


print('[original, vectorized]')
print(np.stack((peak2peak_o[0:10], peak2peak_v[0:10])).T)
print((peak2peak_o.mean(), peak2peak_v.mean()))


# In[40]:


print('[original, vectorized]')
print(np.stack((peak2mean_o[0:10], peak2mean_v[0:10])).T)
print((peak2mean_o.mean(), peak2mean_v.mean()))


# ## Test for bias errors

# In[41]:


from openpiv.pyprocess import correlation_to_displacement, fft_correlate_images, get_field_shape
files = glob('../test14/*')
files_a = files[::2]
files_b = files[1::2]


# In[44]:


bias_error_original = []
bias_error_vectorized = []
window_size = 32
overlap = 16
real_disp = 3
n = 1/32
for i in range(len(files_a)):
    frame_a = imread(files_a[i])
    frame_b = imread(files_b[i])
    n_rows, n_cols = get_field_shape(
        frame_a.shape, 
        window_size, 
        overlap
    )
    aa = sliding_window_array(frame_a, window_size, overlap)
    bb = sliding_window_array(frame_b, window_size, overlap)
    corr = fft_correlate_images(aa, bb, 'circular', False)
    u_o, v_o = correlation_to_displacement(corr, n_rows, n_cols, 'gaussian')
    u_v, v_v = vectorized_correlation_to_displacements(corr, n_rows, n_cols, 'gaussian')
    u_o = u_o[2:-2, 2:-2] # extract valid components
    u_v = u_v[2:-2, 2:-2]
    v_o = v_o[2:-2, 2:-2] 
    v_v = v_v[2:-2, 2:-2]
    bias_error_original.append(np.hypot(real_disp,real_disp) - np.nanmean(np.hypot(u_o, v_o))) 
    bias_error_vectorized.append(np.hypot(real_disp,real_disp) - np.nanmean(np.hypot(u_v, v_v)))
    real_disp += n


# In[45]:


fig, ax = plt.subplots()
ax.set_ylabel('Bias error [px]')
ax.set_xlabel('Real u and v displacements [px]')
ax.plot(np.mgrid[3:4+n:n], bias_error_original)
ax.plot(np.mgrid[3:4+n:n], bias_error_vectorized)
ax.legend(
    ['orginal', 'vectorized'],
    loc = 'upper right'
)


# In[ ]: