# -*- coding: utf-8 -*-
"""
Created on Mon Jan  5 13:11:15 2026

@author: Jiawen Jiang

This script makes the exploded stack diagram using zstack images
It is used in combination with intersection visualisation.py (which produced
                                                              a 2D illustration)
to produce a 3D illustration of how stacks intersect

"""


import tifffile
import matplotlib.pyplot as plt
import numpy as np
from functions import *
import glob
import config

fp = glob.glob(config.stack1gfp_image)
plfp = glob.glob(config.stack1gfp_plist)

fp2 = glob.glob(config.stack2gfp_image)
plfp2 = glob.glob(config.stack2gfp_plist)

gfp_plist = load_plist(plfp)
gfp_plist2 = load_plist(plfp2)

zvals = get_z_pos(gfp_plist)
zvals2 = get_z_pos(gfp_plist2)
zstack_full = []
for i in range(15, 31):  # extended range: original 12-18 plus more planes below

    image = closest_wall(fp, gfp_plist, zvals[i], 0.1, 1)
    zstack_full.append(image)

#%%
import numpy as np
import tifffile
from mayavi import mlab

# ------------------------------------------------------------
# Load and normalize zstack
# ------------------------------------------------------------
zstack = np.array(zstack_full).astype(np.float32)
zstack -= zstack.min()
zstack /= zstack.max()
Z, Y, X = zstack.shape

# Downsample for speed
factor = 1
z_small = zstack[:, ::factor, ::factor]
Zs, Ys, Xs = z_small.shape

# ------------------------------------------------------------
# Create Mayavi figure
# ------------------------------------------------------------
mlab.figure(size=(800, 800), bgcolor=(0.8, 0.8, 0.8))

# ------------------------------------------------------------
# Render z-stack as semi-transparent planes
# ------------------------------------------------------------
spacing = 8
alpha = 1.0
vmax_value = 0.2
scalebar_length = 46   # pixels; change this to adjust the scale bar length
scalebar_width  = 5    # pixels
scalebar_margin = 10   # pixels from bottom-right corner
scalebar_color  = (255, 255, 0, 255)  # RGBA yellow


def stamp_scalebar(image):
    """
    Stamp a sentinel value of 1.0 into the scale bar region of a normalised
    [0,1] float image.  LUT index 255 is later overwritten to the desired
    scale bar colour, so the sentinel maps to yellow (or whatever color is set).
    Image data is first clamped to 254/255 so that bright pixels never
    accidentally reach LUT index 255 and get coloured as the scale bar.
    """
    img = np.clip(image, 0, 254/255)   # reserve index 255 exclusively for the scale bar
    h, w = img.shape
    row_end   = h - scalebar_margin
    row_start = row_end - scalebar_width
    col_end   = w - scalebar_margin
    col_start = col_end - scalebar_length
    img[row_start:row_end, col_start:col_end] = 1.0
    return img


def apply_scalebar_color_to_lut(lut):
    """
    Reserve LUT index 255 (the sentinel stamped by stamp_scalebar) for the
    scale bar colour.  Returns a modified copy of the 256x4 uint8 LUT.
    """
    lut = lut.copy()
    lut[255] = list(scalebar_color)
    return lut


def add_border(Xs, Ys, z, color=(1, 0, 0), width=4.0):
    # rectangle corners
    x = np.array([0, Xs, Xs, 0, 0])
    y = np.array([0, 0, Ys, Ys, 0])
    z = np.full_like(x, z)

    mlab.plot3d(
        x, y, z,
        color=color,
        tube_radius=None,   # flat line, not a tube
        line_width=width
    )


def add_tilted_border(Xo, Yo, Zo, color=(0, 1, 0), width=4.0):
    # corners in meshgrid indexing
    corners = [
        (0, 0),
        (0, -1),
        (-1, -1),
        (-1, 0),
        (0, 0)   # close the loop
    ]

    x = np.array([Xo[i, j] for i, j in corners])
    y = np.array([Yo[i, j] for i, j in corners])
    z = np.array([Zo[i, j] for i, j in corners])

    mlab.plot3d(
        x, y, z,
        color=color,
        tube_radius=None,   # flat line
        line_width=width
    )


