# -*- coding: utf-8 -*-
"""
Created on Wed Jan 14 11:19:17 2026

@author: Jiawen

This script tests for how registration parameters affect diagonal plot - 
and performs registration based on the reduced chi2 values.

"""

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 filepaths and plist
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
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'):
    #this function essentially performs MVF PIV with a selected stack 2 plane
    #defined by tilt index, where 15 is the central plane (0-30)
    #it tries to simplify input by using global parameters where possible, so it only 
    #takes 4 meaningful parameters (spatial transforms), making it easier for optimisation
    xshift *= config.um_per_pixel
    zshift *= config.um_per_pixel
    stack_shift = [xshift,zshift]
    #applying stack shift, and calculating intersection positions
    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_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
    #generating positional grids and applying delta y
    stack_u = []
    stack_v = []
    stack_uerr = []
    stack_verr = []
    # performs PIV only in specific chamber, if needed, to reduce computation time
    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:
                    #no frames selected for the given criterion
                    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)
                    #perform PIV at intersection
                urow.append(u)
                vrow.append(v)
                uerrrow.append(uerr)
                verrrow.append(verr)
        if lowerx < planexpos < upperx: 
            stack_u.append(urow)
            stack_v.append(vrow)
            stack_uerr.append(uerrrow)
            stack_verr.append(verrrow)
    #converting vector field to arrays and correct orientation 
    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:
                    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)
    #converting lists to arrays
    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]
        
        if y1*y2 != 0 and y1err*0 == 0 and y2err*0 == 0:
            stack1y.append(y1)
            stack2y.append(y2)
            stack1yerr.append(y1err)
            stack2yerr.append(y2err)
    #this function can also be used to give vector field measurements
    #but in this script we just need the y velocities
    #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)
def calc_chi2(stack1, stack2, stack1err, stack2err, base_error = 0.0):
    #calculates the reduced chi2 value, applying a base error
    #this base error comes from fundamental PIV uncertainty
    #and was experimentally measured using bead samples with known results
    #this was added because tiny velocity measurements (very common) could bias
    #chi2 value.
    #adding errors in quadrature
    s = np.sqrt((np.sqrt(stack1err**2+base_error**2))**2 + (np.sqrt(stack2err**2+base_error**2))**2)
    c = np.sum((stack1 - stack2)**2 / s**2)
    dof = len(stack1)
    return c/dof
#%% performing PIV and get y velocities needed for diagonal plot
#note the registration parameters here, by default, are already
#optimised ones 
stack1y, stack1yerr, stack2y, stack2yerr = test_plane(15, -19.1, -1.7, 0.4, 1.7)
print('reduced chi2=' + str(calc_chi2(stack1y,stack2y,stack1yerr,stack2yerr,base_error=0.055)))


#%% plotting the diagonal plot
xs = np.linspace(min(stack1y), max(stack1y), 50) * config.um_per_pixel/2
fig, ax = plt.subplots(dpi=300)
ax.plot(xs,xs,color='r',label='y=x')
#plotting velocities, converted to mm/s as PIV measurement gives result in pixel displacements
ax.plot(stack1y*config.um_per_pixel/2, stack2y*config.um_per_pixel/2, '.', color='b')
ax.errorbar(stack1y*config.um_per_pixel/2, stack2y*config.um_per_pixel/2, xerr=(np.abs(stack1yerr)+0.055)*config.um_per_pixel/2, yerr=(np.abs(stack2yerr)+0.055)*config.um_per_pixel/2, fmt=',', color='b')
ax.legend()
ax.set_xlabel(r"stack 1 y-velocity/$mms^{-1}$")
ax.set_ylabel(r"stack 2 y-velocity/$mms^{-1}$")
ax.grid()
ax.text(
    0.98, 0.2,
    r'$\chi^2_\nu = 1.79$',
    ha='right', va='top',
    transform=ax.transAxes
) #showing calculated chi2 value in the plot.

