# -*- coding: utf-8 -*-
"""
Created on Tue Apr  9 13:12:02 2024.

@author: Jiawen Jiang

This script measures the rotation and shift between two zstacks.
It is used to manually register zstacks. Need to set the filepaths to the dataset,
phase values and also z-height. 

NOTE: this script is no longer in use as the method was updated to use the more
advanced, automated SSIM registration. But it is still kept as a tool to 
actually look at XZ cross sections, and producing cross section images
"""

import numpy as np
import plistlib as plib
import glob
import matplotlib.pyplot as plt
from tqdm import tqdm
from scipy import ndimage
from functions import *  
import config
#%%

#loading images and plists
stack1gfptif = glob.glob(config.stack1gfp_image)
stack1gfpp_list = glob.glob(config.stack1gfp_plist)

stack2gfptif = glob.glob(config.stack2gfp_image)
stack2gfpp_list = glob.glob(config.stack2gfp_plist)


stack1gfpplist = []
stack2gfpplist = []

stack1gfpplist = load_plist(stack1gfpp_list)
stack2gfpplist = load_plist(stack2gfpp_list)
    
#%%
#pixel size 0.552 um /0.69um for 10x magnification, 2x2 binning
#pixel width in x is 0.69*312 = 215.28 um
#height in z is 120 um, stretch to get the same aspect ratio
# h/312 = 120/215.28, h = 120*312/215.28 = 173.9 = 174 pixels
h=config.zheight #this is the height of zstack in unit pixels. Change for different dataset.  
phase1 = config.stack_phase_wall #phase value of stack1
phase2 = config.plane_phase_wall #phase value of stack2




z1=get_z_pos(stack1gfpplist)
z2=get_z_pos(stack2gfpplist)
stack1_xz = []
stack2_xz = []

stack1_mip = []
stack2_mip = []
for z in tqdm(z1):
    image = closest_wall(stack1gfptif, stack1gfpplist, z, 0.1, phase1)
    row = get_y_cross_section(image)
    stack1_xz.append(row)
    row = get_y_projection(image)
    stack1_mip.append(row)

for z in tqdm(z2):
    image = closest_wall(stack2gfptif, stack2gfpplist, z, 0.1, phase2)
    row = get_y_cross_section(image)
    stack2_xz.append(row)
    row = get_y_projection(image)
    stack2_mip.append(row)

stack1_xz = np.array(stack1_xz)
stack2_xz = np.array(stack2_xz)
stack1_mip = np.array(stack1_mip)
stack2_mip = np.array(stack2_mip)    

from skimage.transform import resize 

size = (h, len(stack1_xz[0]))
resized_stack1 = resize(stack1_xz, size)
resized_stack2 = resize(stack2_xz, size)
resized_stack1_mip = resize(stack1_mip, size)
resized_stack2_mip = resize(stack2_mip, size)
#%% saving cross sections, if needed

#from PIL import Image

#img = Image.fromarray(resized_stack1)
#img2 = Image.fromarray(resized_stack2)
#img.save("stack1_cross_section.tif")
#img2.save("stack2_cross_section.tif")

#%%
color_stack1 = grey2rgb(resized_stack1,'g')
color_stack2 = grey2rgb(resized_stack2,'m')
color_stack1_mip = grey2rgb(resized_stack1_mip, 'g')
color_stack2_mip = grey2rgb(resized_stack2_mip, 'm')

from matplotlib.widgets import Slider, Button
from scipy.ndimage import rotate

def normalize(img):
    #normalise images to between 0-1 to avoid clipping warning from matplotlib. 
    img = img.astype(np.float32)
    return img / np.max(img)

color_stack1 = normalize(color_stack1)
color_stack2 = normalize(color_stack2)
color_stack1_mip = normalize(color_stack1_mip)
color_stack2_mip = normalize(color_stack2_mip)

