# -*- coding: utf-8 -*-
"""
Created on Fri Mar  8 17:00:04 2024.

@author: Jiawen Jiang
"""

"""
This script contains functions that calculates the intersection positions 
between two z-stacks.

"""
import matplotlib.pyplot as plt
import numpy as np
import config

def line_intersection(line1, line2):
    """
    Calculate the intersection point between two line segments.
    
    This function works by solving P+t.r=Q+u.s, where s and r are vectors, and
    P, Q are starting point of two line segments. t and s are length along each
    line segment that the intersection lies.

    Parameters
    ----------
    line1 : 2x2 matrix containing two endpoints of line segment in [x,z].
    line2 : 2x2 matrix containing two endpoints of line segment in [x,z].

    Returns
    -------
    intersect : intersection position in [x,z]

    """
    #vectors of two line segments
    r = [line1[1][0] - line1[0][0], line1[1][1] - line1[0][1]]
    s = [line2[1][0] - line2[0][0], line2[1][1] - line2[0][1]]
    qp = [line2[0][0] - line1[0][0], line2[0][1] - line1[0][1]]
    #take cross product of two vectors
    def cross(a,b):
        return a[0]*b[1] - a[1]*b[0]
    r_c_s = cross(r,s) #r cross s
    qp_c_r = cross(qp, r) #qp cross r
    qp_c_s = cross(qp, s) #qp cross s
    
    if r_c_s == 0:
        #lines do not intersect
        return None
    t = qp_c_s/r_c_s
    u = qp_c_r/r_c_s
    
    if 0 <= t <= 1 and 0 <= u <= 1:
        intersect = [line1[0][0] + t*r[0], line1[0][1] + t*r[1]]
        return intersect 
    else:
        #lines intersect but beyond the range of line segments
        return None
    
    
def calculate_midpoints(stack_midpoint, distance, num_segments):
    """
    Calculate a list of midpoints of imaging planes in a z-stack.

    Parameters
    ----------
    stack_midpoint : midpoint of the z-stack, in [x,z]
    distance : distance between imaging planes in um, normally set an integer
    num_segments : number of imaging planes, int. 

    Returns
    -------
    midpoints : list of midpoints. 

    """
    #calculate the z position of the first line's midpoint
    start_z = stack_midpoint[1] - (num_segments - 1) / 2 * distance

    # calculate the midpoints of the line segments
    midpoints = [(stack_midpoint[0], start_z + i * distance) for i in range(num_segments)]

    return midpoints


def calculate_positions(midpoints, length, theta):
    """
    Calculate line segments of imaging planes in z-stack.

    Parameters
    ----------
    midpoints : list of midpoints.
    length : x size of image in um.
    theta : rotation angle in degrees, this should be 0 for one stack.

    Returns
    -------
    lines : list of line segments.

    """
    #convert theta from degrees to radians
    theta_rad = np.radians(theta)

    #calculate the displacement of the endpoints from the midpoint
    dx = length / 2 * np.cos(theta_rad)
    dz = length / 2 * np.sin(theta_rad)

    #calculate the positions of the line segments
    lines = []
    for x, z in midpoints:
        x1 = x - dx
        z1 = z - dz
        x2 = x + dx
        z2 = z + dz
        lines.append(((x1, z1), (x2, z2)))

    return lines

def rotate_point(origin, point, angle):
    """
    Rotate a point around an origin by an angle.

    Parameters
    ----------
    origin : origin of rotation, in [x,z].
    point : point to be rotated, in [x,z].
    angle : rotation angle, in degrees.

    Returns
    -------
    rotated point in x, z

    """
    angle = np.radians(angle)
    ox, oz = origin
    px, pz = point

    qx = ox + np.cos(angle) * (px - ox) - np.sin(angle) * (pz - oz)
    qz = oz + np.sin(angle) * (px - ox) + np.cos(angle) * (pz - oz)

    return qx, qz

