# -*- coding: utf-8 -*-
"""
Created on Mon Jan 26 10:05:53 2026

@author: Jiawen

This script plots the XZ cross sections of the two gfp stacks
and illustrate how imaging planes of two stacks intersect.
Cross section images are included in the directory, but could be 
generated using get_y_cross_section() from functions.py, and the zstack
images can be obtained from spatial registration.py

"""

import matplotlib.pyplot as plt
from skimage import io
from functions import *

im1 = io.imread('stack1_cross_section.tif')
im2 = io.imread('stack2_cross_section.tif')

im1 = grey2rgb(im1, 'g') ** 1/0.5
im2 = grey2rgb(im2, 'm') ** 1/0.5

import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import rotate

# Rotate im2 for display
im2_rot = rotate(im2, angle=20, reshape=False, order=1, mode='nearest')

# Image shapes
H1, W1 = im1.shape[:2]
H2, W2 = im2.shape[:2]

# Number of lines
num_lines = 31

# -----------------------------
# Compute horizontal lines for im1
# -----------------------------
ys1 = np.linspace(0, H1 - 1, num_lines)

# -----------------------------
# Compute rotated lines for im2
# -----------------------------
theta = np.deg2rad(20)

# Rotation matrix
R = np.array([[np.cos(theta), -np.sin(theta)],
              [np.sin(theta),  np.cos(theta)]])

# Center of im2
cx, cy = W2 / 2, H2 / 2

# Horizontal line y-positions in im2 BEFORE rotation
ys2 = np.linspace(0, H2 - 1, num_lines)

# For each line, compute rotated endpoints
rotated_lines = []
for y in ys2:
    p1 = np.array([0, y])
    p2 = np.array([W2, y])

    # shift to center
    p1c = p1 - np.array([cx, cy])
    p2c = p2 - np.array([cx, cy])

    # rotate
    p1r = R @ p1c + np.array([cx, cy])
    p2r = R @ p2c + np.array([cx, cy])

    rotated_lines.append((p1r, p2r))

intersection_points = []

for y in ys1:  # horizontal lines from im1
    for (p1, p2) in rotated_lines:

        x1, y1 = p1
        x2, y2 = p2

        # Skip if rotated line is horizontal (parallel)
        if abs(y2 - y1) < 1e-9:
            continue

        # Solve for t where the rotated line hits y
        t = (y - y1) / (y2 - y1)

        # --- FINITE SEGMENT CHECK ---
        if t < 0 or t > 1:
            continue  # intersection lies outside the segment

        # Compute intersection point
        x = x1 + t * (x2 - x1)

        # Optional: also check x is within image bounds
        if x < 0 or x > W1:
            continue

        intersection_points.append((x, y))
#%% plotting images and lines overlay
fig, ax = plt.subplots(figsize=(8,8),dpi=300)

# Plot im1

# Plot rotated im2
ax.imshow(im1 + im2_rot)

# Plot horizontal lines for im1
for y in ys1:
    ax.plot([0, W1], [y, y], color='green', linewidth=1)
ax.plot([0, W1], [y, y], color='green', linewidth=1, label='stack 1 imaging plane')
# Plot rotated lines for im2
for (p1, p2) in rotated_lines:
    ax.plot([p1[0], p2[0]], [p1[1], p2[1]], color='magenta', linewidth=1)
ax.plot([p1[0], p2[0]], [p1[1], p2[1]], color='magenta', linewidth=1, label='stack 2 imaging plane')
for (x, y) in intersection_points:
    ax.plot(x, y, 'ro', markersize=2)
ax.plot(x, y, 'ro', markersize=2, label='intersection')
plt.legend()

SCALEBAR_LENGTH_PX = 46   # length in image pixels
SCALEBAR_LABEL = "20 µm"
img_h, img_w = im1.shape[0:2]
margin = 10
sb_x0 = img_w - margin - SCALEBAR_LENGTH_PX
sb_x1 = img_w - margin
sb_y  = img_h - margin

ax.plot([sb_x0, sb_x1], [sb_y, sb_y], color='white', linewidth=2)
ax.text((sb_x0 + sb_x1) / 2, sb_y - 5, SCALEBAR_LABEL,
           color='white', fontsize=8, ha='center', va='bottom')

ax.set_xlim(0, W1)
ax.set_ylim(H1, 0)
plt.xlabel('x/pixel')
plt.ylabel('z/pixel')
#ax.axis('off')

plt.show()