'''
@author: Jiawen Jiang

This script visualise the velocity vector fields from two stacks, overlayed
on GFP images processed by affine transform. It demonstrates how measurements 
needed for multiview velocity fusion (MVF) were obtained from 2D PIV measurements.

'''
import numpy as np
import glob
import matplotlib.pyplot as plt

from tqdm import tqdm
from functions import *  
from stack_intersect_calculation import *
import pickle
import os
import config
#%%
stack1tif = glob.glob(config.stack1gata_image)
stack1p_list = glob.glob(config.stack1gata_plist)
stack2tif = glob.glob(config.stack2gata_image)
stack2p_list = glob.glob(config.stack2gata_plist)
stack1plist = load_plist(stack1p_list)
stack2plist = load_plist(stack2p_list)
#%%
stack1z = get_z_pos(stack1plist)
stack2z = get_z_pos(stack2plist)

xlength = config.xlength
um_per_pixel = config.um_per_pixel
xlength = xlength*um_per_pixel #in unit um
ylength = config.ylength #in unit pixel

#tilt_index = 15 
midz = config.midz #midpoint of z positions
stack_phase = config.stack_phase #phase of stack1 gata
plane_phase = config.plane_phase #phase of stack2 gata
nplanes=config.nplanes #how many imaging planes
plane_sep = config.plane_sep #separation between imaging planes
rotation = config.rotation #rotation angle in degrees between two stacks
#nofrows = config.nofrows #how many rows of vectors in y, could be changed to get different field resolution
nofrows = 25 #just for this test - could be reduced for actual field calculation for all phases
#stack_shift = config.stack_shift #shift in stack1 in [x,z] measured from spatial registration
#y_shift = config.y_shift #shift in y direction of stack2
SW_size = config.SW_size #size of search window in pixels
IW_size = config.IW_size #size of interrogation window in pixels 
SNR_p2m_threshold = config.SNR_p2m_threshold
SNR_p2p_threshold = config.SNR_p2p_threshold
bs_iteration = config.bootstrap_iteration
current_phase = 1.1
phase_shift = 4.2
#%%
def test_plane(tilt_index = 15, rotation = rotation, xshift = 0, y_shift = 0, zshift = 0, chamber='all'):
    xshift *= config.um_per_pixel
    zshift *= config.um_per_pixel
    stack_shift = [xshift,zshift] #loading in stack registration parameters
    #calculating intersection positions between two z-stacks
    intersections, tilt_line, flat_stack = stack_line_intersection((xlength/2, midz), plane_sep, nplanes, rotation, xlength, tilt_index, stack_shift)
    stack_x, stack_z = get_stack_x_grid(intersections, nofrows) #generate gridpoints of intersections
    #generating position grid using calculated intersections
    plane_x_line = distance_along_line(intersections, tilt_line)
    plane_x = get_plane_x_grid(plane_x_line, nofrows)
    plane_y = get_y_grid(0, ylength, len(stack_x[0]), nofrows)
    stack_y = plane_y - y_shift
    stack_u = []
    stack_v = []
    stack_uerr = []
    stack_verr = []
    #setting boundaries for measuring flow within specific chambers, or the whole heart
    if chamber == 'all':
        lowerx, upperx = 0, 10000
    elif chamber == 'atrium':
        lowerx, upperx = 200, 10000
    elif chamber == 'ventricle':
        lowerx, upperx = 0, 200
    for k in range(0, len(stack_z)):
        first_frames, second_frames = select_image(stack1tif, stack1plist, stack_z[k], 0.05, current_phase, np.pi/60, limit=25, condition = config.pair_selection_condition_stack1)
        #loading frame pairs
        urow=[]
        vrow=[]
        uerrrow = []
        verrrow = []
        for j in range(0, len(stack_y)):
            xpos = stack_x[0][k] #for each z plane, get all [x,y] coordinates of intersections
            ypos = stack_y[j][0]
            
            planexpos = plane_x[0][k]
            if lowerx < planexpos < upperx:
                if len(first_frames) == 0:
                    #if no frames exist for the given selection criteria
                    u = 0
                    v = 0
                    uerr=0
                    verr=0
                else:
                    #performs cross correlation at the given position
                    #PIV function has updated to return uncertainty measurements as well
                    u, v, uerr, verr = correlate_at_pos(first_frames, second_frames, xpos, ypos, SW_size,IW_size, SNR_p2m_threshold, SNR_p2p_threshold)
                urow.append(u)
                vrow.append(v)
                uerrrow.append(uerr)
                verrrow.append(verr)
        if lowerx < planexpos < upperx: 
            #recording measurements within given range
            stack_u.append(urow)
            stack_v.append(vrow)
            stack_uerr.append(uerrrow)
            stack_verr.append(verrrow)
    #process measured velocity vector field to consistent shapes    
    stack_u = np.array(stack_u).T
    stack_v = np.array(stack_v).T
    stack_uerr = np.array(stack_uerr).T
    stack_verr = np.array(stack_verr).T
    xs,ys,us,vs = tools.transform_coordinates(stack_x, stack_y, stack_u, stack_v)
    xserr, yserr, userr, vserr = tools.transform_coordinates(stack_x, stack_y, stack_uerr, stack_verr)
    #transform coordinates such that the origin is at the bottom left
    first_frames, second_frames = select_image(stack2tif, stack2plist, stack2z[tilt_index], 0.05, current_phase+phase_shift, np.pi/60, limit=25, condition = config.pair_selection_condition_stack2)
    #loading and shifting frame pairs from stack 2
    
    plane_u = [] #calculate intersections on a single plane from stack2
    plane_v = []
    plane_uerr = []
    plane_verr = []
    for i in range(0, len(plane_x)):
        urow = []
        vrow = []
        uerrrow = []
        verrrow = []
        for j in range(0, len(plane_x[0])):
            if lowerx < plane_x[i][j] < upperx:
                if len(first_frames) == 0:
                    #if no frames exist for the given selection criteria
                    u = 0
                    v = 0
                    uerr = 0
                    verr = 0
                else:
                    u, v, uerr, verr = correlate_at_pos(first_frames, second_frames, plane_x[i][j], plane_y[i][j], SW_size,IW_size, SNR_p2m_threshold, SNR_p2p_threshold)
                    #perform PIV at intersection
                urow.append(u)
                vrow.append(v)
                uerrrow.append(uerr)
                verrrow.append(verr)
        plane_u.append(urow)
        plane_v.append(vrow)
        plane_uerr.append(uerrrow)
        plane_verr.append(verrrow)
        
    plane_u = np.array(plane_u)
    plane_v = np.array(plane_v)
    plane_uerr = np.array(plane_uerr)
    plane_verr = np.array(plane_verr)
    xp,yp,up,vp = tools.transform_coordinates(plane_x, plane_y, plane_u, plane_v)
    xperr, yperr, uperr, vperr = tools.transform_coordinates(plane_x, plane_y, plane_uerr, plane_verr)
    #shift vector field if needed, normally this is not used as spatial
    #alignment of zstacks should already give correct intersections
    stack2_x_field = np.roll(up, (0,0), axis=(0,1))
    stack1_x_field = us
    stack2_y_field = np.roll(vp, (0,0), axis=(0,1))
    stack1_y_field = vs
    
    stack1y = []
    stack2y = []
    stack1yerr = []
    stack2yerr = []
    
    for i in range(0, len(stack1_y_field.flatten())):
        y1 = stack1_y_field.flatten()[i]
        y2 = stack2_y_field.flatten()[i]
        y1err = vserr.flatten()[i]
        y2err = vperr.flatten()[i]
        #removing vector measurements that give 0 - invalid measurements
        if y1*y2 != 0 and y1err*0 == 0 and y2err*0 == 0:
            stack1y.append(y1)
            stack2y.append(y2)
            stack1yerr.append(y1err)
            stack2yerr.append(y2err)
    
    return [xp,yp,up,vp], [xp,yp,us,vs], np.mean(first_frames, axis=0)
    #return np.array(stack1y), np.array(stack1yerr), np.array(stack2y), np.array(stack2yerr)