plt.show()



#%% calculating the variations in chi2 value with reg parameters

#defining range of varying parameters
rotation_angles = np.arange(-24,-13,1)
x_shifts = np.arange(-10,11,1)
y_shifts = np.arange(-10,11,1)
z_shifts = np.arange(-10,11,1)
#defining optimal values obtained from either SSIM or Gauss-Seidel
#note values here are already optimised, so they will produce chi2 curves with these minimum
best_angle = -19.1
best_dx = -1.7
best_dy = 0.4
best_dz = 1.7
#indices of planes to sample from, cannot be too much as it takes too long to run
tilt_indices = [10,15,20]
#%% calculating chi2 values with each parameter varying, and the other three constant
angle_chi = []
for angle in tqdm(rotation_angles):
    stack1ys = np.array([])
    stack2ys = np.array([])
    stack1yerrs = np.array([])
    stack2yerrs = np.array([])
    for tilt in tilt_indices:
        #collecting all measurements on three planes together for chi2 calculation
        stack1y, stack1yerr, stack2y, stack2yerr = test_plane(tilt, angle, best_dx, best_dy, best_dz,'ventricle')
        stack1ys = np.concatenate((stack1ys, stack1y))
        stack2ys = np.concatenate((stack2ys, stack2y))
        stack1yerrs = np.concatenate((stack1yerrs, stack1yerr))
        stack2yerrs = np.concatenate((stack2yerrs, stack2yerr))
        
    angle_chi.append(calc_chi2(stack1ys, stack2ys, stack1yerrs, stack2yerrs, base_error=0.055))

dx_chi = []
for dx in tqdm(x_shifts):
    stack1ys = np.array([])
    stack2ys = np.array([])
    stack1yerrs = np.array([])
    stack2yerrs = np.array([])
    for tilt in tilt_indices:
        stack1y, stack1yerr, stack2y, stack2yerr = test_plane(tilt, best_angle, dx, best_dy, best_dz,'ventricle')
        stack1ys = np.concatenate((stack1ys, stack1y))
        stack2ys = np.concatenate((stack2ys, stack2y))
        stack1yerrs = np.concatenate((stack1yerrs, stack1yerr))
        stack2yerrs = np.concatenate((stack2yerrs, stack2yerr))
    dx_chi.append(calc_chi2(stack1ys, stack2ys, stack1yerrs, stack2yerrs, base_error=0.055))

dy_chi = []
for dy in tqdm(y_shifts):
    stack1ys = np.array([])
    stack2ys = np.array([])
    stack1yerrs = np.array([])
    stack2yerrs = np.array([])
    for tilt in tilt_indices:
        stack1y, stack1yerr, stack2y, stack2yerr = test_plane(tilt, best_angle, best_dx, dy, best_dz,'ventricle')
        stack1ys = np.concatenate((stack1ys, stack1y))
        stack2ys = np.concatenate((stack2ys, stack2y))
        stack1yerrs = np.concatenate((stack1yerrs, stack1yerr))
        stack2yerrs = np.concatenate((stack2yerrs, stack2yerr))
    dy_chi.append(calc_chi2(stack1ys, stack2ys, stack1yerrs, stack2yerrs, base_error=0.055))

dz_chi = []
mean_error_z = []
for dz in tqdm(z_shifts):
    stack1ys = np.array([])
    stack2ys = np.array([])
    stack1yerrs = np.array([])
    stack2yerrs = np.array([])
    for tilt in tilt_indices:
        stack1y, stack1yerr, stack2y, stack2yerr = test_plane(tilt, best_angle, best_dx, best_dy, dz,'ventricle')
        stack1ys = np.concatenate((stack1ys, stack1y))
        stack2ys = np.concatenate((stack2ys, stack2y))
        stack1yerrs = np.concatenate((stack1yerrs, stack1yerr))
        stack2yerrs = np.concatenate((stack2yerrs, stack2yerr))
    dz_chi.append(calc_chi2(stack1ys, stack2ys, stack1yerrs, stack2yerrs, base_error=0.055))
    mean_error_z.append(0.5*(np.mean(stack1yerr) + np.mean(stack2yerr)))


