# -*- coding: utf-8 -*-
"""
Created on Fri Nov 28 17:01:29 2025

@author: Jiawen

This script loads the field maps, and visualise the orthogonal projections
of the 3D-3C vector field.

"""


from mpl_toolkits.mplot3d import Axes3D
import glob
import matplotlib.pyplot as plt
from tqdm import tqdm
import pickle
import numpy as np
import config
from natsort import natsorted
#loading the field map
stack_phase = config.stack_phase
fp = r'field_map_snapshot/stack1phase58.npz' #change this to load another field file
data = np.load(fp)

x_stack = []
y_stack = []
z_stack = []

xf_stack = []
yf_stack = []
zf_stack = []

#loading vector field into relevant arrays
for i in range(config.nplanes):
    plane = [data[f"plane_{i}_arr_{j}"] for j in range(9)]
    x_stack.append(plane[0])
    y_stack.append(plane[1])
    z_stack.append(plane[2])
    xf_stack.append(plane[3])
    yf_stack.append(plane[4])
    zf_stack.append(plane[5])


#%%
import matplotlib
from matplotlib.colors import Normalize
from mayavi import mlab


xp3d=[]
yp3d=[]
zp3d=[]
xf3d=[]
yf3d=[]
zf3d=[]

#this removes all vectors with 0 velocity to improve visibility of 3D field.
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
    #2ms between frames -> 1ms interpulse, 1ms exposure
    
    xf3d+=x_field_um.flatten().tolist()
    yf3d+=y_field_um.flatten().tolist()
    zf3d+=z_field_um.flatten().tolist()
    #print(zp.shape)
xf3d = np.array(xf3d)
yf3d = np.array(yf3d)
zf3d = np.array(zf3d)
speed = np.sqrt(xf3d**2 + yf3d**2 + zf3d**2)
speed += 0.0001
xf3d = np.nan_to_num(xf3d)
yf3d = np.nan_to_num(yf3d)
zf3d = np.nan_to_num(zf3d)

xnon0 = []
ynon0 = []
znon0 = []
xfnon0 = []
yfnon0 = []
zfnon0 = []
#filtering measurements to remove 0 or invalid measurements so
#points do not obstruct the render
#however this interferes with field interpolation so do not use
#unless needed.
for i in range(0, len(speed)): 
    if speed[i]>0.0001:
        xnon0.append(xp3d[i])
        ynon0.append(yp3d[i])
        znon0.append(zp3d[i])
        xfnon0.append(xf3d[i])
        yfnon0.append(yf3d[i])
        zfnon0.append(-zf3d[i])
#field converted to the form of 1d array - where x, y, z are positions, 
#xf, yf, zf are velocity components. 6 corresponding elements from 6 arrays give 
#one vector. This was done because the array format used in plt.quiver() was
#giving problems with Mayavi.

#briefly visualise in mayavi, if needed.
#mlab.figure(bgcolor=(0,0,0),size=(800, 800)) 

#mlab.points3d(xnon0, ynon0, znon0)
#mlab.quiver3d(xp3d,yp3d,zp3d,xf3d,yf3d,zf3d)
#mlab.flow(xp3d,yp3d,zp3d,xf3d,yf3d,zf3d)
#plt.colorbar(quiver,ax)

#mlab.quiver3d(xnon0,ynon0,znon0,xfnon0,yfnon0,zfnon0)
#mlab.flow(xnon0,ynon0,znon0,xfnon0,yfnon0,zfnon0)
#mlab.show()
#%% UI that inspects the cross sections of the heart
import numpy as np
from scipy.ndimage import zoom
from scipy.interpolate import griddata
from scipy.interpolate import interp1d
from skimage.exposure import equalize_adapthist

