# -*- coding: utf-8 -*-
"""
Created on Wed Feb  4 10:43:13 2026

@author: Jiawen Jiang

This script loads in the calculated 3D vector field and visualise it in 3D 
using Mayavi.

"""
from glob import glob
import pickle
from natsort import natsorted
import numpy as np
from scipy.ndimage import zoom
from scipy.interpolate import griddata
from scipy.interpolate import interp1d
import config
from mayavi import mlab
from functions import *
import matplotlib.pyplot as plt
def load_field(fp):
    #load in vector field data and perform some processing
    #such as converting pixel/ms to mm/s
    #fp: folder containing all .pkl files
    xp3d=[]
    yp3d=[]
    zp3d=[]
    xf3d=[]
    yf3d=[]
    zf3d=[]
    dx3d=[]
    dy3d=[]
    dz3d=[]
    data = np.load(fp)
    for i in range(config.nplanes):
        plane = [data[f"plane_{i}_arr_{j}"] for j in range(9)]
            
        xp3d+=plane[0].flatten().tolist()
        yp3d+=plane[1].flatten().tolist()
        zp3d+=plane[2].flatten().tolist()
        #onverts unit of velocity to mm/s
        x_field_um = plane[3]*config.um_per_pixel/2
        y_field_um = plane[4]*config.um_per_pixel/2
        z_field_um = plane[5]*config.um_per_pixel/2
        #1.5ms between frames -> 1ms interpulse, 1ms exposure
        dx_field_um = plane[6]*config.um_per_pixel/2
        dy_field_um = plane[7]*config.um_per_pixel/2
        dz_field_um = plane[8]*config.um_per_pixel/2
        
        xf3d+=x_field_um.flatten().tolist()
        yf3d+=y_field_um.flatten().tolist()
        zf3d+=z_field_um.flatten().tolist()
        
        dx3d+=dx_field_um.flatten().tolist()
        dy3d+=dy_field_um.flatten().tolist()
        dz3d+=dz_field_um.flatten().tolist()
        #print(zp.shape)
    xf3d = np.array(xf3d)
    yf3d = np.array(yf3d)
    zf3d = np.array(zf3d)
    xf3d = np.nan_to_num(xf3d)
    yf3d = np.nan_to_num(yf3d)
    zf3d = np.nan_to_num(zf3d)
    dx3d = np.nan_to_num(dx3d)
    dy3d = np.nan_to_num(dy3d)
    dz3d = np.nan_to_num(dz3d)
    return [xp3d, yp3d, zp3d, xf3d, yf3d, zf3d, dx3d, dy3d, dz3d]

