# -*- coding: utf-8 -*-
"""
Created on Mon Apr  6 14:51:06 2026

@author: Jiawen Jiang

This script performs registration using SSIM between stacks.
It works by transforming stack 1, and calculating SSIM with stack 2.
The function has 4 spatial parameters as input, and use Bayesian optimiser to 
find the optimum. Phase shift is independently optimised first because SSIM 
is extremely sensitive to phase shift, and we only need to use rough estimate 
for spatial parameters (or just keep them at 0, the phase SSIM curve would 
                        still have a sharp, distinct peak)

"""

import config
from functions import *
from glob import glob
import numpy as np
from scipy.ndimage import zoom, affine_transform
import matplotlib.pyplot as plt
from skimage.metrics import structural_similarity as ssim

#%% loading images and plists
stack1tif = glob(config.stack1gfp_image)
stack2tif = glob(config.stack2gfp_image)
stack1pl = load_plist(glob(config.stack1gfp_plist))
stack2pl = load_plist(glob(config.stack2gfp_plist))

#%% build z-stack
def get_zstack(tif, plist, phase, target_zdim = config.zheight, order = 1):
    #returns a zstack at a particular phase
    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

    zoom_factor = target_zdim / zdim
    vol_resampled = zoom(vol, (zoom_factor, 1, 1), order=order)

    return vol_resampled

def get_full_heartbeat(tif, plist, phase_interval = 0.1, target_zdim = config.zheight, order = 1):
    #returns zstacks for full heartbeat, with given interval
    phase_stacks = []
    phases = np.arange(0.1, 2*np.pi, step=phase_interval)
    zloc = get_z_pos(plist)
    for phase in phases:
        zstack = []
        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

        zoom_factor = target_zdim / zdim
        vol_resampled = zoom(vol, (zoom_factor, 1, 1), order=order)
        phase_stacks.append(vol_resampled)
    return phase_stacks
#%% applying affine transform to rotate the z-stack
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

def aff_transform_3rotation(zstack,
                  angle_x_deg=0,
                  angle_y_deg=0,
                  angle_z_deg=0,
                  tx=0, ty=0, tz=0,
                  order=1):
    """
    Apply 3D affine transform to a z-stack:
    rotations around x, y, z axes + translation.

    Parameters
    ----------
    zstack : list of 2D numpy arrays
        Input z-stack (Z, Y, X).
    angle_x_deg, angle_y_deg, angle_z_deg : float
        Rotation angles around x, y, z axes (degrees).
    tx, ty, tz : float
        Translation in x, y, z directions (in voxels).
        Positive shifts move the volume toward +x, +y, +z.
    order : int
        Interpolation order.

    Returns
    -------
    list of 2D numpy arrays
        Transformed z-stack.
    """

    vol = np.stack(zstack, axis=0)  # (Z, Y, X)
    zdim, ydim, xdim = vol.shape

    # --- Convert to radians ---
    ax = np.deg2rad(angle_x_deg)
    ay = np.deg2rad(angle_y_deg)
    az = np.deg2rad(angle_z_deg)

    # --- Rotation matrices ---
    Rx = np.array([
        [1,          0,           0],
        [0,  np.cos(ax), -np.sin(ax)],
        [0,  np.sin(ax),  np.cos(ax)]
    ])

    Ry = np.array([
        [ np.cos(ay), 0, np.sin(ay)],
        [ 0,          1,         0 ],
        [-np.sin(ay), 0, np.cos(ay)]
    ])

    Rz = np.array([
        [ np.cos(az), -np.sin(az), 0],
        [ np.sin(az),  np.cos(az), 0],
        [ 0,           0,          1]
    ])

    # --- Combined rotation (Rz * Ry * Rx) ---
    R = Rz @ Ry @ Rx

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

    # --- Translation vector (z, y, x) ---
    T = np.array([tz, ty, tx])

    # --- 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


def find_valid_zrange(zstack, pad_value=0):
    """
    Given a rotated z-stack (list of 2D arrays), find the z-range
    where no padding was applied.

    Returns
    -------
    zmin, zmax : int
        Inclusive range of slices that contain no padding.
    """
    vol = np.stack(zstack, axis=0)  # (Z, Y, X)

    # Boolean mask: True if slice contains ANY padding
    slice_has_padding = np.any(vol == pad_value, axis=(1, 2))

    # Valid slices are those WITHOUT padding
    valid = ~slice_has_padding

    if not np.any(valid):
        return None, None  # everything padded

    # First and last valid slice
    zmin = np.argmax(valid)
    zmax = len(valid) - 1 - np.argmax(valid[::-1])

    return zmin, zmax