#%% finding minima in chi2 curve using parabolic interpolation
def parabolic_minima(x, y):
    x = np.asarray(x, float)
    y = np.asarray(y, float)

    i = np.argmin(y)
    if i == 0 or i == len(y)-1:
        return x[i]

    f1, f2, f3 = y[i-1], y[i], y[i+1]
    d = f1 - 2*f2 + f3
    if d == 0:
        return x[i]

    dx = 0.5 * (f1 - f3) / d
    return x[i] + dx * (x[i+1] - x[i])

#
print(parabolic_minima(rotation_angles, angle_chi))
print(parabolic_minima(x_shifts, dx_chi))
print(parabolic_minima(y_shifts, dy_chi))
print(parabolic_minima(z_shifts, dz_chi))



#%% plotting the resulting chi2 curve and the initial ssim optimum
#if the chi2 optimum was used in the previous section, plots should show
#minima of chi2 curves perfectly aligned with defined reg parameters.
ssim_op = [-18.6, -4.6, 2.8, 2.1] #these are the optimum and their uncertainty from SSIM-optimisation
angle_std = 0.91
dx_std = 0.58
dy_std = 0.37
dz_std = 1.35
fig, ax = plt.subplots(2,2, figsize=(12,10), dpi=500)
ax[0][0].plot(rotation_angles, angle_chi, '.', label=r"$\chi^2$", markersize = 10)
ax[0][0].axvline(best_angle, linestyle='-', color='g', label='$\chi^2$ optimum')
ax[0][0].axvspan(ssim_op[0]-angle_std, ssim_op[0]+angle_std, color='m', alpha=0.3)
ax[0][0].axvline(ssim_op[0], linestyle='dotted', color='m', label='SSIM optimum')
#ax[0][0].plot(angles, similarity_angle_p, label='patch-wise SSIM')
ax[0][0].set_xlabel('rotation/degree')
ax[0][0].set_ylabel(r'reduced $\chi^2$')
ax[0][0].legend()

ax[0][1].plot(x_shifts, dx_chi, '.', label=r"$\chi^2$", markersize = 10)
ax[0][1].axvline(best_dx, linestyle='-', color='g', label='$\chi^2$ optimum')
ax[0][1].axvspan(ssim_op[1]-dx_std, ssim_op[1]+dx_std, color='m', alpha=0.3)
ax[0][1].axvline(ssim_op[1], linestyle='dotted', color='m', label='SSIM optimum')
#ax[0][1].plot(shifts, similarity_dx_p, label='patch-wise SSIM')
ax[0][1].set_xlabel('x shift/pixel')
ax[0][1].set_ylabel(r'reduced $\chi^2$')
ax[0][1].legend()

ax[1][0].plot(y_shifts, dy_chi, '.', label=r"$\chi^2$", markersize = 10)
ax[1][0].axvline(best_dy, linestyle='-', color='g', label='$\chi^2$ optimum')
ax[1][0].axvspan(ssim_op[2]-dy_std, ssim_op[2]+dy_std, color='m', alpha=0.3)
ax[1][0].axvline(ssim_op[2], linestyle='dotted', color='m', label='SSIM optimum')
#ax[1][0].plot(shifts, similarity_dy_p, label='patch-wise SSIM')
ax[1][0].set_xlabel('y shift/pixel')
ax[1][0].set_ylabel(r'reduced $\chi^2$')
ax[1][0].legend()

