# -*- coding: utf-8 -*-
"""
Created on Mon Nov 17 10:33:57 2025.

@author: Jiawen Jiang
"""

'''

This script includes functions that perform the following:
    - loading images and plists from z-stack imaging,
    - preprocessing images,
    - various PIV calculations
    - calculating 3D velocity components

It includes mostly standard Python packages, additionally:
    - OpenPIV: only used for slicing images into small windows and finding peaks
    in correlation matrices.
    - j_py_sad_correlation: custom package made by J.Taylor. Used to calculate
    SAD faster than standard Python method, although Numpy could equivalenty be used
    for marginally slower computation. 
    - plistlib: for reading plist files
    
'''
import numpy as np
from skimage import io
from openpiv import tools, pyprocess
import j_py_sad_correlation as jps
import time
import os
import pickle
import random
import plistlib as plib
from natsort import natsorted
import tifffile
#%%
def load_plist(fp):
    """
    Load plist files that contain metadata for images.

    Parameters
    ----------
    fp : file path of plists, can be a single plist or a list of file names 
    generated by glob().

    Returns
    -------
    plist : plists 

    """
    plist=[]
    if type(fp) == list:
        fp = natsorted(fp)
        for i in range(0, len(fp)):
            with open(fp[i], 'rb') as myfile:
                plist.append(plib.load(myfile))
            myfile.close()
        return plist
    
    else:
        with open(fp, 'rb') as myfile:
            plist=plib.load(myfile)
        myfile.close()
        return plist

def get_z_pos(plist, tolerance = 0.001, min_hold=100):
    """
    Idenfity the z positions of imaging planes in z-stack.

    Parameters
    ----------
    plist : list, plists returned from load_plist
    tolerance : float, optional
        max separation between z values at one z plane. The default is 0.001.
    min_hold : float, optional
        min number of frames within z range required to identify a z plane. The default is 3.

    Returns
    -------
    paused_positions : list of z plane positions.

    """
    all_z = []
    all_plists = plist
    
    for i in range(0,len(all_plists)):
        for j in range(0, len(all_plists[i].get('frames'))):
            all_z.append(all_plists[i].get('frames')[j].get('stage_positions').get('last_known_z'))
    all_z = np.array(all_z)   
    zdiff = np.abs(np.diff(all_z))
    #mark frames at z planes where z values are constant
    stable = np.concatenate([[False], zdiff < tolerance])
    paused_positions = []
    start = None
    for i, is_stable in enumerate(stable):
        if is_stable and start is None:
            start = i
        elif not is_stable and start is not None:
            if i - start >= min_hold:
                paused_z = np.mean(all_z[start:i])
                paused_positions.append(round(paused_z,1))
            start = None
    # Handle case where last segment is stable
    if start is not None and len(all_z) - start >= min_hold:
        paused_z = np.mean(all_z[start:])
        paused_positions.append(round(paused_z,1))

    return paused_positions

def select_image(tif_list, plist, z_value, z_range, phase, phase_range, condition=0, limit=20):
    """
    Select PIV frame pairs within a given z and phase range.

    Parameters
    ----------
    tif_list : list of image file paths.
    plist : list of plists.
    z_value : chosen z position.
    z_range : chosen z range.
    phase : chosen heart phase.
    phase_range : chosen phase range.
    condition : 1 or 0, determines if the sequence start with a frameA or frameB. The default is 0.
    limit : int, number of frame pairs returned. The default is 20.

    Returns
    -------
    first_frames : list of frameA.
    second_frames : list of frameB.
    """
    first_frames = []
    second_frames = []
    if phase<0:
        phase+=np.pi*2
    elif phase>np.pi*2:
        phase-=np.pi*2
    ilist=[]
    jlist=[]
    #the loop gathers the index information of frames within the chosen z&phase range
    tif_list = natsorted(tif_list) #put list of image file paths in natural order
    for i in range(0, len(plist)):
        tempj = []
        for j in range(0, len(plist[i].get('frames'))):
            
            current_pl = plist[i].get('frames')[j]
            if current_pl.get('frame_number')%2 == condition and current_pl.get('frame_number')>0: #only obtain phase info from frameB
                image_phase = current_pl.get('phase_from_offline_sync_analysis')
                z_pos = current_pl.get('stage_positions').get('last_known_z')
                if type(image_phase) == float: #added as some phase values = None, when heart sync fails
                    if image_phase<(phase+phase_range) and image_phase>(phase-phase_range):
                        if z_pos<z_value+z_range and z_pos>z_value-z_range:
                            if sum(len(sublist) for sublist in jlist) < limit and j!=0:
                                tempj.append(j)
            
        if len(tempj) > 0:
            jlist.append(tempj)
            ilist.append(i)
    #load in images at certain indices of certain tiff file
    #for this analysis, typically only one tiff file is loaded
    for i in range(0, len(ilist)):
        with tifffile.TiffFile(tif_list[ilist[i]]) as tif:
            for j in range(0, len(jlist[i])):
                first_frames.append(tif.pages[jlist[i][j]-1].asarray())
                second_frames.append(tif.pages[jlist[i][j]].asarray())
    
    
    return first_frames, second_frames