#initial settings
init_angle = -18
init_shift_x = -4
init_shift_y = -1
init_brightness1 = 2.0
init_brightness2 = 2.0

#whether to display cross section or MIP
use_mip = False

fig, axes = plt.subplots(1,2,dpi=200)
plt.subplots_adjust(left=0.25, bottom=0.45)
ax = axes[1]
axes[0].set_ylabel('z')
fig.text(0.55, 0.37, 'x')
#display images
red_img = ax.imshow(color_stack1 * init_brightness1, cmap='Reds', alpha=0.5, vmin=0, vmax=1, origin='lower')
green_img = ax.imshow(color_stack2 * init_brightness2, cmap='Greens', alpha=0.5, vmin=0, vmax=1, origin='lower')

fig.suptitle('XZ cross section of z-stacks')
ax.set_title("registered")
axes[0].set_title('unregistered')

red_img1 = axes[0].imshow(color_stack1 * init_brightness1, cmap='Reds', alpha=0.5, vmin=0, vmax=1, origin='lower')
green_img1 = axes[0].imshow(color_stack2 * init_brightness2, cmap='Greens', alpha=0.5, vmin=0, vmax=1, origin='lower')

#axes for sliders
ax_angle = plt.axes([0.25, 0.30, 0.65, 0.03])
ax_shift_x = plt.axes([0.25, 0.25, 0.65, 0.03])
ax_shift_y = plt.axes([0.25, 0.20, 0.65, 0.03])
ax_brightness1 = plt.axes([0.25, 0.15, 0.65, 0.03])
ax_brightness2 = plt.axes([0.25, 0.10, 0.65, 0.03])

#axes for buttons
ax_button = plt.axes([0.25, 0.02, 0.2, 0.05])
toggle_button = Button(ax_button, 'Toggle MIP')

#sliders
s_angle       = Slider(ax_angle, 'rotation/degree', -180, 180, valinit=init_angle)
s_shift_x     = Slider(ax_shift_x, 'shift x/pixel', -50, 50, valinit=init_shift_x)
s_shift_y     = Slider(ax_shift_y, 'shift z/pixel', -50, 50, valinit=init_shift_y)
s_brightness1 = Slider(ax_brightness1, 'brightness green', 0, 6.0, valinit=init_brightness1)
s_brightness2 = Slider(ax_brightness2, 'brightness magenta', 0, 6.0, valinit=init_brightness2)

def get_active_images():
    #check whether MIP or cross section should be displayed
    if use_mip:
        return color_stack1_mip, color_stack2_mip
    else:
        return color_stack1, color_stack2

#update plot
def update(val):
    angle = s_angle.val
    shift_x = int(s_shift_x.val)
    shift_y = int(s_shift_y.val)
    brightness1 = s_brightness1.val
    brightness2 = s_brightness2.val

    img1, img2 = get_active_images()

    #rotation
    rotated_green = rotate(img2, -angle, reshape=False, order=1) #use -angle to make it consistent with angle used to calculate intersection
    rotated_green = np.clip(rotated_green * brightness2, 0, 1) #contrast adjustment

    #shift
    shifted_red = np.roll(img1, shift=(shift_y, shift_x), axis=(0, 1))
    shifted_red = np.clip(shifted_red * brightness1, 0, 1) #contrast adjustment

    #update plot
    red_img.set_array(shifted_red)
    green_img.set_array(rotated_green)

    fig.canvas.draw_idle()

#toggle MIP display
def toggle(event):
    global use_mip
    use_mip = not use_mip
    if use_mip:
        ax.set_title("registered")
    else:
        ax.set_title("registered")
    update(None)

toggle_button.on_clicked(toggle)

# Connect sliders to update function
s_angle.on_changed(update)
s_shift_x.on_changed(update)
s_shift_y.on_changed(update)
s_brightness1.on_changed(update)
s_brightness2.on_changed(update)

plt.show()
    
    
    
    
    
    
    
    
    
    
    
    