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

# # Tutorial on using OpenPIV with PIV Uncertainty Quantification
# 
# Authors: @lento234, @alexlib
# 
# Documentation: https://pivuq.readthedocs.io/en/latest/

# In[14]:


# If you have not installed pivuq yet:
# !conda activate openpiv 
# !pip install pivuq # or !pip install git+https://github.com/lento234/pivuq


# In[15]:


get_ipython().run_line_magic('load_ext', 'watermark')


# In[16]:


import pivuq
from openpiv import tools, pyprocess, scaling, validation, filters
import numpy as np
import matplotlib.pyplot as plt


# In[17]:


get_ipython().run_line_magic('watermark', '-iv')


# ## Obtain images and ground truth data from pivuq 

# In[18]:


# we can run it from any folder
# path = os.path.dirname(os.path.abspath(__file__))

frame_a  = tools.imread( 'https://github.com/lento234/pivuq/raw/main/examples/data/particledisparity_code_testdata/B00010.tif')
frame_b  = tools.imread( 'https://github.com/lento234/pivuq/raw/main/examples/data/particledisparity_code_testdata/B00011.tif')

image_pair = np.clip(np.array([frame_a, frame_b]), 0, 255).astype('uint8')


# In[19]:


plt.figure(figsize=(12,12))
plt.imshow(np.stack([image_pair[0],0*image_pair[0],image_pair[1]],axis=2)*3)


# ## PIV cross-correlation using OpenPIV

# In[20]:


get_ipython().run_cell_magic('time', '', "\nwindow_size = 24\noverlap = int(window_size * 0.5)\nsearch_area_size = 32\n\nu, v, sig2noise = pyprocess.extended_search_area_piv(\n    frame_a, frame_b, \n    window_size=window_size, \n    overlap=overlap, \n    dt=1,\n    search_area_size=search_area_size,\n    sig2noise_method='peak2peak')\n\n# print(u,v,sig2noise)\n\nx, y = pyprocess.get_coordinates(image_size=frame_a.shape, search_area_size=search_area_size, overlap=overlap)\nu, v, mask = validation.sig2noise_val(u, v, sig2noise, threshold = 1.3)\nu, v, mask = validation.global_val(u, v, (-1000, 2000), (-1000, 1000) )\nu, v = filters.replace_outliers(u, v, method='localmean', max_iter=10, kernel_size=2)\nx, y, u, v = scaling.uniform(x, y, u, v, scaling_factor = 1)#96.52)\n\n# Save\ntools.save(x, y, u, v, mask, 'test_uq.vec')\n")


# In[21]:


fig, ax = plt.subplots(figsize=(5,5))
tools.display_vector_field('test_uq.vec', scale=50, width=0.0035, ax = ax)


# ## Uncertainty quantification
# 
# Based on paper and source code:
# 
# * Sciacchitano, A., Wieneke, B., & Scarano, F. (2013). PIV uncertainty quantification by image matching. Measurement Science and Technology, 24 (4). https://doi.org/10.1088/0957-0233/24/4/045302.
# * http://piv.de/uncertainty/?page_id=221
# 
# 
# The mean of disparity set inside a window is defined as:
# $$
# \mu = \frac{1}{N}\sum_{i\in N} c_i d_i,
# $$
# where $c_i = \sqrt{\Pi(x_i)}$ for $i=1,2,...,N$.
# 
# The standard deviation of disparity set inside a window is defined as:
# $$
# \sigma = \sqrt{\frac{\sum_{i\in N}c_i (d_i - \mu)^2}{\sum_{i\in N}c_i}}.
# $$
# 
# Thus, the instantaneous error (estimate) vector is defined as:
# $$
# \hat{\boldsymbol{\delta}} = \{\hat{\delta}_u,\hat{\delta}_v\} = \sqrt{\boldsymbol{\mu}^2 + \left(\frac{\boldsymbol{\sigma}}{\sqrt{N}}\right)^2}
# $$
# 

# ## Loading PIV vectors

# In[22]:


data = np.loadtxt('test_uq.vec', skiprows=1).T

I, J = np.unique(data[0]).shape[0], np.unique(data[1]).shape[0]
X = np.reshape(data[0], (I, J)) # x-coordinates
Y = np.reshape(data[1], (I, J)) # y-coordinates
U = np.stack((np.reshape(data[2], (I, J)), np.reshape(data[3], (I, J))))
# X_i, Y_i = np.meshgrid(np.arange(I)*overlap, np.arange(J)*overlap)


# In[23]:


fig, axes = plt.subplots(ncols=3, sharex=True, sharey=True, figsize=(15, 5))

for i, (ax, var) in enumerate(zip(axes[:2], ["U", "V"])):
    im = ax.pcolormesh(X, Y, U[i])
    fig.colorbar(im, ax=ax)
    ax.set(title=f"${var}$")

ax = axes[-1]
im = ax.pcolormesh(X, Y, np.linalg.norm(U, axis=0), vmax=3)
fig.colorbar(im, ax=ax)
ax.set(title="Magnitude")

for ax in axes:
    ax.set(xlabel="$i$", ylabel="$j$")


# ## UQ using PIVUQ

# In[24]:


get_ipython().run_cell_magic('time', '', 'X_d, Y_d, delta, N, mu, sigma = pivuq.disparity.sws(\n    image_pair,\n    U,\n    window_size=window_size, # Similar to PIV window size\n    window="gaussian", # Best\n    radius=1,\n    sliding_window_subtraction=True, \n    ROI=[10, 450, 220, 430],\n    velocity_upsample_kind="linear",\n    warp_direction="center", # Depends on original PIV algorithm: "center" is typical\n    warp_order=-1, # Whittaker interpolation\n    warp_nsteps=1, # Depends on original PIV: 1 step is standard\n)\n')


# ### Plot disparity fields

# In[25]:


fig, axes = plt.subplots(ncols=4, sharex=True, sharey=True, figsize=(20, 5))

# Magnitude
ax = axes[0]
im = ax.pcolormesh(X, Y, np.linalg.norm(U, axis=0))
fig.colorbar(im, ax=ax)
ax.set(title="Velocity magnitude")

# Disparity error components
for i, (ax, var) in enumerate(zip(axes[1:3], ["x", "y"])):
    im = ax.contourf(X_d, Y_d, delta[i], np.linspace(-0.5, 0.5, 11))
    fig.colorbar(im, ax=ax)
    ax.set(title=f"$\delta_{var}$")

# Disparity error norm
ax = axes[3]
im = ax.contourf(X_d, Y_d, np.linalg.norm(delta, axis=0), np.linspace(0, 1, 11))
fig.colorbar(im, ax=ax)
ax.set(title="Disparity error norm $|\delta|$");


# ### Plot disparity histogram

# In[26]:


fig, axes = plt.subplots(ncols=3, figsize=(15, 3))

for i, (ax, label) in enumerate(zip(axes[:2], [r"$\delta_x$ (px)", r"$\delta_y$ (px)"])):
    values = delta[i].ravel() # Ignoring *zero*-disparity region
    ax.hist(values[np.abs(values) > 0], bins=100, density=True)
    ax.set(title=label)

ax = axes[-1]
values = np.linalg.norm(delta, axis=0).ravel()
ax.hist(values[np.abs(values) > 0], bins=50, density=True)
ax.set(title="$|\delta|$ (px)");