#%% calculating vector field using PIV
field1, field2, image = test_plane(15, -19.1, -1.7, 0.4, 1.7)


#%% loading packages for image visualisation
from glob import glob
from scipy.ndimage import zoom, affine_transform
import matplotlib.pyplot as plt
from skimage.metrics import structural_similarity as ssim

#%%
stack1gfptif = glob(config.stack1gfp_image)
stack2gfptif = glob(config.stack2gfp_image)
stack1gfppl = load_plist(glob(config.stack1gfp_plist))
stack2gfppl = load_plist(glob(config.stack2gfp_plist))
def get_zstack(tif, plist, phase, target_zdim = config.zheight, order = 1):
    #extract zstack at a particular phase from the dataset
    #upsampled in the axial direction for consistent voxel aspect ratio
    zstack = []
    zloc = get_z_pos(plist)
    for z in zloc:
        image = closest_wall(tif, plist, z, 0.1, phase)
        zstack.append(image)
    
    vol = np.stack(zstack, axis=0)
    zdim, ydim, xdim = vol.shape
    #upsampling using linear interpolation
    zoom_factor = target_zdim / zdim
    vol_resampled = zoom(vol, (zoom_factor, 1, 1), order=order)

    return vol_resampled
def aff_transform(zstack, angle_deg, 
                                  tx=0, ty=0, tz=0, 
                                  order=1):
    """
    Rotate a z-stack around the y-axis and apply translation.
    
    Parameters
    ----------
    zstack : list of 2D numpy arrays
        Input z-stack (already resampled if needed).
    angle_deg : float
        Rotation angle around the y-axis.
    tx, ty, tz : float
        Translation in x, y, z directions (in voxels).
        Positive values shift the image toward +x, +y, +z.
    order : int
        Interpolation order for rotation/translation.
    
    Returns
    -------
    list of 2D numpy arrays
        Transformed z-stack.
    """
    vol = np.stack(zstack, axis=0)  # (Z, Y, X)
    zdim, ydim, xdim = vol.shape

    # --- Rotation matrix around y-axis ---
    theta = np.deg2rad(angle_deg)
    R = np.array([
        [ np.cos(theta), 0, np.sin(theta)],
        [ 0,             1, 0            ],
        [-np.sin(theta), 0, np.cos(theta)]
    ])

    # --- Center of rotation ---
    center = np.array([zdim/2, ydim/2, xdim/2])

    # --- Translation vector ---
    # Note: affine_transform expects translation in *input* coordinates
    T = np.array([tz, ty, tx])  # (z, y, x)

    # --- Compute final offset ---
    # x_in = R @ (x_out - center) + center - T
    offset = center - R @ center - T

    # --- Apply affine transform ---
    rotated = affine_transform(
        vol,
        R,
        offset=offset,
        output_shape=vol.shape,
        order=order,
        mode='constant',
        cval=0
    )

    return rotated
