import numpy as np


def fast_voxel_intersect(start, end, origin, step, shape) -> tuple[list, list, list]:
    """Compute the voxels intersected by a line segment.

    Args:
        start (array-like): start point of line segment
        end (array-like): end point of line segment
        origin (array-like): origin of voxel grid
        step (array-like): step size of voxel grid

    Returns:
        list: list of intersection points
        list: list of intersected voxels
        list: list of voxels idx
    """

    # Convert to numpy arrays
    start = np.asarray(start)
    end = np.asarray(end)
    origin = np.asarray(origin)
    step = np.asarray(step)

    # Translate line segment to voxel grid
    start = start - origin
    end = end - origin

    # Initialize list of intersected voxels
    intersections = []
    voxels = []
    voxels_idx = []

    # Compute direction of line segment
    direction = end - start
    global_distance = np.linalg.norm(direction, axis=0)
    if global_distance == 0:
        return intersections, voxels, voxels_idx
    direction = direction / global_distance
    non_zero_direction = (direction != 0).argmax()

    # Compute the sign of the direction
    direction_signs = np.sign(direction)
    is_negative = direction_signs < 0

    # Initialize current position to start
    position = start.copy()

    # Main loop
    while True:
        # Compute the distance to the next boundaries
        next_boundaries = np.divide(position + step * direction_signs, step)
        distances = ((1 - is_negative) * np.floor(next_boundaries) +
                     is_negative * np.ceil(next_boundaries)) * step - position

        # Determine the nearest boundary to be reached
        boundary_distances = np.abs(distances / direction)
        clothest_boundary = np.argmin(boundary_distances)
        clothest_boundary_distance = boundary_distances[clothest_boundary]

        # Check if we are done
        distance_to_end = abs((end[non_zero_direction] - position[non_zero_direction]) / direction[non_zero_direction])
        if clothest_boundary_distance > distance_to_end:
            break

        # Update position
        position = position + clothest_boundary_distance * direction

        # Correct position to be on boundary
        position[clothest_boundary] = round(
            position[clothest_boundary] / step[clothest_boundary]) * step[clothest_boundary]

        # Get corresponding voxel
        on_boundary = np.mod(position, step) == 0
        voxel = np.floor(position) - is_negative * on_boundary * step

        # Add voxel to list
        idx = np.floor_divide(voxel, step).astype(int)
        if np.any(idx < 0) or np.any(idx >= shape):
            continue
        intersections.append(position + origin)
        voxels.append(voxel + origin)
        voxels_idx.append(idx)

    return intersections, voxels, voxels_idx


if __name__ == '__main__':
    import matplotlib.pyplot as plt

    def update_figure():
        positions, voxels, voxels_idx = fast_voxel_intersect(start, end, origin, step, shape)

        plt.clf()

        # Plot hitted voxels
        for voxel in voxels:
            plt.fill([voxel[0], voxel[0] + step[0], voxel[0] + step[0], voxel[0]],
                     [voxel[1], voxel[1], voxel[1] + step[1], voxel[1] + step[1]],
                     color='#e25', alpha=0.5)
        
        for voxel_id in voxels_idx:
            plt.fill([
                        origin[0] + voxel_id[0] * step[0], origin[0] + (voxel_id[0] + 1) * step[0],
                        origin[0] + (voxel_id[0] + 1) * step[0], origin[0] + voxel_id[0] * step[0]
                    ], [
                        origin[1] + voxel_id[1] * step[1], origin[1] + voxel_id[1] * step[1],
                        origin[1] + (voxel_id[1] + 1) * step[1], origin[1] + (voxel_id[1] + 1) * step[1]
                    ], color='#2e3', alpha=0.5)

        # Plot line segment
        plt.plot([start[0], end[0]], [start[1], end[1]], 'k-')
        plt.plot(start[0], start[1], 'go')
        plt.plot(end[0], end[1], 'ro')

        # Plot intersection points
        for pos in positions:
            plt.plot(pos[0], pos[1], 'bo')

        # Plot voxel grid
        plt.axis('equal')
        plt.xlim((-10, 10))
        plt.ylim((-10, 10))
        plt.xticks(origin[0] + step[0] * np.arange(shape[0] + 1))
        plt.yticks(origin[1] + step[1] * np.arange(shape[1] + 1))
        plt.grid()
        plt.draw()

    def onkey(event):
        global start, end
        if event.key == ' ':
            start = np.random.rand(2) * 20 - 10
            end = np.random.rand(2) * 20 - 10
            update_figure()

    # Define voxel grid
    origin = np.array([-5., -5.])
    step = np.array([1.0, 1.0])
    shape = (10, 10)

    # Define segment
    # start = np.random.rand(2) * 20 - 10
    # end = np.random.rand(2) * 20 - 10
    start = np.array([2.5, -3.5])
    end = np.array([7.0, 3.0])

    # Plot
    fig = plt.figure()
    fig.canvas.mpl_connect('key_press_event', onkey)
    update_figure()
    plt.show()