def patchwise_ssim(img1, img2, num_patches):
    #calculates patchwise ssim, was used for testing
    #but found no clear difference than simple ssim
    ny, nx = num_patches
    H, W = img1.shape[:2]

    ph = H // ny
    pw = W // nx

    off_y = (H - ph * ny) // 2
    off_x = (W - pw * nx) // 2

    vals = []
    for i in range(ny):
        for j in range(nx):
            y0 = off_y + i * ph
            x0 = off_x + j * pw
            vals.append(ssim(
                img1[y0:y0+ph, x0:x0+pw],
                img2[y0:y0+ph, x0:x0+pw],
                channel_axis=-1 if img1.ndim == 3 else None,
                data_range=img1.max() - img1.min()
            ))
    return np.mean(vals)

def get_ssim(stack1, stack2, patch_wise = False, patch_num = (0,0), chamber = 'all'):
    #calculates SSIM between given stacks, has the option to use patchwise approach
    #and selecting specific heart chamber, default is whole heart.
    if patch_wise == False:
        sim = []
        for i in range(0, len(stack1)):
            if chamber == 'all':
                img1 = stack1[i][10:-10, 10:-10] #cropped down such that SSIM not calculated in paddded regions
                img2 = stack2[i][10:-10, 10:-10]
            elif chamber == 'ventricle':
                img1 = stack1[i][10:-10, 10:200]
                img2 = stack2[i][10:-10, 10:200]
            elif chamber == 'atrium':
                img1 = stack1[i][120:-10, 200:-10]
                img2 = stack2[i][120:-10, 200:-10]
            data_range = img1.max() - img1.min()
            score = ssim(img1, img2, data_range=data_range)
            sim.append(score)
        return np.mean(sim)
    else:
        sim = []
        for i in range(0, len(stack1)):
            img1 = stack1[i][10:-10, 10:-10] #cropped down such that SSIM not calculated in paddded regions
            img2 = stack2[i][10:-10, 10:-10]
            score = patchwise_ssim(img1, img2, patch_num)
            sim.append(score)
        return np.mean(sim)
def get_sad(stack1, stack2):
    #calculates the sum of absolute differences between two stacks
    #tested as an alternative to ssim but do not work very well.
    sad = []
    for i in range(0, len(stack1)):
        img1 = stack1[i]
        img2 = stack2[i]
        score = np.sum(np.abs(img1 - img2))
        sad.append(score)
    return np.mean(sad)

#%% constucting z-stacks at given phase
#phase values can be changed.
zstack1 = get_zstack(stack1tif, stack1pl, 1.1)
zstack2 = get_zstack(stack2tif, stack2pl, 5.3)
#zstackrot = rotate_zstack_yaxis(zstack1, -20)

#%% example of optimising rotation angle
from tqdm import tqdm
angles = np.arange(-23,-12,step = 0.2)
shifts = np.arange(-10,10,step=0.2)
similarity_angle = []
similarity_angle_p = []
best_angle = -18.790033457835186
best_dx = -6.75070879947233
best_dy = 3.364315202210409
best_dz = 0.6170800885745784
for i in tqdm(range(0, len(angles))):
    zstackrot = aff_transform(zstack1, angles[i], tx=best_dx, ty=best_dy, tz=best_dz)
    similarity_angle.append(get_ssim(zstackrot[120:160], zstack2[120:160]))
    #similarity_angle_p.append(get_ssim(zstackrot[120:160], zstack2[120:160], patch_wise=True, patch_num=(10,10)))
#optimal parameters found from the Bayesian optimiser below - for plotting
#four curves and checking whether the calculated optimum is correct

similarity_dx = []
similarity_dy = []
similarity_dz = []

similarity_dx_p = []
similarity_dy_p = []
similarity_dz_p = []
for i in tqdm(range(0, len(shifts))):
    zstackrot = aff_transform(zstack1, best_angle, tx = shifts[i], ty=best_dy, tz=best_dz)
    similarity_dx.append(get_ssim(zstackrot[120:160], zstack2[120:160]))
    #similarity_dx_p.append(get_ssim(zstackrot[120:160], zstack2[120:160], patch_wise=True, patch_num=(10,10)))
    