def orthogonalize(field_param, nx=16, error=False):
    """
    Orthogonalize a 3D vector field by interpolating only in x direction.
    field_param: 9 element list of the field at one phase, 3 pos, 3 vel, 3 error
    """
    x, y, z, xf, yf, zf, dx, dy, dz = field_param
    x = np.array(x)
    y = np.array(y)
    z = np.array(z)
    xf = np.array(xf)
    yf = np.array(yf)
    zf = np.array(zf)
    dx = np.array(dx)
    dy = np.array(dy)
    dz = np.array(dz)
    #loading vector field into individual arrays
    y_unique = np.unique(y)
    z_unique = np.unique(z)

    # Common orthogonal x grid
    x_min, x_max = np.min(x), np.max(x)
    x_grid = np.linspace(x_min, x_max, nx)

    # Output arrays
    Xg, Yg, Zg = np.meshgrid(x_grid, y_unique, z_unique, indexing='ij')
    U = np.full_like(Xg, np.nan, dtype=float)
    V = np.full_like(Xg, np.nan, dtype=float)
    W = np.full_like(Xg, np.nan, dtype=float)
    
    dU = np.full_like(Xg, np.nan, dtype=float)
    dV = np.full_like(Xg, np.nan, dtype=float)
    dW = np.full_like(Xg, np.nan, dtype=float)

    # Loop over each (y,z) plane
    for yi_idx, yi in enumerate(y_unique):
        for zi_idx, zi in enumerate(z_unique):
            mask = (y == yi) & (z == zi)
            if np.sum(mask) > 1:
                xs = x[mask]
                us = xf[mask]
                vs = yf[mask]
                ws = zf[mask]
                dus = dx[mask]
                dvs = dy[mask]
                dws = dz[mask]

                sort_idx = np.argsort(xs)
                xs, us, vs, ws, dus, dvs, dws = xs[sort_idx], us[sort_idx], vs[sort_idx], ws[sort_idx], dus[sort_idx], dvs[sort_idx], dws[sort_idx]

                f_u = interp1d(xs, us, kind='linear', bounds_error=False, fill_value=np.nan)
                f_v = interp1d(xs, vs, kind='linear', bounds_error=False, fill_value=np.nan)
                f_w = interp1d(xs, ws, kind='linear', bounds_error=False, fill_value=np.nan)
                f_du = interp1d(xs, dus, kind='linear', bounds_error=False, fill_value=np.nan)
                f_dv = interp1d(xs, dvs, kind='linear', bounds_error=False, fill_value=np.nan)
                f_dw = interp1d(xs, dws, kind='linear', bounds_error=False, fill_value=np.nan)

                # Assign directly using integer indices
                U[:, yi_idx, zi_idx] = f_u(x_grid)
                V[:, yi_idx, zi_idx] = f_v(x_grid)
                W[:, yi_idx, zi_idx] = f_w(x_grid)
                
                dU[:, yi_idx, zi_idx] = f_du(x_grid)
                dV[:, yi_idx, zi_idx] = f_dv(x_grid)
                dW[:, yi_idx, zi_idx] = f_dw(x_grid)
                
                
    if error: #if you want uncertainty as well
        return Xg, Yg, Zg/config.um_per_pixel, U, V, W, dU, dV, -dW
    else: 
        return Xg, Yg, Zg/config.um_per_pixel, U, V, W
#%% loading and orthogonalising flow field
fp = r'field_map_snapshot/stack1phase58.npz' #change 58 to 11 for the field at phase 1.1
flow_field = load_field(fp)
wall_phase = 5.8 #phase of the GFP wall - note this is given in terms of stack1's phase
                 #to be consistent with image fp, but this is actually applied to stack2
                 #with phase shift automatically applied
ortho_field = orthogonalize(flow_field)
#ortho_field_error = orthogonalize(flow_field)
#%% loading gfp images
stack2gfptif = glob(config.stack2gfp_image)
stack2gfpp_list = glob(config.stack2gfp_plist)
stack2plist = load_plist(stack2gfpp_list)
#%% locating z plane locations
stack2z = get_z_pos(stack2plist)

def get_zstack(phase, zvals):
    stack = []
    for z in zvals:
        image = closest_wall(stack2gfptif, stack2plist, z, 0.05, phase)
        stack.append(image)
    
    return stack
def resample_zstack(images, num_new_planes, method='linear'):
    """
    Interpolate a z-stack to a desired axial dimension.

    Parameters
    ----------
    images : list of np.ndarray
        List of 2D image arrays (all same shape).
    num_new_planes : int
        Desired number of planes in the resampled stack.
    method : str
        Interpolation order: 'nearest' (0), 'linear' (1), 'cubic' (3).

    Returns
    -------
    np.ndarray
        Resampled stack with shape (num_new_planes, ny, nx).
    """
    stack = np.stack(images, axis=0)  # shape (nz, ny, nx)
    nz, ny, nx = stack.shape

    # Compute zoom factors
    zoom_factors = (num_new_planes / nz, 1, 1)

    # Map method string to order
    order_map = {'nearest': 0, 'linear': 1, 'cubic': 3}
    order = order_map.get(method, 1)

    # Interpolate
    resampled_stack = zoom(stack, zoom_factors, order=order)
    return resampled_stack