def save_frames(first_frames, second_frames, phase, z, stack_num = 1):
    """
    Cache frame pairs obtained from select_image.

    Parameters
    ----------
    first_frames : list of frameA.
    second_frames : list of frameB.
    phase : phase value, just for labelling.
    z : z value, just for labelling.
    stack_num : label of the folder storing the stack, The default is 'stack_1/'.

    Returns
    -------
    None.

    """
    if stack_num == 1:
        stack_num = 'stack_1/'
    elif stack_num == 2:
        stack_num = 'stack_2/'
    else:
        print('please input correct stack number 1 or 2')
        return
    path = 'frame_cache/' + stack_num + str(phase) + '/' + str(z) + '/'
    ifExist = os.path.exists(path)
    if not ifExist:
        os.makedirs(path)
    first_path = path + 'first.npz'
    second_path = path + 'second.npz'
    
    first_stack = np.stack(first_frames)
    second_stack = np.stack(second_frames)
    np.savez(first_path, stack=first_stack)
    np.savez(second_path, stack=second_stack)
    return

def load_frames(phase, z, stack_num = 1):
    """
    Load cached frame pairs obtained from select_image.

    Parameters
    ----------
    phase : phase value, need to be exactly the same as the one used in caching.
    z : z value, need to be exactly the same as the one used in caching.
    stack_num : label of the folder storing the stack, The default is 'stack_1/'.

    Returns
    -------
    None.

    """
    if stack_num == 1:
        stack_num = 'stack_1/'
    elif stack_num == 2:
        stack_num = 'stack_2/'
    else:
        print('please input correct stack number 1 or 2')
        return
    path = 'frame_cache/' + stack_num + str(phase) + '/' + str(z) + '/'
    first_path = path + 'first.npz'
    second_path = path + 'second.npz'
    
    first_exist = os.path.exists(first_path)
    second_exist = os.path.exists(second_path)
    if not (first_exist and second_exist):
        print('No images found for given phase and z')
        return
    else:
        data = np.load(first_path)
        first_frames = data['stack']
        data = np.load(second_path)
        second_frames = data['stack']
        
        
    
    return first_frames, second_frames

def extended_window_correlate(win1, win2, smallsize):
    """
    Cross correlate a small section of win1 (IW) with win2.

    Parameters
    ----------
    win1 : 2x2 array of window image.
    win2 : 2x2 array of window image.
    smallsize : size of interrogation window (IW) in [y,x].

    Returns
    -------
    mat : 2x2 correlation matrix.

    """
    midy = int(round(len(win1)/2))
    midx = int(round(len(win1[0])/2))
    halfsizex = int(round(smallsize[1]/2))
    halfsizey = int(round(smallsize[0]/2)) #size of IW
    iw = win1[midy-halfsizey:midy+halfsizey, midx-halfsizex:midx+halfsizex] #crops out IW
    mat = np.zeros((len(win2)-len(iw), len(win2[0])-len(iw[0])), dtype=float)
    for i in range(0, len(win2)): #calculate SAD for each IW position in win2
        for j in range(0, len(win2[0])):
            if i+len(iw) < len(win2) and j+len(iw[0]) < len(win2[0]):
                sad = jps.sad_correlation(iw, win2[i:i+len(iw), j:j+len(iw[0])])[0][0]
                mat[i][j] = -sad
    mat = mat - mat.min() #this inverts the correlation matrix because we want to find the 
                          #minimum SAD, but openpiv function finds the peak
    return mat
    