from skimage import io
def resample_zstack(images, num_new_planes, method='linear'):
    """
    Interpolate a z-stack to a desired axial dimension using linear interpolation.

    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

zstack = io.imread('phase_10.0gfpstack.tif')
flip_y = []
for frame in zstack:
    flip_y.append(np.flipud(frame))
zstack = resample_zstack(flip_y, config.zheight)


def orthogonalize(x, y, z, xf, yf, zf, nx=50):
    """
    Orthogonalize a 3D vector field by interpolating only in x direction.
    """
    x = np.array(x)
    y = np.array(y)
    z = np.array(z)
    xf = np.array(xf)
    yf = np.array(yf)
    zf = np.array(zf)

    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)

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

                sort_idx = np.argsort(xs)
                xs, us, vs, ws = xs[sort_idx], us[sort_idx], vs[sort_idx], ws[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)

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

    return Xg, Yg, Zg, U, V, W


#orthogonalising vector field
X, Y, Z, U, V, W = orthogonalize(xp3d, yp3d, zp3d, xf3d, yf3d, zf3d, nx=16)
Z = Z/config.um_per_pixel #converting z from microns to pixels



#%% overlaying gfp image on vector field, interpolated
import numpy as np
import matplotlib.pyplot as plt

def plot_zstack(images, X, Y, Z, U, V, W, slice_idx, clamp=True, brightness=1, gamma=1, scalebar_px=23.2):
    """
    Overlay vector field arrows onto z-stack cross-sections.
    Slice indices [ix, iy, iz] map to physical coordinates via X, Y, Z,
    which are then rounded to integers and used to index the image stack.

    Parameters
    ----------
    images : list of 2D np.ndarray
        Z-stack (each entry is a 2D image for one z-plane). All images must share shape (ny, nx).
    X, Y, Z : np.ndarray
        Coordinate arrays (shape: (nx, ny, nz), using 'ij' indexing). Can be floats.
    U, V, W : np.ndarray
        Vector components on the same grid (shape: (nx, ny, nz)).
    slice_idx : (ix, iy, iz)
        Indices into X, Y, Z from which to read the physical positions.
        These positions are rounded to integers for indexing into the image stack.
    clamp : bool
        If True, out-of-range rounded indices are clamped to valid bounds instead of erroring.

    Notes
    -----
    - XY panel uses image at z_pos (images[z_pos]) with quiver (U,V) at [:, :, iz].
    - ZY panel uses column x_pos across all z (stack[:, :, x_pos].T) with quiver (W,V) at [ix, :, :].
    - XZ panel uses row y_pos across all z (stack[:, y_pos, :]) with quiver (U,W) at [:, iy, :].
    """
    ix, iy, iz = map(int, slice_idx)
    def add_scalebar(ax, x_end, y_bar, label):
        #add scale bar for images, label is hard coded as it will be the same
        #for all my figures, but can be conveniently made to vary
        #it works just by ploting a solid line segment on the figure
        ax.plot([x_end - scalebar_px, x_end], [y_bar, y_bar], color='black', lw=3, solid_capstyle='butt')
        ax.text(x_end - scalebar_px / 2, y_bar + 5, label, color='black', ha='center', va='bottom', fontsize=7)

    # Read physical positions from coordinate arrays at the given indices
    x_pos_f = X[ix, iy, iz]
    y_pos_f = Y[ix, iy, iz]
    z_pos_f = Z[ix, iy, iz]

    # Round to integer indices for the image stack and rows/columns
    x_pos = int(round(x_pos_f))
    y_pos = int(round(y_pos_f))
    z_pos = int(round(z_pos_f))

    # Build image stack (nz, ny, nx)
    stack = np.stack(images, axis=0)
    stack = np.clip((stack*brightness)**1/gamma, np.min(stack), np.max(stack))
    nz, ny, nx = stack.shape

    # Bounds checking / clamping
    def _clamp(val, hi):
        return min(max(val, 0), hi - 1)

    if clamp:
        x_pos = _clamp(x_pos, nx)
        y_pos = _clamp(y_pos, ny)
        z_pos = _clamp(z_pos, nz)
    else:
        if not (0 <= x_pos < nx and 0 <= y_pos < ny and 0 <= z_pos < nz):
            raise IndexError(
                f"Rounded indices out of range: x={x_pos} (0..{nx-1}), y={y_pos} (0..{ny-1}), z={z_pos} (0..{nz-1})"
            )

    fig, axes = plt.subplots(2, 2, figsize=(12, 12), dpi=300)

    # --- XY plane: image at z_pos with quiver at iz ---
    ax_xy = axes[0][0]
    ax_xy.imshow(stack[z_pos, :, :], cmap='gray_r', origin='lower')
    Xxy, Yxy, Uxy, Vxy = X[:, :, iz], Y[:, :, iz], U[:, :, iz], V[:, :, iz]
    q1 = ax_xy.quiver(Xxy, Yxy, Uxy, Vxy, color='red', angles='xy', scale_units='xy', scale=scale, width=width)
    #quiverkey shows the scale bar for vectors
    key1 = ax_xy.quiverkey(q1, X=60, Y=10, U=5, label='5 mm/s',coordinates='data')
    key1.text.set_color('red')
    ax_xy.set_title(f"XY: z={z_pos_f:.3f}")
    ax_xy.set_xlabel("X"); ax_xy.set_ylabel("Y"); ax_xy.set_aspect('equal')
    ax_xy.axhline(y=y_pos_f, color='b',linestyle='--')
    ax_xy.axvline(x=x_pos_f, color='g',linestyle='--')
    add_scalebar(ax_xy, 380, 15, '20 µm')
    ax_xy.set_ylim(0, 400)
    ax_xy.set_xlim(0, 400)

    # --- ZY plane: column x_pos across z, quiver at ix ---
    ax_zy = axes[0][1]
    ax_zy.imshow(stack[:, :, x_pos].T, cmap='gray_r', origin='lower', aspect=nz/ny)
    Zzy, Yzy, Wzy, Vzy = Z[ix, :, :], Y[ix, :, :], W[ix, :, :], V[ix, :, :]
    q2 = ax_zy.quiver(Zzy, Yzy, Wzy, Vzy, color='g', angles='xy', scale_units='xy', scale=scale, width=width)
    key2 = ax_zy.quiverkey(q2, X=60, Y=10, U=5, label='5 mm/s',coordinates='data')
    key2.text.set_color('green')
    add_scalebar(ax_zy, 260, 15, '20 µm')
    ax_zy.set_title(f"ZY: x={x_pos_f:.3f}")
    ax_zy.set_xlabel("Z"); ax_zy.set_ylabel("Y"); ax_zy.set_aspect('equal')
    ax_zy.set_ylim(0, 400)

    # --- XZ plane: row y_pos across z, quiver at iy ---
    ax_xz = axes[1][0]
    ax_xz.imshow(stack[:, y_pos, :], cmap='gray_r', origin='lower', aspect=nz/nx)
    Xxz, Zxz, Uxz, Wxz = X[:, iy, :], Z[:, iy, :], U[:, iy, :], W[:, iy, :]
    q3 = ax_xz.quiver(Xxz, Zxz, Uxz, Wxz, color='b', angles='xy', scale_units='xy', scale=scale, width=width)
    key3 = ax_xz.quiverkey(q3, X=60, Y=10, U=5, label='5 mm/s',coordinates='data')
    key3.text.set_color('b')
    add_scalebar(ax_xz, 380, 15, '20 µm')
    ax_xz.set_title(f"XZ: y={y_pos_f:.3f}")
    ax_xz.set_xlabel("X"); ax_xz.set_ylabel("Z"); ax_xz.set_aspect('equal')
    ax_xz.set_xlim(0,400)
    fig.delaxes(axes[1][1])
    #do not show the bottom right plot as it is empty
    plt.tight_layout()
    plt.show()

slice_idx = [4,7,15]  # indices for y, z, x planes, change for different cross sections
scale = 0.08 #scale of vectors, larger -> smaller vectors
width=0.005 #width of vectors

#plotting cross sections with vector field overlay
plot_zstack(zstack, X, Y, Z, U, V, -W, slice_idx, brightness=6, gamma = 0.8, scalebar_px=46.4)



