import numpy as np


def fast_voxel_intersect(start, end, origin, step) -> tuple[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
    """

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

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

    # Compute the sign of the direction
    direction_signs = np.sign(direction)
    is_positive = direction_signs > 0
    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 = (is_positive * 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[0] - position[0]) / direction[0])
        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
        boundary_reached_negativly = np.zeros(start.shape, dtype=bool)
        boundary_reached_negativly[clothest_boundary] = is_negative[clothest_boundary]
        voxel = np.floor(position) - boundary_reached_negativly * step

        # Add voxel to list
        intersections.append(position + origin)
        voxels.append(voxel + origin)

    return intersections, voxels


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

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

        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)

        # 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))
        xmin, xmax = plt.xlim()
        ymin, ymax = plt.ylim()
        plt.xticks(np.arange(xmin + origin[0],
                             xmax + origin[0] + step[0], step[0]))
        plt.yticks(np.arange(ymin + origin[1],
                             ymax + origin[1] + step[1], step[1]))
        plt.grid()
        plt.draw()

    def onclick(event):
        global start, end
        # if event.button == 1:
        #     start = np.array([event.xdata, event.ydata])
        # elif event.button == 3:
        #     end = np.array([event.xdata, event.ydata])
        start = np.random.rand(2) * 10 - 5
        end = np.random.rand(2) * 10 - 5
        update_figure()

    # Define voxel grid
    origin = np.array([.1, -.3])
    step = np.array([1.0, 1.0])

    # Define segment
    start = np.random.rand(2) * 10 - 5
    end = np.random.rand(2) * 10 - 5

    # Plot
    fig = plt.figure()
    fig.canvas.mpl_connect('button_press_event', onclick)
    update_figure()
    plt.show()