def extended_search_correlate(first_set, second_set, winsize=[0,0], smallwinsize=[0,0], overlap=0, s2n_threshold = 1.0, s2n_p2p = 1.0, intfilter=False):
    """
    Calculate the vector field using correlation averaged PIV.

    Parameters
    ----------
    first_set : list of frameA.
    second_set : list of frameB.
    winsize : size of search window in [y,x]. The default is [0,0].
    smallwinsize : size of interrogation window in [y,x]. The default is [0,0].
    overlap : size of overlap between search windows. The default is 0.
    s2n_threshold : SNR peak to mean threshold to filter out noise. The default is 1.0.
    s2n_p2p : SNR peak to peak threshold to filter out noise. The default is 1.0.
    intfilter : whether to switch on intensity filter. Window with avg pixel
    intensity smaller than image avg will be filtered out. The default is False.

    Returns
    -------
    x : x coordinates of vector points.
    y : y coordinates of vector points.
    u : x velocities of vector points.
    v : y velocities of vector points.

    """
    t1 = time.perf_counter()
    corr_mat = []
    image = []
    for i in range(0, len(first_set)):
        #print(i)
        wincor = [] 
        #producing windows from frame pairs
        win1 = pyprocess.sliding_window_array(first_set[i],winsize,overlap)
        win2 = pyprocess.sliding_window_array(second_set[i],winsize,overlap)
        image.append(win1)
        for j in range(0, len(win1)):
            
            #calculating correlation matrices for windows
            wincor.append(extended_window_correlate(win1[j], win2[j], smallwinsize))
        corr_mat.append(wincor)
    image = np.array(image)
    avg_img = image.mean(axis = 0)
    corr_mat = np.array(corr_mat)  
    mean_corr = corr_mat.mean(axis=0) #taking averaged correlation matrix
    grid = pyprocess.get_field_shape(first_set[0].shape,search_area_size=winsize,overlap=overlap)
    nrows, ncols = grid[0], grid[1]
    #using openpiv to calculate displacement from correlation matrices
    u,v = pyprocess.correlation_to_displacement(mean_corr, nrows, ncols)
    x,y = pyprocess.get_rect_coordinates(first_set[0].shape, winsize, overlap)
    #calculating SNR
    s2np2m = pyprocess.sig2noise_ratio(mean_corr, 'peak2mean')
    s2np2p = pyprocess.sig2noise_ratio(mean_corr, 'peak2peak')
    #using SNR filter to set vectors with large noise to 0
    for i in range(0,len(s2np2m)):
        if s2np2m[i]<s2n_threshold or s2np2p[i]<s2n_p2p:
                
            u[i//ncols][i%ncols] = 0
            v[i//ncols][i%ncols] = 0
    #applying intensity filter
    if intfilter:
        for i in range(0, len(avg_img)):
            if avg_img[i].mean() < np.mean(avg_img.flatten()):
                u[i//ncols][i%ncols] = 0
                v[i//ncols][i%ncols] = 0
    x, y, u, v = tools.transform_coordinates(x, y, u, v)
    t2 = time.perf_counter()
    print('Time taken to process correlation: ', str(t2 - t1))
    return x,y,u,v

def bg_subtraction(images):
    """
    subtract background from a group of images.

    Parameters
    ----------
    images : list of images.

    Returns
    -------
    subtracted : list of images with background subtracted.

    """
    subtracted = []
    avim = np.median(images, axis=0)
    for image in images:
        subtracted.append(np.clip(image-avim, 0, np.max(image)))
    return subtracted
    
    
    
    
def correlate_at_pos(first_set, second_set, xpos, ypos, winsize=[0,0], smallwinsize=[0,0], p2p = 1.0, p2m = 1.0, intfilter=True, retmat=False, bs_iter=3000):
    """
    Perform correlation averaged PIV at a given position, useful for multiview PIV.
    
    As it calculates flow velocity at calculated intersection positions between 
    two z-stacks

    Parameters
    ----------
    first_set : list of frameA.
    second_set : list of frameB.
    xpos : x position.
    ypos : y position.
    winsize : size of search window in [y,x]. The default is [0,0].
    smallwinsize : size of interrogation window in [y,x]. The default is [0,0].
    p2p : SNR threshold peak to peak. The default is 1.0.
    p2m : SNR threshold peak to mean. The default is 1.0.
    intfilter : whether to switch on intensity filter. The default is True.
    retmat : whether to return correlation matrix for debugging purposes. The default is False.
    
    Returns
    -------
    u : x displacement, float.
    v : y displacement, float.

    """
    #calculating size of search window in x and y
    half_sizex = int(round(winsize[1]*0.5))
    half_sizey = int(round(winsize[0]*0.5))
    #return 0 if given position is too close to the edge of the image
    if xpos-half_sizex < 0 or ypos-half_sizey < 0:
        return 0,0,0,0
    if xpos+half_sizex >= len(first_set[0][0]) or ypos+half_sizey >= len(first_set[0]):
        return 0,0,0,0
    matrices = []
    images = []
    windows = []
    for i in range(0, len(first_set)):
        im1 = first_set[i]
        im2 = second_set[i]
        ylower = int(round(ypos-half_sizey))
        yupper = int(round(ypos+half_sizey))
        xlower = int(round(xpos-half_sizex))
        xupper = int(round(xpos+half_sizex))
        #producing search window
        win1 = im1[ylower:yupper, xlower:xupper]
        win2 = im2[ylower:yupper, xlower:xupper]
        images.append(im1)
        windows.append(win1)
        #passing search windoe to correlation function to calculate correlation matrices
        matrices.append(extended_window_correlate(win1, win2, smallwinsize))
        
    if intfilter and np.mean(windows) < np.mean(images):
        #return 0 if the window intensity is lower than image average
        return 0,0,0,0
    matrices = np.stack(matrices)
    
    x_dis = []
    y_dis = []
    #randomly sample subsets from the list of all matrices
    for _ in range(bs_iter):
        idx = np.random.choice(matrices.shape[0], matrices.shape[0], replace=True)
        chosen_matrices = matrices[idx]
        avg_corr = np.mean(chosen_matrices, axis=0)
        u,v = pyprocess.correlation_to_displacement(np.array([avg_corr]), 1, 1)
        x_dis.append(u[0][0])
        y_dis.append(v[0][0])
    #xmean = np.mean(x_dis)
    #ymean = np.mean(y_dis)
    xstd = np.std(x_dis)
    ystd = np.std(y_dis)
    
    mat = matrices.mean(axis=0) #average correlation matrix
    if retmat == True:
        return mat
    s2n_p2p = pyprocess.sig2noise_ratio(np.array([mat]), 'peak2peak')
    s2n_p2m = pyprocess.sig2noise_ratio(np.array([mat]), 'peak2mean')
    if s2n_p2p < p2p or s2n_p2m < p2m:
        #return 0 if SNR too low
        return 0,0,0,0
    u,v = pyprocess.correlation_to_displacement(np.array([mat]), 1, 1)
    if retmat==True:
        return mat
    return u[0,0], v[0,0], xstd, ystd
    
def mon_bootstrap(first_set, second_set, xpos, ypos, winsize=[0,0], smallwinsize=[0,0], samplesize=20, iteration=5000,p2p=1.0,p2m=1.0):
    """
    Calculate uncertainty in corr. avg. PIV measurement using m-out-of-n bootstrapping.
    
    This may become useful in cases where normal bootstrapping fails, but it is
    not trivial to choose suitable samplesize (m). 

    Parameters
    ----------
    first_set : list of frameA.
    second_set : list of frameB.
    xpos : x position, float.
    ypos : y position, float.
    winsize : size of search window in [y,x]. The default is [0,0].
    smallwinsize : size of interrogation window in [y,x]. The default is [0,0].
    samplesize : number of frame pairs in one subset. The default is 20.
    iteration : number of iterations for bootstrapping. The default is 5000.
    p2p : SNR threshold peak to peak. The default is 1.0.
    p2m : SNR threshold peak to mean. The default is 1.0.

    Returns
    -------
    x_dis : list of x displacement values calculated from N=iteration subsets
    y_dis : list of y displacement values

    """
    if len(first_set) < samplesize+4:
        print('insufficient pair, reducing sample size')
        samplesize = len(first_set) - 4
        #dynamically reducing sample size if there are not enough frame pairs
        #but should not be an issue anymore with the new high FPS camera
    #same procedure as correlate_at_pos to calculate correlation matrices
    half_sizex = int(round(winsize[1]*0.5))
    half_sizey = int(round(winsize[0]*0.5))
    if xpos-half_sizex < 0 or ypos-half_sizey < 0:
        return [0,0]
    if xpos+half_sizex >= len(first_set[0][0]) or ypos+half_sizey >= len(first_set[0]):
        return [0,0]
    matrices = []
    images = []
    windows = []
    for i in range(0, len(first_set)):
        im1 = first_set[i]
        im2 = second_set[i]
        ylower = int(round(ypos-half_sizey))
        yupper = int(round(ypos+half_sizey))
        xlower = int(round(xpos-half_sizex))
        xupper = int(round(xpos+half_sizex))
        
        win1 = im1[ylower:yupper, xlower:xupper]
        win2 = im2[ylower:yupper, xlower:xupper]
        images.append(im1)
        windows.append(win1)
        
        matrices.append(extended_window_correlate(win1, win2, smallwinsize))
        
    matrices = np.stack(matrices)
    x_dis = []
    y_dis = []
    #randomly sample subsets from the list of all matrices
    for _ in range(iteration):
        idx = np.random.choice(matrices.shape[0], samplesize, replace=False)
        chosen_matrices = matrices[idx]
        avg_corr = np.mean(chosen_matrices, axis=0)
        #mat = np.array(avg_corr) #correlation averaging
        u,v = pyprocess.correlation_to_displacement(np.array([avg_corr]), 1, 1)
        #s2n_p2p = pyprocess.sig2noise_ratio(np.array([avg_corr]), 'peak2peak')
        #s2n_p2m = pyprocess.sig2noise_ratio(np.array([avg_corr]), 'peak2mean')
        #gives 0 if SNR too low
        #if s2n_p2p < p2p or s2n_p2m < p2m:
            #x_dis.append(0)
            #y_dis.append(0)
        #else:
            #x_dis.append(u[0][0])
            #y_dis.append(v[0][0])
        x_dis.append(u[0][0])
        y_dis.append(v[0][0])
    #xmean = np.mean(x_dis)
    #ymean = np.mean(y_dis)
    xstd = np.std(x_dis)
    ystd = np.std(y_dis)
    #return x_dis, y_dis
    #can be modified to return calculated mean and std
    #return [xmean,xstd],[ymean,ystd]  
    return xstd, ystd

def bootstrap(first_set, second_set, xpos, ypos, winsize=[0,0], smallwinsize=[0,0], iteration=5000,p2p=1.0,p2m=1.0):
    """
    Calculate uncertainty in corr. avg. PIV measurement using bootstrapping.
    
    CORRECTED VERSION - subsampling N from N with replacement

    Parameters
    ----------
    first_set : list of frameA.
    second_set : list of frameB.
    xpos : x position, float.
    ypos : y position, float.
    winsize : size of search window in [y,x]. The default is [0,0].
    smallwinsize : size of interrogation window in [y,x]. The default is [0,0].
    samplesize : number of frame pairs in one subset. The default is 20.
    iteration : number of iterations for bootstrapping. The default is 5000.
    p2p : SNR threshold peak to peak. The default is 1.0.
    p2m : SNR threshold peak to mean. The default is 1.0.

    Returns
    -------
    x_dis : list of x displacement values calculated from N=iteration subsets
    y_dis : list of y displacement values

    """
    #same procedure as correlate_at_pos to calculate correlation matrices
    half_sizex = int(round(winsize[1]*0.5))
    half_sizey = int(round(winsize[0]*0.5))
    if xpos-half_sizex < 0 or ypos-half_sizey < 0:
        return [0,0]
    if xpos+half_sizex >= len(first_set[0][0]) or ypos+half_sizey >= len(first_set[0]):
        return [0,0]
    matrices = []
    images = []
    windows = []
    for i in range(0, len(first_set)):
        im1 = first_set[i]
        im2 = second_set[i]
        ylower = int(round(ypos-half_sizey))
        yupper = int(round(ypos+half_sizey))
        xlower = int(round(xpos-half_sizex))
        xupper = int(round(xpos+half_sizex))
        
        win1 = im1[ylower:yupper, xlower:xupper]
        win2 = im2[ylower:yupper, xlower:xupper]
        images.append(im1)
        windows.append(win1)
        
        matrices.append(extended_window_correlate(win1, win2, smallwinsize))
        
    matrices = np.stack(matrices)
    x_dis = []
    y_dis = []
    #randomly sample subsets from the list of all matrices
    for _ in range(iteration):
        idx = np.random.choice(matrices.shape[0], matrices.shape[0], replace=True)
        chosen_matrices = matrices[idx]
        avg_corr = np.mean(chosen_matrices, axis=0)
        #mat = np.array(avg_corr) #correlation averaging
        u,v = pyprocess.correlation_to_displacement(np.array([avg_corr]), 1, 1)
        #s2n_p2p = pyprocess.sig2noise_ratio(np.array([avg_corr]), 'peak2peak')
        #s2n_p2m = pyprocess.sig2noise_ratio(np.array([avg_corr]), 'peak2mean')
        #gives 0 if SNR too low
        #if s2n_p2p < p2p or s2n_p2m < p2m:
            #x_dis.append(0)
            #y_dis.append(0)
        #else:
            #x_dis.append(u[0][0])
            #y_dis.append(v[0][0])
        x_dis.append(u[0][0])
        y_dis.append(v[0][0])
    #xmean = np.mean(x_dis)
    #ymean = np.mean(y_dis)
    xstd = np.std(x_dis)
    ystd = np.std(y_dis)
    return x_dis, y_dis
    #can be modified to return calculated mean and std
    #return [xmean,xstd],[ymean,ystd]  
    #return xstd, ystd

def calc_true(stack1, stack2, angle=20):
    """
    Calclate the z-velocity from x velocities from two stacks.

    Parameters
    ----------
    stack1 : x velocity from stack 1, float.
    stack2 : x velocity from stack2, float.
    angle : rotation angle between two stacks in degrees. The default is 20.

    Returns
    -------
    [x velocity, z velocity]

    """
    #rotation angle in degrees
    angle_rad = np.deg2rad(angle)

    #calculate endpoints
    P1 = (stack1 * np.cos(angle_rad), stack1 * np.sin(angle_rad))  
    P2 = (stack2, 0) 

    
    m_normal_1 = -1 / np.tan(angle_rad)#gradient of the normal to stack 1 line
    b_normal_1 = P1[1] - m_normal_1 * P1[0]#y intercept of the normal 


    x_normal_2 = P2[0]

    #calculating the intersection of two normal lines
    x_intersection = x_normal_2
    y_intersection = m_normal_1 * x_intersection + b_normal_1

    intersection_point = (x_intersection, y_intersection)
    return intersection_point #this would return [x_vel, z_vel]

def calc_true_error(params, errors, gridsize=[100,100,10]):
    """
    Calculate uncetainty in z velocity component using a numerical approach.

    Parameters
    ----------
    params : list [u1,u2,angle] where u1 u2 are x velocities from stack 1 and 2, angle is rotation between stacks.
    errors : list [du1,du2,dangle], error in each parameter.
    gridsize : sample size for numerical calculation of error. The default is [100,100,10].

    Returns
    -------
    x_err : error in x velocity.
    z_err : error in z velocity.

    """
    xvals = []
    zvals = []
    u1s, u2s, angle = params
    u1err, u2err, angleerr = errors
    #generating sample values of u1, u2 and angle within uncertainty range
    stack1 = np.linspace(u1s-u1err, u1s+u1err, gridsize[0])
    stack2 = np.linspace(u2s-u2err, u2s+u2err, gridsize[1])
    angles = np.linspace(angle-angleerr, angle+angleerr, gridsize[2])
    #calculating z component using all possible combinations of values
    for u1 in stack1:
        for u2 in stack2:
            for theta in angles:
                result = calc_true(u1, u2, theta)
                xvals.append(result[0])
                zvals.append(result[1])
    #finding the max&min to calculate uncertainty range
    xerr = (np.max(xvals) - np.min(xvals))/2
    zerr = (np.max(zvals) - np.min(zvals))/2
    return np.abs(xerr), np.abs(zerr)


def calculate_zvel(u1, u2, angle=20):
    """
    Calclate the z-velocity from x velocities from two stacks.
    
    Uses a analytical approach derived from rotation matrix, unlike the previous
    function which uses graphical approach. By using analytical approach, correct
    error propagation can be used. 

    Parameters
    ----------
    u1 : x velocity from stack 1, float.
    u2 : x velocity from stack2, float.
    angle : rotation angle between two stacks in degrees. The default is 20.

    Returns
    -------
    [x velocity, z velocity]

    """
    angle = np.deg2rad(angle)
    w = (u1 - u2*np.cos(angle))/np.sin(angle)
    return [u2, w]
    
def calculate_zerr(params, error):
    """
    Calculate uncetainty in z velocity component using a analytical approach.
    
    As z velocity can be calculated from an analytical equation derived from
    rotation matrix, its uncertainty can be propagated from uncertainties in
    u1, u2 and angle.
    
    Parameters
    ----------
    params : list [u1,u2,angle] where u1 u2 are x velocities from stack 1 and 2, angle is rotation between stacks.
    errors : list [du1,du2,dangle], error in each parameter.
    
    Returns
    -------
    du2 : error in x velocity.
    dz : error in z velocity.

    """
    u1, u2, angle = params
    du1, du2, da = error
    angle = np.deg2rad(angle)
    da = np.deg2rad(da)
    
    dz_du1 = 1/np.sin(angle)
    dz_du2 = -np.cos(angle)/np.sin(angle)
    dz_da = u2 - u1*np.cos(angle)/np.sin(angle)**2 + u2*np.cos(angle)**2/np.sin(angle)**2
    
    dz = np.sqrt((dz_du1*du1)**2 + (dz_du2*du2)**2 + (dz_da*da)**2)
    return [du2, dz]

def grey2rgb(image, color = 'r'):
    """
    Convert greyscale image to rgb.

    Parameters
    ----------
    image : image array in greyscale.
    color : color in 'r' or 'g' for red or green. The default is 'r'.

    Returns
    -------
    colored image array.

    """
    image = np.array(image)
    empty = np.zeros((len(image), len(image[0]), 3))
    m = image.max()
    for i in range(0,len(image)):
        for j in range(0, len(image[i])):
            if color == 'r':
                empty[i][j][0] = image[i][j]/m
            if color == 'g':
                empty[i][j][1] = image[i][j]/m
            if color == 'm':
                empty[i][j][0] = image[i][j]/m
                empty[i][j][2] = image[i][j]/m
    return np.array(empty)

def build_rgb(gata, gfp):
    """
    Build colored image from red and green images of cell and heart wall.

    Parameters
    ----------
    gata : greyscale image of blood cells.
    gfp : greyscale image of heart wall.

    Returns
    -------
    newim : colored image.

    """
    newim = grey2rgb(gata, 'r') + grey2rgb(gfp, 'g')
    return newim



def closest_wall(image_path, plist, z_value, z_range, phase, phase_range=0.1, non_piv=True):
    """
    Obtain the GFP image closest to the given phase and z position.
    
    The image with the closest phase to the chosen phase within z range will
    be selected. This is because many images would have the same z position
    but phase values would be more unpredictable. 

    Parameters
    ----------
    image_path : list of file names of GFP images.
    plist : plists of GFP images.
    z_value : chosen z position.
    z_range : chosen range of z position, reduce for more precision.
    phase : chosen phase value.
    phase_range : chosen phase range. 
    non_piv : whether images were taken in pulsed mode/PIV mode. The default is True.

    Returns
    -------
    one GFP image array.

    """
    image_path = natsorted(image_path)
    #fixing phase values outside range 0-2pi
    if phase>np.pi*2:
        phase = phase - 2*np.pi
    
    if phase<0:
        phase = phase+2*np.pi
    ijlist = [] #recording the indices of all frames within z&phase range
    residual = [] #recording the difference between actual phase and chosen phase
    wall_plists = plist
    tif_list_wall = image_path
    imagephases=[]
    height = []
    for i in range(0, len(wall_plists)):
        for j in range(0, len(wall_plists[i].get('frames'))):
            if non_piv:
                image_phase = wall_plists[i].get('frames')[j].get('phase_from_offline_sync_analysis')
                z_pos = wall_plists[i].get('frames')[j].get('stage_positions').get('last_known_z')
                if z_pos<z_value+z_range and z_pos>z_value-z_range and image_phase != None:
                    residual.append(np.abs(image_phase - phase)) #get phase diff
                    imagephases.append(image_phase)
                    ijlist.append([i,j])
                    height.append(z_pos)
    ind = residual.index(min(residual)) #index of image with closest phase
    #frame = io.imread(tif_list_wall[ijlist[ind][0]])
    #frame = frame[ijlist[ind][1]]
    #fetch specific frame in the large tiff file for faster speed
    with tifffile.TiffFile(tif_list_wall[ijlist[ind][0]]) as tif:
        frame = tif.pages[ijlist[ind][1]].asarray()
    #frame = np.flipud(frame)
    
    return frame #, imagephases[ind], height[ind]

def get_y_projection(image):
    """
    Produce the maximum intensity projection along the y axis for a single image.

    Parameters
    ----------
    image : 2x2 image array.

    Returns
    -------
    array : 1D array of a row.

    """
    mip = np.max(image, axis=0) #0 for projection in y, 1 for x
    return mip
    
def get_y_cross_section(image):
    """
    Produce a row for the x-z cross section from a single image.
    
    Repeat this for images at different z values to form a x-z cross section, at a given y plane.

    Parameters
    ----------
    image : 2x2 image array.

    Returns
    -------
    array : 1D array of a row.

    """
    array = []
    ypos = int(round(0.5*len(image))) #by default, cross section taken at the center of zstack
    for row in image.T:
        array.append(row[ypos])
    
    return array
def roll_frames(first_frames, second_frames, roll):
    """
    Shift frame pairs in the y direction.
    
    This aligns two stacks in the y direction - it is easier to preprocess
    images than dealing with intersection calculations, as y-positions of 
    intersects can be kept the same between two stacks. Shift below the size of
    SW does not affect PIV because PIV cannot be done at the edge of the image
    anyway.

    Parameters
    ----------
    first_frames : list of frameA.
    second_frames : list of frameB.
    roll : shift in y, in pixels.

    Returns
    -------
    new_first : list of shifted frameA.
    new_second : list of shifted frameB.

    """
    new_first=[]
    new_second=[]
    for frame in first_frames: 
        new_first.append(np.roll(frame, (roll, 0), axis=(0, 1)))
    for frame in second_frames:
        new_second.append(np.roll(frame, (roll, 0), axis=(0, 1)))
    
    #new_first = np.array(new_first)
    #new_second = np.array(new_second)
    return new_first, new_second

def calculate_u1(u2, z, angle=20):
    #temporary function to get stack1 velocity from calculated result of 
    #3 component velocity, because stack1 velocity was not recorded in 
    #calculated vector fields.
    #for the purpose of propagating error in x velocity to z velocity. 
    angle = np.deg2rad(angle) 
    u1 = z * np.sin(angle) + u2*np.cos(angle)
    return u1