ax[1][1].plot(z_shifts, dz_chi, '.', label=r"$\chi^2$", markersize = 10)
ax[1][1].axvline(best_dz, linestyle='-', color='g', label='$\chi^2$ optimum')
ax[1][1].axvspan(ssim_op[3]-dz_std, ssim_op[3]+dz_std, color='m', alpha=0.3)
ax[1][1].axvline(ssim_op[3], linestyle='dotted', color='m', label='SSIM optimum')
#ax[1][1].plot(shifts, similarity_dz_p, label='patch-wise SSIM')
ax[1][1].set_xlabel('z shift/pixel')
ax[1][1].set_ylabel(r'reduced $\chi^2$')
ax[1][1].legend()

for axes in ax.flatten():
    axes.set_ylim(1.0, 3.0)


plt.show()
#%% Optimising chi2 curves using Gauss-Seidel method for simultaneous optimisation
def chi2function(rotation, dx, dy, dz):
    #defining a optimiser function to take in 4 variables and returning a single score
    #essentially a refined version of test_plane
    tilt_indices = [10,15,20]
    y1s,y2s,y1es,y2es = np.array([]),np.array([]),np.array([]),np.array([])
    for tilt in tilt_indices:
        r1, r1e, r2, r2e = test_plane(tilt, rotation, dx, dy, dz, chamber = 'ventricle')
        y1s = np.concatenate((y1s, r1))
        y2s = np.concatenate((y2s, r2))
        y1es = np.concatenate((y1es, r1e))
        y2es = np.concatenate((y2es, r2e))
    return(calc_chi2(y1s, y2s, y1es, y2es, base_error=0.055))

# -----------------------------
# USER INPUTS
# -----------------------------

# Initial parameters
params = np.array([-18.6, -4.6, 2.8, 2.1])

# Bounds for each variable (integer grid)
bounds_list = [(-24, -16), (-5, 5), (-5, 5), (-5, 5)]

# Number of full iterations
iterations = 5

# Storage for history
history = [params.copy()]


# -----------------------------
# PARABOLIC MINIMUM HELPER
# -----------------------------
def parabolic_minimum(x_vals, y_vals):
    #calculates minimum by local parabolic interpolation, around the given position
    x0, x1, x2 = x_vals
    y0, y1, y2 = y_vals

    numerator = (x2**2 - x1**2) * y0 + (x0**2 - x2**2) * y1 + (x1**2 - x0**2) * y2
    denominator = (x2 - x1) * y0 + (x0 - x2) * y1 + (x1 - x0) * y2

    if denominator == 0:
        return x1

    return 0.5 * numerator / denominator


# -----------------------------
# MAIN OPTIMISATION LOOP
# -----------------------------
for it in tqdm(range(iterations)):

    # Loop over each variable sequentially
    for var_index in range(4):

        low, high = bounds_list[var_index]
        grid = np.arange(low, high + 1)

        # Evaluate function on integer grid
        values = []
        for g in grid:
            p = params.copy()
            p[var_index] = g
            values.append(chi2function(*p))
        values = np.array(values)

        # Find lowest grid point - this give global minimum
        i = np.argmin(values)

        # Boundary case: no interpolation possible
        if i == 0 or i == len(grid) - 1:
            print('boundary, consider stopping and resetting bounds')
            best_x = grid[i]
        else:
            # Fit parabola using 3 points
            x_vals = grid[i-1:i+2]
            y_vals = values[i-1:i+2]
            best_x = parabolic_minimum(x_vals, y_vals)

        # Round to 1 decimal place - speed up convergence, and it is not necessary to
        # reach precision smaller than 0.1 pixel. If a higher precision is desired,
        # remove this line and the code will converge in more iterations to more decimal places
        best_x = np.round(best_x, 1)

        # Update parameter
        params[var_index] = best_x

    # Record after full sweep
    history.append(params.copy())
#output - history contains all parameter values in every iteration, can be used to check when the optimisation converged
#5 iterations were performed but convergence 
#resulting params can be fed back into previous sections to get the chi2 curves, with minima coinciding the calculated optimum as expected. 
print("Final parameters:", params)
print("History:")
for h in history:
    print(h)