def calculate_rotated_midpoints(stack_midpoint, distance, num_segments, theta):
    """
    Calculate midpoints for a rotated stack.

    Parameters
    ----------
    stack_midpoint : a list of midpoints of imaging planes in a z-stack.
    distance : separation between imaging planes in um, normally set as integer.
    num_segments : number of imaging planes.
    theta : rotation angle in degrees.

    Returns
    -------
    rotated_midpoints : a list of midpoints of the rotated stack.

    """
    #calculate the unrotated midpoints
    start_z = stack_midpoint[1] - (num_segments - 1) / 2 * distance
    midpoints = [(stack_midpoint[0], start_z + i * distance) for i in range(num_segments)]

    

    #calculate the rotated midpoints
    rotated_midpoints = [rotate_point(stack_midpoint, point, theta) for point in midpoints]

    return rotated_midpoints



def stack_line_intersection(stack_midpoint, distance, num_segments, rotation, x_length_um, tilted_index, stack_shift=(0,0)):
    """
    Calculate intersection positions between one tilted imaging plane and a zstack.
    
    This function can now only calculate the intersection between one imaging 
    plane in the rotated stack with the horizontal stack - as in practice 
    we need to loop through each z plane anyway, so it was decided not to calculate
    all intersections between two stacks at once. 

    Parameters
    ----------
    stack_midpoint : midpoint of the horizontal z-stack, in [x,z].
    distance : separation between imaging planes, in um. 
    num_segments : number of imaging planes.
    rotation : rotation angle between two z-stacks, in degrees.
    x_length_um : length of the image in x, in um.
    tilted_index : index of the imaging plane of the rotated stack
    stack_shift : shift in the horizontal stack, in[x,z]. The default is (0,0).

    Returns
    -------
    [1] : a list of intersection positions of the horizontal stack.
    [2] : line segment of the tilted plane, 2x2 matrix containing two endpoints.
    [3] : line segments of the horizontal stack, list of 2x2 tuple each containing two endpoints.

    """
    #calculating relevant midpoints of stacks
    unshifted_mid = calculate_midpoints(stack_midpoint, distance, num_segments)
    unshifted_stack = calculate_positions(unshifted_mid, x_length_um, 0)
    midpoints = calculate_midpoints(np.add(stack_midpoint, stack_shift), distance, num_segments)
    rotated_midpoints = calculate_rotated_midpoints(stack_midpoint, distance, num_segments, rotation)
    stack1 = calculate_positions(midpoints, x_length_um, 0)
    stack2 = calculate_positions(rotated_midpoints, x_length_um, rotation)
    line2 = stack2[tilted_index]
    grid = []
    #calculating intersection between tilted line and every horizontal line
    for line1 in stack1:
        intersection = line_intersection(line1, line2)
        if intersection is not None:
            grid.append(intersection)
        #else:
            #grid.append(None)
    minus_shift = -np.array(stack_shift) 
    return np.add(grid, minus_shift), np.add(stack2[tilted_index], minus_shift), unshifted_stack

def distance_along_line(intersection, line):
    """
    Calculate the coordinate of intersection, in the prospective of the tilted line.

    Parameters
    ----------
    intersection : a list of intersection positions.
    line : line segment of the tilted line.

    Returns
    -------
    1D list of positions along the tilted line.

    """
    coord = []
    start = line[0]
    for point in intersection:
        dist = ((point[0]-start[0])**2+(point[1]-start[1])**2)**0.5
        coord.append(dist)
    if coord[-1] < coord[0]:
        return np.flipud(coord)
    else:
        return coord

def um_to_pixel(coord, um_per_pixel=config.um_per_pixel):
    """
    Convert length unit from um to pixel because that is required for PIV.

    Parameters
    ----------
    coord : TYPE
        DESCRIPTION.
    um_per_pixel : TYPE, optional
        DESCRIPTION. The default is 0.645.

    Returns
    -------
    new_coord : TYPE
        DESCRIPTION.

    """
    new_coord = []
    for num in coord:
        new_coord.append(round(num/um_per_pixel))
    return new_coord