for i in tqdm(range(0, len(shifts))):
    zstackrot = aff_transform(zstack1, best_angle, tx=best_dx, ty = shifts[i], tz=best_dz)
    similarity_dy.append(get_ssim(zstackrot[120:160], zstack2[120:160]))
    #similarity_dy_p.append(get_ssim(zstackrot[120:160], zstack2[120:160], patch_wise=True, patch_num=(10,10)))

for i in tqdm(range(0, len(shifts))):
    zstackrot = aff_transform(zstack1, best_angle, tx=best_dx, ty = best_dy, tz = shifts[i])
    similarity_dz.append(get_ssim(zstackrot[120:160], zstack2[120:160]))
    #similarity_dz_p.append(get_ssim(zstackrot[120:160], zstack2[120:160], patch_wise=True, patch_num=(10,10)))

    
#%% plotting SSIM curves
fig, ax = plt.subplots(2,2, figsize=(12,10), dpi=500)
ax[0][0].plot(angles, similarity_angle, label='SSIM')
ax[0][0].axvline(best_angle, linestyle=':', color='r', label='BO')
#ax[0][0].plot(angles, similarity_angle_p, label='patch-wise SSIM')
ax[0][0].set_xlabel('rotation/degree')
ax[0][0].set_ylabel('SSIM')
ax[0][0].legend()

ax[0][1].plot(shifts, similarity_dx, label='SSIM')
ax[0][1].axvline(best_dx, linestyle=':', color='r', label='BO')
#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('SSIM')
ax[0][1].legend()

ax[1][0].plot(shifts, similarity_dy, label='SSIM')
ax[1][0].axvline(best_dy, linestyle=':', color='r', label='BO')
#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('SSIM')
ax[1][0].legend()

ax[1][1].plot(shifts, similarity_dz, label='SSIM')
ax[1][1].axvline(best_dz, linestyle=':', color='r', label='BO')
#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('SSIM')
ax[1][1].legend()

plt.show()

#%% optimising parameters

#defining optimiser functions that can take in parameter and return a single score
def test_ssim(rot, dx, dy, dz):
    zstackrot_ = aff_transform(zstack1, rot, tx = dx, ty = dy, tz = dz)
    score = get_ssim(zstackrot_[120:160], zstack2[120:160], chamber = 'atrium')
    return score

def test_ssim_3rot(rx, ry, rz, dx, dy, dz):
    zstackrot_ = aff_transform_3rotation(zstack1, rx, ry, rz, tx = dx, ty = dy, tz = dz)
    score = get_ssim(zstackrot_[120:160], zstack2[120:160])
    return score

#%% Bayesian optimisation using scikit-optimization
from skopt import gp_minimize
from skopt.space import Real

def objective(x):
    return -test_ssim(*x) #negative as gp_minimize finds the minimum 

space = [
    #Real(-5,5), #these commented variables are for 3 rotation optimisation, which is not needed 
    Real(-22, -12),
    #Real(-5,5), 
    Real(-10, 2),
    Real(-2, 8),
    Real(-10, 10)
]

res = gp_minimize(objective, space, n_calls=100) #minimum finding, change ncalls for no. iterations
#note iterations is not larger -> better, larger iterations may give worse result.
best_params = res.x #get optimised parameters, that can be fed back into previous section
#to get plots of SSIM curves
best_value = -res.fun
#%% test for convergence
#this code shows Bayesian is not convergent for 100 iterations
#but provides a good enough estimate at 100 iterations
iterations = np.arange(10,110,10).astype(int)
params = []
highest_ssim = []
for i in tqdm(iterations):
    res = gp_minimize(objective, space, n_calls=i)
    best_params = res.x
    best_value = -res.fun
    params.append(best_params)
    highest_ssim.append(best_value)

fig, ax = plt.subplots(1,2,dpi = 300)
ax[0].plot(iterations, np.array(params).T[3])
ax[1].plot(iterations, np.array(params).T[2])
plt.show()



#%% visualising registered images
original = zstack2[135][120:-10, 200:-10]
reconstructed = aff_transform(zstack1, best_params[0], best_params[1], best_params[2], best_params[3])[135][120:-10, 200:-10]
#%%
overlay = np.clip((grey2rgb(original, 'm') + grey2rgb(reconstructed, 'g'))**1/0.5, 0, 65535)

plt.imshow(overlay)
plt.show()