def get_cross_sections(zstack, nx=200, ny=200, nz=100):
    xy = zstack[nz, :, :] 
    xz = zstack[:, ny, :] 
    yz = zstack[:, :, nx]
    return xy, xz, yz

def warp_phase(phase, shift = config.phase_shift):
    #applying phase shift and bringing it within 0-2pi
    return (phase + shift) % (2 * np.pi)

#%%

def make_inverted_gray_lut():
    #customising LUT to display the heart wall as black pixels on a white bg
    lut = np.zeros((256, 4), dtype=np.uint8)
    for i in range(256):
        val = 255 - i  # invert: low input → white, high input → black
        lut[i] = [val, val, val, 255]
    return lut
def get_roll():
    cam = fig.scene.renderer.active_camera
    # VTK stores the view-up vector; compute roll from it
    import numpy as np
    view_up = np.array(cam.view_up)
    right = np.cross(np.array(cam.direction_of_projection), view_up)
    right /= np.linalg.norm(right)
    roll = np.degrees(np.arctan2(view_up[1], right[1]))
    return roll
# --- User defined ---
SCALEBAR_LENGTH = 46   # in image/world pixels
SCALEBAR_THICKNESS = 5 # in image/world pixels
SCALEBAR_MARGIN = 10   # distance from edge

def add_image_scalebar(image, length, thickness, margin):
    """Paint a black stripe into the bottom-left corner of a 2D image array."""
    out = image.copy()
    out[margin:margin+thickness, margin:margin+length] = np.max(image)  # black in inverted gray LUT
    return out
#where do you want the cross sections on the 3D render
zpos = 140
xpos = 110
ypos = 111
fig = mlab.figure(bgcolor=(0,0,0))

# Enable orthographic projection
fig.scene.renderer.active_camera.parallel_projection = True
# Enable terrain mode that keep roll at 0, if required
#fig.scene.interactor.interactor_style = tvtk.InteractorStyleTerrain()
#generating upsampled z stack
zstack = get_zstack(warp_phase(wall_phase), stack2z)
expanded_zstack = resample_zstack(zstack, config.zheight)
#get xy cross section
xyplane = get_cross_sections(expanded_zstack, nz = zpos)[0]

xyplane = np.clip(xyplane*2, np.min(xyplane), np.max(xyplane))
xyplane = np.flipud(xyplane).T #enhancing contrast and sorting into correct orientation
xyplane = add_image_scalebar(xyplane, SCALEBAR_LENGTH, SCALEBAR_THICKNESS, SCALEBAR_MARGIN)
img = mlab.imshow(xyplane)
img.module_manager.scalar_lut_manager.lut.table = make_inverted_gray_lut()
img.actor.position=(xyplane.shape[1]*0.5, xyplane.shape[0]*0.5, zpos) #displaying cross section
img.actor.property.opacity = 1.0

x, y, z, u, v, w = ortho_field
q = mlab.quiver3d(x,y,z,u,v,-w) #displaying vector field
cb = mlab.colorbar(object=q, title="flow speed/mm s⁻¹", orientation="vertical", nb_labels=5)
cb.scalar_bar.label_text_property.font_size = 16   # tick labels
cb.scalar_bar.title_text_property.font_size = 16  # main label
cb.scalar_bar.unconstrained_font_size = True
q.module_manager.scalar_lut_manager.lut_mode = 'inferno' #colormap for vectors
q.module_manager.vector_lut_manager.lut_mode = 'inferno'
fig.scene.render() #flushing colormap to make sure it is displaying given color

