# -*- coding: utf-8 -*-
"""
Created on Tue Jan  6 15:53:52 2026

@author: Jiawen Jiang

This script demonstrates how bootstrapping works at measuring the uncertainty
in PIV measurements. It performs PIV at a defined position on a set of selected
frame pairs (at the same phase of heartbeat, on the same imaging plane) and 
shows the bootstrapped distribution of the measurements.

"""
import matplotlib.pyplot as plt
import numpy as np
from functions import *
import glob
from config import *

#loading image filepaths and plists
stack1gatafp = glob.glob(stack1gata_image)
stack2gatafp = glob.glob(stack2gata_image)
stack1gataplist = load_plist(glob.glob(stack1gata_plist))
stack2gataplist = load_plist(glob.glob(stack2gata_plist))

stack1gfpfp = glob.glob(stack1gfp_image)
stack2gfpfp = glob.glob(stack2gfp_image)
stack1gfpplist = load_plist(glob.glob(stack1gfp_plist))
stack2gfpplist = load_plist(glob.glob(stack2gfp_plist))


#%% get z positions of all planes
stack1z = get_z_pos(stack1gfpplist)
stack2z = get_z_pos(stack2gfpplist)
#%% get gfp images of the heart wall
stack1wall = closest_wall(stack1gfpfp, stack1gfpplist, stack1z[15], 0.1, 0.1)
stack2wall = closest_wall(stack2gfpfp, stack2gfpplist, stack2z[15], 0.1, plane_phase)

#%% selecing frame pairs for PIV
first1, second1 = select_image(stack1gatafp, stack1gataplist, stack1z[15], 0.1, 0.1, 0.05)
#first2, second2 = select_image(stack2gatafp, stack2gataplist, stack2z[15], 0.1, plane_phase, 0.05)
#select from stack2 if desired.

#%% performing bootstrapping, note this is a separate function, as the updated
#PIV function already does bootstrapping together. But they do the same thing.
#Also as bootstrapping is a random process, re-running would yield slightly 
#different results, but the statistics should be unchanged.
from scipy.stats import norm
#bootstrap, getting distribution in x and y displacements
u_sub, v_sub = bootstrap(first1, second1, 100, 150, [48,48], [24,24], iteration=50000)
mu_sub, sigma_sub = norm.fit(v_sub) #fit a normal distribution to x and y
mu_usub, sigma_usub = norm.fit(u_sub)
x_sub = np.linspace(np.min(v_sub), np.max(v_sub), 100) #generating datapoints to plot normal line
xu_sub = np.linspace(np.min(u_sub), np.max(u_sub), 100)
pdf_sub = norm.pdf(x_sub, mu_sub, sigma_sub) #generating normal distribution
pdf_usub = norm.pdf(xu_sub,mu_usub, sigma_usub)

#%% plot the distribution of measured x and y displacements
#accompanied by an image showing the position of the PIV interrogation window
import matplotlib.patches as patches
import matplotlib.lines as mlines
fig, ax = plt.subplots(1,2, figsize=(13,6), dpi=300)
rect = patches.Rectangle(
    (100-12, 150-12),      # bottom‑left corner
    24, 24, # box size
    linewidth=1,
    edgecolor='y',
    facecolor='none',
    label = 'I.W.'# transparent fill
)
dotted_rect = patches.Rectangle(
    (100-24, 150-24),      # bottom‑left corner
    48, 48, # box size
    linewidth=1,
    edgecolor='y',
    facecolor='none',
    linestyle='dotted',
    label = 'S.R.'# transparent fill
)
ax[0].hist(v_sub, bins=100, density=True, alpha=0.5)
ax[0].hist(u_sub, bins=100, density=True, alpha=0.5)
ax[0].plot(xu_sub, pdf_usub, color='tab:orange', label='x')
ax[0].set_xlabel('displacement/pixel')
ax[0].set_ylabel('density')
#ax[0].set_title('bootstrapped '+f"mean {mu_sub:.3f}, std {sigma_sub:.3f}")
ax[0].plot(x_sub, pdf_sub, color='tab:blue', label='y')
ax[0].legend()

ax[1].imshow(grey2rgb(np.mean(first1, axis=0), 'r') + grey2rgb(stack1wall, 'g')**0.6)
ax[1].add_patch(rect)
ax[1].add_patch(dotted_rect)
#adding scale bar
SCALEBAR_LENGTH_PX = 46   # length in image pixels
SCALEBAR_LABEL = "20 µm"
img_h, img_w = stack1wall.shape
margin = 10 #distance to edge of image
sb_x0 = img_w - margin - SCALEBAR_LENGTH_PX
sb_x1 = img_w - margin
sb_y  = img_h - margin

ax[1].plot([sb_x0, sb_x1], [sb_y, sb_y], color='white', linewidth=2)
ax[1].text((sb_x0 + sb_x1) / 2, sb_y - 5, SCALEBAR_LABEL,
           color='white', fontsize=8, ha='center', va='bottom')

plt.legend()
plt.show()