#%% generating upsampled z-stacks, performing affine transform on the z stack
zstack1 = get_zstack(stack1gfptif, stack1gfppl, 1.1)
zstack2 = get_zstack(stack2gfptif, stack2gfppl, 5.3)
original = zstack2[135]
reconstructed = aff_transform(zstack1, -19.1,-1.7, 0.4, 1.7)[135]
#%% generating coloured image to be displayed
overlay = np.clip((grey2rgb(original, 'm') + grey2rgb(reconstructed, 'g'))**1/0.5, 0, 65535)



#%% displaying images and the vector field overlay

x,y,u,v = field1
x2,y2,u2,v2 = field2
fig, ax = plt.subplots(dpi=300)
ax.imshow(overlay)
ax.quiver(x2,np.flipud(y2),u2*config.um_per_pixel/2,v2*config.um_per_pixel/2,scale=10,color='g',label='stack 1')
q = ax.quiver(x,np.flipud(y),u*config.um_per_pixel/2,v*config.um_per_pixel/2,scale=10,color='m',label='stack 2')

ax.quiverkey(q, X=0.1, Y=0.02, U=1, label='1 mm/s',color='white',labelcolor='white')
plt.legend()
SCALEBAR_LENGTH_PX = 46   # length in image pixels
SCALEBAR_LABEL = "20 µm"
img_h, img_w = overlay.shape[0:2]
margin = 10 #distance between scale bar to edge 
sb_x0 = img_w - margin - SCALEBAR_LENGTH_PX
sb_x1 = img_w - margin
sb_y  = img_h - margin
#positions of the scale bar, plotting scale bar
ax.plot([sb_x0, sb_x1], [sb_y, sb_y], color='white', linewidth=2)
ax.text((sb_x0 + sb_x1) / 2, sb_y - 5, SCALEBAR_LABEL,
           color='white', fontsize=8, ha='center', va='bottom')
plt.ylim(404,0) #prevent white stripe from showing due to vectors close to the edge.
plt.show()

#%% plotting the intersecting lines to illustrate the relative imaging planes involved in this MVF
L = 172 # length of lines - 400 pixels -> 172um
spacing = 4# separation between planes, 4um
n_lines = 31# number of planes
angle_deg = -19# sample rotation

# y positions centered around 0
ys = (np.arange(n_lines) - n_lines//2) * spacing

# line endpoints
x0, x1 = -L/2, L/2

# rotation matrix
theta = np.deg2rad(angle_deg)
R = np.array([[np.cos(theta), -np.sin(theta)],
              [np.sin(theta),  np.cos(theta)]])

fig, ax = plt.subplots(dpi=300)

# creating legend
ax.plot([], [], color='green', lw=2, label='stack 1')
ax.plot([], [], color='magenta', lw=2, label='stack 2')

# plotting horizontal stack 1 in green
for y in ys:
    ax.plot([x0, x1], [y, y], color='green', lw=1)

# plotting rotated stack 2 in magenta, looping through the verticle positions
for i, y in enumerate(ys):
    p1 = np.array([x0, y])
    p2 = np.array([x1, y])
    p1r = R @ p1
    p2r = R @ p2
    alpha = 1.0 if i == n_lines//2 else 0.3 #central line highlighted to show the actual plane involved from stack 2
    ax.plot([p1r[0], p2r[0]], [p1r[1], p2r[1]],
            color='magenta', lw=1, alpha=alpha)

# to make the aspect ratio correct, so the illustration has the correct
# scale showing the actual position of imaging planes
ax.set_aspect('equal', adjustable='box')
ax.set_xlim(-120, 120)
ax.set_ylim(-120, 120)

#only showing two axis lines 
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)


ax.set_xticks([])
ax.set_yticks([])

# XZ view, so horizontal is x, verticle is z
ax.set_xlabel("x", fontsize=14)
ax.set_ylabel("z", fontsize=14)

# legend
ax.legend(loc='upper right')

plt.show()






