def get_y_grid(start, stop, xdim, ydim):
    """
    Generate a grid of y positions for PIV, result can be directly used in quiver.
    
    Because stack rotation is y-invanriant, y positions can be set freely, use 
    bigger ydim for higher vector density. It in some ways increase y resolution.

    Parameters
    ----------
    start : starting y position, in pixel.
    stop : ending y position in pixel
    xdim : number of x positions in the grid, depends on number of intersections.
    ydim : number of y positions, can be freely set.

    Returns
    -------
    2x2 matrix of y coordinates.

    """
    ygrid = []
    yline = np.linspace(start, stop, ydim, endpoint=True)
    new_yline = []
    for y in yline:
        new_yline.append(round(y))
    yline = new_yline
    for i in range(0, xdim):
        ygrid.append(yline)
    
    ygrid = np.array(ygrid)
    return ygrid.T

def get_plane_x_grid(coord, ydim):
    """
    Generate a grid of x coordinates, for the tilted plane.

    Parameters
    ----------
    coord : a list of x positions of intersection points, in um, can use result
    from distance_along_line.
    ydim : number of y positions, need to be consistent with other functions.

    Returns
    -------
    2x2 matrix of x coordinates.

    """
    xgrid = []
    for i in range(0, ydim):
        xgrid.append(um_to_pixel(coord))
    return np.array(xgrid)

def get_stack_x_grid(intersection, ydim):
    """
    Generate a grid of x coordinates, for the horizontal stack.

    Parameters
    ----------
    intersection : a list of intersection points, can use result from stack_line_intersection
    ydim : number of y positions, need to be consistent with other functions.

    Returns
    -------
    xgrid : 2x2 matrix of x coordinates
    rounded_z : a list of z positions of imaging planes.

    """
    xgrid = []
    xline = np.array(intersection).T[0]
    zline = np.array(intersection).T[1]
    if xline[-1] - xline[0] < 0:
        xline = np.flipud(xline)
        zline = np.flipud(zline)
        #this makes sure x values increase from left to right
    for i in range(0, ydim):
        xgrid.append(um_to_pixel(xline))
    
    rounded_z = []
    for i in range(0, len(zline)):
        rounded_z.append(round(zline[i]))
    rounded_z = np.array(rounded_z)
    
    return xgrid, rounded_z
#%%
'''
#this section checks that intersection is being calculated correctly. 
xlength = config.xlength
um_per_pixel = config.um_per_pixel
xlength = xlength*um_per_pixel #in unit um
ylength = config.ylength #in unit pixel

tilt_index = 15 
midz = config.midz #midpoint of z positions
stack_phase = config.stack_phase #phase of stack1 gata
plane_phase = config.plane_phase #phase of stack2 gata
nplanes=config.nplanes #how many imaging planes
plane_sep = config.plane_sep #separation between imaging planes
rotation = config.rotation #rotation angle in degrees between two stacks
nofrows = config.nofrows #how many rows of vectors in y, could be changed to get different field resolution
stack_shift = config.stack_shift #shift in stack1 in [x,z] measured from spatial registration
y_shift = config.y_shift #shift in y direction of stack2
grid, tilt_line, stack = stack_line_intersection((xlength/2, midz), plane_sep, nplanes, rotation, xlength, tilt_index, [0,5])
    

# plot the lines and intersection points
fig, ax = plt.subplots()
for line in stack:
    ax.plot([line[0][0], line[1][0]], [line[0][1], line[1][1]], 'k-')

ax.plot([tilt_line[0][0], tilt_line[1][0]], [tilt_line[0][1], tilt_line[1][1]], 'k-')


for point in grid:    
    if point is not None:
        ax.plot(point[0], point[1], 'ro')
plt.show()
'''