#get yz cross section
yzplane = get_cross_sections(expanded_zstack, nx=xpos)[2]
#enhancing contrast and correcting orientation
yzplane = np.clip(yzplane*2, np.min(yzplane), np.max(yzplane))
yzplane = np.flipud(yzplane.T)
yzplane = add_image_scalebar(yzplane, SCALEBAR_LENGTH, SCALEBAR_THICKNESS, SCALEBAR_MARGIN)
Y, Z = yzplane.shape
yy, zz = np.meshgrid(np.arange(Y), np.arange(Z), indexing='ij')
xx = np.full_like(yy, xpos, dtype=float)
#as yz plane is orthogonal to xy plane, it is displayed as a mesh colored by the image data
img_obj = mlab.mesh(xx, yy, zz, scalars=yzplane)
img_obj.module_manager.scalar_lut_manager.lut.table = make_inverted_gray_lut()
img_obj.actor.property.opacity = 1.0

#get xz cross section
xzplane = get_cross_sections(expanded_zstack, ny=404-ypos)[1]
#enhancing contrast and correcting orientation
xzplane = np.clip(xzplane*2, np.min(xzplane), np.max(xzplane))
xzplane = xzplane.T
xzplane = add_image_scalebar(xzplane, SCALEBAR_LENGTH, SCALEBAR_THICKNESS, SCALEBAR_MARGIN)
X, Z = xzplane.shape
xx, zz = np.meshgrid(np.arange(X), np.arange(Z), indexing='ij')
yy = np.full_like(xx, ypos, dtype=float)
img_obj = mlab.mesh(xx, yy, zz, scalars=xzplane)
img_obj.module_manager.scalar_lut_manager.lut.table = make_inverted_gray_lut()
img_obj.actor.property.opacity = 1.0

#adding border in xy plane
H, W = xyplane.shape
z_xy = zpos

x0, x1 = 0, W
y0, y1 = 0, H

mlab.plot3d([x0, x1], [y0, y0], [z_xy, z_xy], color=(1,0,0), tube_radius=None)
mlab.plot3d([x0, x1], [y1, y1], [z_xy, z_xy], color=(1,0,0), tube_radius=None)
mlab.plot3d([x0, x0], [y0, y1], [z_xy, z_xy], color=(1,0,0), tube_radius=None)
mlab.plot3d([x1, x1], [y0, y1], [z_xy, z_xy], color=(1,0,0), tube_radius=None)

#adding border in yz plane
Y, Z = yzplane.shape
x_yz = xpos

y0, y1 = 0, Y
z0, z1 = 0, Z

mlab.plot3d([x_yz, x_yz], [y0, y1], [z0, z0], color=(0,1,0), tube_radius=None)
mlab.plot3d([x_yz, x_yz], [y0, y1], [z1, z1], color=(0,1,0), tube_radius=None)
mlab.plot3d([x_yz, x_yz], [y0, y0], [z0, z1], color=(0,1,0), tube_radius=None)
mlab.plot3d([x_yz, x_yz], [y1, y1], [z0, z1], color=(0,1,0), tube_radius=None)

#adding border in xz plane
X, Z = xzplane.shape
y_xz = ypos

x0, x1 = 0, X
z0, z1 = 0, Z

mlab.plot3d([x0, x1], [y_xz, y_xz], [z0, z0], color=(0,0,1), tube_radius=None)
mlab.plot3d([x0, x1], [y_xz, y_xz], [z1, z1], color=(0,0,1), tube_radius=None)
mlab.plot3d([x0, x0], [y_xz, y_xz], [z0, z1], color=(0,0,1), tube_radius=None)
mlab.plot3d([x1, x1], [y_xz, y_xz], [z0, z1], color=(0,0,1), tube_radius=None)

txt = mlab.text(0.2, 0.2, "", width=0.2)
txt.actor.text_property.font_size = 8
txt.actor.property.color = (1, 0, 0)

def update_text():
    #display the current viewing orientation such that figures
    #captured from 3D render is reproducible
    az, el, _, _ = mlab.view()
    roll = get_roll()
    txt.text = f"Az={az:.1f} El={el:.1f} Ro={roll:.1f}"

fig.scene.interactor.add_observer("InteractionEvent", lambda *args: update_text())
mlab.show()

