# -*- coding: utf-8 -*-
"""
Created on Wed Jul 17 14:34:34 2024.

@author: Jiawen Jiang

This script calculates the intersections between two zstacks, and perform correlation
averaged PIV at intersection positions, generating a 3C-3D vector field. Each given phase
of the heartbeat generates a single .npz file, which contains 9 entries: three positional
arrays, three velocity arrays and three uncertainty arrays.

Note it saves filenames with phase value*10 to avoid decimal point in filenames
E.g. phase 5.8 -> phase58

"""

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
#%%loading image file paths and plist files
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)
#%% loading configurations for the dataset
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
stack_phase = 5.8
plane_phase = stack_phase+config.phase_shift #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
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_sample = config.bootstrap_samplesize
bs_iteration = config.bootstrap_iteration
phase_shift = config.phase_shift
#%% generating 3D field
whole_field_information = [] #list to collect field maps on all planes
for p in tqdm(range(0, nplanes)): 
    tilt_index = int(p) #go through each plane from stack2
    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
    plane_y = get_y_grid(0, ylength, len(stack_x[0]), nofrows)
    stack_y = plane_y - y_shift #applying delta y
    #stack_y = get_y_grid(0, ylength, len(stack_x[0]), nofrows)
    stack_u = []
    stack_v = []
    stack_uerr = []
    stack_verr = []
    for k in range(0, len(stack_z)):
        first_frames, second_frames = select_image(stack1tif, stack1plist, stack_z[k], 0.05, stack_phase, 2*np.pi/120, 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]
            if len(first_frames) == 0:
                u = 0
                v = 0
                uerr=0
                verr=0
            else:
                u, v, uerr, verr = correlate_at_pos(first_frames, second_frames, xpos, ypos, SW_size,IW_size, SNR_p2m_threshold, SNR_p2p_threshold, bs_iter=bs_iteration)
                #perform PIV at intersect
            urow.append(u)
            vrow.append(v)
            uerrrow.append(uerr)
            verrrow.append(verr)
            
        stack_u.append(urow)
        stack_v.append(vrow)
        stack_uerr.append(uerrrow)
        stack_verr.append(verrrow)
        
    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, stack_phase+phase_shift, 2*np.pi/120, limit=25, condition = config.pair_selection_condition_stack2)
    #loading and shifting frame pairs from stack 2
    plane_x_line = distance_along_line(intersections, tilt_line)
    plane_x = get_plane_x_grid(plane_x_line, nofrows)
    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 len(first_frames) == 0:
                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, bs_iter=bs_iteration)
                #uerr, verr = bootstrap(first_frames, second_frames, plane_x[i][j], plane_y[i][j], SW_size, IW_size, bs_sample, bs_iteration, SNR_p2p_threshold, SNR_p2m_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, stack_y, plane_u, plane_v)
    xperr, yperr, uperr, vperr = tools.transform_coordinates(plane_x, stack_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
    #create flow field grid
    z_field = np.zeros(stack2_x_field.shape)
    x_field = np.zeros(z_field.shape)
    y_field = np.zeros(z_field.shape)
    
    zerr_field = np.zeros(z_field.shape)
    xerr_field = np.zeros(z_field.shape)
    yerr_field = np.zeros(z_field.shape)
    for i in range(0, len(z_field)):
        for j in range(0, len(z_field[0])):
            if stack2_x_field[i][j] * stack1_x_field[i][j] != 0 and stack2_y_field[i][j]*stack1_y_field[i][j]>0:
                true_vec = calc_true(stack1_x_field[i][j], stack2_x_field[i][j], rotation)
                #calculate z component at every position
                z_field[i][j] = true_vec[1] #which is basically x velocity measured from stack 2
                y_field[i][j] = 0.5*(stack2_y_field[i][j]+stack1_y_field[i][j]) #average of two y velocities
                x_field[i][j] = true_vec[0]
                error_params = [stack1_x_field[i][j],stack2_x_field[i][j],config.rotation]
                errors = [userr[i][j], uperr[i][j], 0.1] #error in rotation is 0.1 - as optimisation only go down to .1 precision
                error_in_x, error_in_z = calculate_zerr(error_params, errors)
                zerr_field[i][j] = error_in_z
                xerr_field[i][j] = error_in_x
                yerr_field[i][j] = 0.5*(np.sqrt(vserr[i][j]**2 + vperr[i][j]**2))#average of two y uncertainties

    zp=np.zeros(xp.shape) + tilt_index*plane_sep
    #create position grid - since the 3d field is in the coord. system of stack2
    #z grid would simply be increasing at fixed interval = plane separation
    field_at_plane = [xp,yp,zp,x_field,y_field,z_field,xerr_field,yerr_field,zerr_field]
    whole_field_information.append(field_at_plane) #array containing all field information
    #old code to save a lot of pkl files, not ideal
    #path = 'field_map/stack1phase' + str(stack_phase) + '/'
    #ifExist = os.path.exists(path)
    #if not ifExist:
    #    os.makedirs(path)
    #first_path = path + str(tilt_index) + '.pkl'
    #
    #with open(first_path, 'wb') as file:
    #    pickle.dump(field_at_plane, file)
    #file.close()
    #saving the field map information, labelled by stack1 phase.
#%% saving the field information
path = 'field_map_snapshot/'
if not os.path.exists(path):
    os.makedirs(path)
save_dict = {}
for i, plane in enumerate(whole_field_information):
    for j, arr in enumerate(plane):
        save_dict[f"plane_{i}_arr_{j}"] = arr
#array 1,2,3 : position grid in x,y,z
#array 4,5,6 : velocity grid in x,y,z
#array 7,8,9 : error grid in x,y,z
filename = 'field_map_snapshot/stack1phase' + str(int(stack_phase*10)) + '.npz'
np.savez(filename, **save_dict)



'''

to recover the saved field map:

data = np.load(filename)

all_planes = []
for i in range(nplanes):
    field_at_plane = [data[f"plane_{i}_arr_{j}"] for j in range(9)]
    all_planes.append(field_at_plane)

'''
