#%% phase registration - normally this is done before previous sections
#as it is apparent from previous sections we already know a best phase shift (1.1 stack1 and 5.3 stack2)
phase_shifts = np.arange(0, np.pi*2, 0.1)
similarity_phase = [] #similarity scores
for shift in tqdm(phase_shifts): #looping over all phase shifts
    phase = 0.1 #note this only compares all phase shifts for a specific phase of stack1
    #more rigorously this should be done for all phases, and get an average, but it will take
    #way too long, and not necessary considering the sensitivity of ssim to phase shift.
    zstack1 = get_zstack(stack1tif, stack1pl, phase)
    zstack2 = get_zstack(stack2tif, stack2pl, phase + shift)
    zstackrot = aff_transform(zstack1, best_params[0], best_params[1], best_params[2], best_params[3])
    score = get_ssim(zstackrot[120:160], zstack2[120:160])
    similarity_phase.append(score)

#%% visualising phase curve
fig, ax = plt.subplots(dpi=300)
ax.plot(phase_shifts, similarity_phase, '.')
ax.set_xlabel('phase shift/rad')
ax.set_ylabel('SSIM')
plt.show()
#the peak should be sharp enough that one can visually tell the peak
#parabolic peak fitting could be done, but that accuracy is not necessary
#%% calculating different BO at different phase of heartbeat
from skopt import gp_minimize
from skopt.space import Real
from tqdm import tqdm
phases = np.arange(0.1, np.pi*2, 0.2)
all_best_angle = []
all_best_dx = []
all_best_dy = []
all_best_dz = []
all_best_val = []
for i in tqdm(range(0, len(phases))): #loop over all phases
    zstack1 = get_zstack(stack1tif, stack1pl, phases[i])
    zstack2 = get_zstack(stack2tif, stack2pl, phases[i] + 4.2)#4.2 is the phase shift measured
    def test_ssim(rot, dx, dy, dz):
        zstackrot_ = aff_transform(zstack1, rot, tx = dx, ty = dy, tz = dz)
        score = get_ssim(zstackrot_[120:160], zstack2[120:160], chamber = 'atrium')
        return score
    def objective(x):
        return -test_ssim(*x)

    space = [
        Real(-22, -12),
        Real(-10, 10),
        Real(0, 20),
        Real(-10, 10)
    ]

    res = gp_minimize(objective, space, n_calls=100)

    best_params = res.x
    best_value = -res.fun
    all_best_angle.append(best_params[0])
    all_best_dx.append(best_params[1])
    all_best_dy.append(best_params[2])
    all_best_dz.append(best_params[3])
    all_best_val.append(best_value)

#%% plotting the distribution of all reg parameters, and the 
#std gives uncertainties of measured parameters.
def trim_std(x):
    #removing the extremes so they do not affect std calculation
    x = np.asarray(x)
    mean = np.mean(x)
    distances = np.abs(x - mean)
    ind = np.argsort(distances)[-2:]
    x_trimmed = np.delete(x, ind)

    return np.std(x_trimmed, ddof=1)

fig, ax = plt.subplots(2,2, figsize=(8,10))
ax[0][0].hist(all_best_angle, color='r')
ax[0][1].hist(all_best_dx, color='r')
ax[1][0].hist(all_best_dy, color='r')
ax[1][1].hist(all_best_dz, color='r')

ax[0][0].set_title('mean '+str(round(np.mean(all_best_angle),2)) + ' std '+ str(round(trim_std(all_best_angle),2)))
ax[0][1].set_title('mean '+str(round(np.mean(all_best_dx),2)) + ' std '+ str(round(trim_std(all_best_dx),2)))
ax[1][0].set_title('mean '+str(round(np.mean(all_best_dy),2)) + ' std '+ str(round(trim_std(all_best_dy),2)))
ax[1][1].set_title('mean '+str(round(np.mean(all_best_dz),2)) + ' std '+ str(round(trim_std(all_best_dz),2)))

ax[0][0].set_xlim(-22,-12)
ax[0][1].set_xlim(-10, 10)
ax[1][0].set_xlim(0,20)
ax[1][1].set_xlim(-10,10)

ax[0][0].set_xlabel('BO rotation/degree')
ax[0][1].set_xlabel('BO x shift/pixel')
ax[1][0].set_xlabel('BO y shift/pixel')
ax[1][1].set_xlabel('BO z shift/pixel')
fig.suptitle('SSIM optimised registration parameters at different phase in atrium')
plt.show()