#%% animating the rotation of the vector field.
# note the animation works by creating Mayavi render, capturing a frame
# and caching the frame locally, then compiling all cached frames into a video
# works for 16GB memory, but reduce fps or resolution if running into memory issues.
import numpy as np
from mayavi import mlab
import imageio

# -----------------------------------------------
# User settings
# -----------------------------------------------
seconds_per_rotation = 5.0   # seconds per cycle; total = 2× this
fps = 20
output_path = "rotation_animation.mp4"
resolution = (720,720)    # ← change this to adjust output resolution (width, height)

zpos = 140
xpos = 110
ypos = 111

fig = mlab.figure(bgcolor=(0,0,0), size=resolution)

zstack = get_zstack(warp_phase(wall_phase), stack2z)
expanded_zstack = resample_zstack(zstack, config.zheight)

# XY cross section
xyplane = get_cross_sections(expanded_zstack, nz=zpos)[0]
xyplane = np.clip(xyplane*2, np.min(xyplane), np.max(xyplane))
xyplane = np.flipud(xyplane).T
xyplane = add_image_scalebar(xyplane, SCALEBAR_LENGTH, SCALEBAR_THICKNESS, SCALEBAR_MARGIN)
img = mlab.imshow(xyplane)
img.module_manager.scalar_lut_manager.lut.table = make_inverted_gray_lut()
img.actor.position = (xyplane.shape[1]*0.5, xyplane.shape[0]*0.5, zpos)
img.actor.property.opacity = 1.0

# Vector field with inferno colormap and colorbar
x, y, z, u, v, w = ortho_field
q = mlab.quiver3d(x, y, z, u, v, -w)
cb = mlab.colorbar(object=q, title="flow speed/mm s⁻¹", orientation="vertical", nb_labels=5)
cb.scalar_bar.label_text_property.font_size = 16
cb.scalar_bar.title_text_property.font_size = 16
cb.scalar_bar.unconstrained_font_size = True
q.module_manager.scalar_lut_manager.lut_mode = 'inferno'
q.module_manager.vector_lut_manager.lut_mode = 'inferno'
fig.scene.render()

# YZ cross section
yzplane = get_cross_sections(expanded_zstack, nx=xpos)[2]
yzplane = np.clip(yzplane*2, np.min(yzplane), np.max(yzplane))
yzplane = np.flipud(yzplane.T)
yzplane = add_image_scalebar(yzplane, SCALEBAR_LENGTH, SCALEBAR_THICKNESS, SCALEBAR_MARGIN)
Y, Z = yzplane.shape
yy, zz = np.meshgrid(np.arange(Y), np.arange(Z), indexing='ij')
xx = np.full_like(yy, xpos, dtype=float)
img_obj = mlab.mesh(xx, yy, zz, scalars=yzplane)
img_obj.module_manager.scalar_lut_manager.lut.table = make_inverted_gray_lut()
img_obj.actor.property.opacity = 1.0

# XZ cross section
xzplane = get_cross_sections(expanded_zstack, ny=404-ypos)[1]
xzplane = np.clip(xzplane*2, np.min(xzplane), np.max(xzplane))
xzplane = xzplane.T
xzplane = add_image_scalebar(xzplane, SCALEBAR_LENGTH, SCALEBAR_THICKNESS, SCALEBAR_MARGIN)
X, Z = xzplane.shape
xx, zz = np.meshgrid(np.arange(X), np.arange(Z), indexing='ij')
yy = np.full_like(xx, ypos, dtype=float)
img_obj = mlab.mesh(xx, yy, zz, scalars=xzplane)
img_obj.module_manager.scalar_lut_manager.lut.table = make_inverted_gray_lut()
img_obj.actor.property.opacity = 1.0

# Borders
H, W = xyplane.shape
x0, x1, y0, y1 = 0, W, 0, H
mlab.plot3d([x0,x1],[y0,y0],[zpos,zpos], color=(1,0,0), tube_radius=None)
mlab.plot3d([x0,x1],[y1,y1],[zpos,zpos], color=(1,0,0), tube_radius=None)
mlab.plot3d([x0,x0],[y0,y1],[zpos,zpos], color=(1,0,0), tube_radius=None)
mlab.plot3d([x1,x1],[y0,y1],[zpos,zpos], color=(1,0,0), tube_radius=None)