def make_green_lut():
    lut = np.zeros((256, 4), dtype=np.uint8)
    for i in range(256):
        t = i / 255.0
        lut[i] = [
            0,
            int(255 * t),   # G increases
            0,
            255
        ]
    return lut


def make_magenta_lut():
    lut = np.zeros((256, 4), dtype=np.uint8)
    for i in range(256):
        t = i / 255.0
        lut[i] = [
            int(255 * t),   # R increases
            0,              # G stays zero
            int(255 * t),   # B increases
            255             # full opacity
        ]
    return lut


for i, slice_2d in enumerate(z_small):
    z = i * spacing
    slice_2d = np.clip(slice_2d / vmax_value, 0, 1)
    slice_2d = stamp_scalebar(slice_2d)
    slice_fixed = slice_2d.T

    obj = mlab.imshow(
        slice_fixed,
        colormap='gray',
        opacity=alpha,
        extent=[0, Xs, 0, Ys, z, z]  # flat plane at z
    )
    add_border(Xs, Ys, z, color=(0, 1, 0), width=4.0)
    lut = apply_scalebar_color_to_lut(make_green_lut())
    obj.module_manager.scalar_lut_manager.lut.table = lut


# ------------------------------------------------------------
# Load and render multiple magenta planes (current + all below)
# ------------------------------------------------------------
theta = np.deg2rad(-20)   # rotation around y-axis
x_shift = -20             # horizontal offset for tilted planes
z_shift = 20              # vertical offset for top magenta plane

# Collect overlay images: zvals2[12] is the original plane, then 13, 14, ... are below
magenta_zindices = list(range(15, 31))  # index 12 = top plane, 13..23 = below

# Pre-load all overlay images
overlay_images = []
for zi in magenta_zindices:
    ov = closest_wall(fp2, gfp_plist2, zvals2[zi], 0.1, 5.4).astype(np.float32)
    ov -= ov.min()
    ov /= ov.max()
    overlay_images.append(ov)

# Use shape from first overlay to build meshgrid once
overlay_ref = overlay_images[0][::factor, ::factor]
H, W = overlay_ref.shape
n_planes = len(overlay_images)

# Meshgrid in image coordinates
u, v = np.meshgrid(np.arange(W), np.arange(H))

# Stack pivot = centre of the magenta stack in its own local space
x0 = Xs / 2
stack_depth = (n_planes - 1) * spacing   # total depth of the magenta stack
z0_stack = stack_depth / 2               # pivot at mid-depth of the stack

# Build each plane's pre-rotation coordinates:
#   x, y stay the same for every plane; z steps from 0 downward
# Then rotate the *whole* stack around its own centre (z0_stack)
for plane_idx, ov_img in enumerate(overlay_images):
    overlay_small = np.clip(ov_img[::factor, ::factor] / vmax_value, 0, 1)
    overlay_small = stamp_scalebar(overlay_small)

    # Local z position within the stack (0 = top, increasing downward)
    z_local = plane_idx * spacing

    # Centre, rotate around y-axis as a rigid body
    x_c = u - x0
    z_c = z_local - z0_stack          # centred on stack mid-depth

    x_rot = x_c * np.cos(theta) + z_c * np.sin(theta)
    z_rot = -x_c * np.sin(theta) + z_c * np.cos(theta)

    # Translate back, apply positional offsets
    Xo = x_rot + x0 + x_shift
    Yo = v
    Zo = z_rot + z0_stack + z_shift   # restore pivot, then shift into scene

    obj2 = mlab.mesh(
        Xo, Yo, Zo,
        scalars=overlay_small,
        colormap='bone',
        opacity=1
    )
    add_tilted_border(Xo, Yo, Zo, color=(1, 0, 1), width=4.0)
    lut = apply_scalebar_color_to_lut(make_magenta_lut())
    obj2.module_manager.scalar_lut_manager.lut.table = lut

mlab.view(azimuth=60, elevation=30, distance='auto')
fig = mlab.gcf()
fig.scene.camera.parallel_projection = False
mlab.show()