Y, Z = yzplane.shape
y0, y1, z0, z1 = 0, Y, 0, Z
mlab.plot3d([xpos,xpos],[y0,y1],[z0,z0], color=(0,1,0), tube_radius=None)
mlab.plot3d([xpos,xpos],[y0,y1],[z1,z1], color=(0,1,0), tube_radius=None)
mlab.plot3d([xpos,xpos],[y0,y0],[z0,z1], color=(0,1,0), tube_radius=None)
mlab.plot3d([xpos,xpos],[y1,y1],[z0,z1], color=(0,1,0), tube_radius=None)

X, Z = xzplane.shape
x0, x1, z0, z1 = 0, X, 0, Z
mlab.plot3d([x0,x1],[ypos,ypos],[z0,z0], color=(0,0,1), tube_radius=None)
mlab.plot3d([x0,x1],[ypos,ypos],[z1,z1], color=(0,0,1), tube_radius=None)
mlab.plot3d([x0,x0],[ypos,ypos],[z0,z1], color=(0,0,1), tube_radius=None)
mlab.plot3d([x1,x1],[ypos,ypos],[z0,z1], color=(0,0,1), tube_radius=None)

# -----------------------------------------------
# Axes orientation indicator
#
# Mayavi's built-in orientation axes widget is a
# VTK vtkOrientationMarkerWidget — it lives in
# screen space so it stays in the corner and
# automatically rotates with the camera.
# -----------------------------------------------
scene = fig.scene
axes = mlab.orientation_axes()


# -----------------------------------------------
# Camera helper
# -----------------------------------------------
mlab.view(azimuth=0, elevation=0)
fig.scene.render()

camera   = fig.scene.camera
focal    = np.array(camera.focal_point)
distance = camera.distance

def set_camera(pos, up):
    camera.position       = focal + pos
    camera.focal_point    = focal
    camera.view_up        = up
    camera.clipping_range = (distance * 0.01, distance * 10)
    fig.scene.render()

def capture_frame():
    # mlab.screenshot honours the figure's size= dimensions
    return mlab.screenshot(antialiased=True)

start_pos = np.array([0.0, 0.0,  distance])
start_up  = np.array([0.0, 1.0,  0.0])

# -----------------------------------------------
# Build frames (held in memory)
# -----------------------------------------------
frames_per_cycle = int(seconds_per_rotation * fps)
frames = []

def capture_and_store(pos, up):
    set_camera(pos, up)
    frames.append(capture_frame())

# Rotation 1: around the Y-axis (camera orbits in XZ plane)
for i in range(frames_per_cycle):
    theta = 2 * np.pi * i / frames_per_cycle
    pos = np.array([ distance * np.sin(theta),
                     0.0,
                     distance * np.cos(theta)])
    up  = np.array([0.0, 1.0, 0.0])
    capture_and_store(pos, up)

# Rotation 2: around the X-axis (camera orbits in YZ plane)
for i in range(frames_per_cycle):
    phi = 2 * np.pi * i / frames_per_cycle
    pos = np.array([ 0.0,
                    -distance * np.sin(phi),
                     distance * np.cos(phi)])
    up  = np.array([ 0.0,
                     np.cos(phi),
                     np.sin(phi)])
    capture_and_store(pos, up)

# -----------------------------------------------
# Write MP4 directly from memory
# -----------------------------------------------
writer = imageio.get_writer(output_path, fps=fps, codec="libx264",
                             quality=8, pixelformat="yuv420p")
for frame in frames:
    writer.append_data(frame)
writer.close()

print(f"Saved {output_path}  ({len(frames)} frames, {len(frames)/fps:.1f}s @ {fps} fps)")

set_camera(start_pos, start_up)
mlab.show()