from ij import IJ, ImagePlus, ImageStack
from ij.plugin.frame import RoiManager
from ij.plugin import Duplicator, Concatenator
from ij.gui import Roi, PointRoi, WaitForUserDialog

imp = IJ.getImage()
dimC, dimZ, dimT = imp.getDimensions()[2:]

if dimZ == 1:
    mode = 'T'
    dimStack = dimT
elif dimT == 1:
    mode = 'Z'
    dimStack = dimZ
else:
    raise Exception('Cant process films of stacks for now')

roi = imp.getRoi()
if roi is None or not isinstance(roi, PointRoi):
    msg = 'The selection is not a MultiPoint ROI'
    WaitForUserDialog('Error', msg)
    raise Exception(msg)

point_list = roi.getContainedPoints()
position_list = [None] * dimStack
roi_dimensions = None

for point, z in zip(point_list, [roi.getPointPosition(i) for i in range(len(point_list))]):
    imp.setSlice(z)
    if mode == 'T':
        stack_position = imp.getT() - 1
    elif mode == 'Z':
        stack_position = imp.getZ() - 1
    print(len(position_list), stack_position)
    position_list[stack_position] = (point.x, point.y)

final_position_list = [None] * dimStack
final_multipoint_roi = PointRoi()
reference_point = None
for stack_position in range(dimStack):
    if position_list[stack_position] is None:
        previous_frame_position = stack_position
        while position_list[previous_frame_position] is None:
            previous_frame_position -= 1
        next_frame_position = stack_position
        while position_list[next_frame_position] is None:
            next_frame_position += 1
        x = position_list[previous_frame_position][0] * (next_frame_position - stack_position) / (
                next_frame_position - previous_frame_position) \
            + position_list[next_frame_position][0] * (stack_position - previous_frame_position) / (
                    next_frame_position - previous_frame_position)
        y = position_list[previous_frame_position][1] * (next_frame_position - stack_position) / (
                next_frame_position - previous_frame_position) \
            + position_list[next_frame_position][1] * (stack_position - previous_frame_position) / (
                    next_frame_position - previous_frame_position)
        final_position_list[stack_position] = (x, y)
    else:
        if reference_point is None:
            reference_point = position_list[stack_position]

        final_position_list[stack_position] = position_list[stack_position]
        final_multipoint_roi.addPoint(reference_point[0], reference_point[1], stack_position + 1)

imp_crops = []
x0, y0 = final_position_list[0]
for stack_position in range(dimStack):
    if mode == 'Z':
        imp_slice = Duplicator().run(imp, 1, dimC, stack_position + 1, stack_position + 1, 1, 1)
    elif mode == 'T':
        imp_slice = Duplicator().run(imp, 1, dimC, 1, 1, stack_position + 1, stack_position + 1)
    x, y = final_position_list[stack_position]
    IJ.run(imp_slice, "Translate...", "x=" + str(x0 - x) + " y=" + str(y0 - y) + " interpolation=None slice")
    imp_crops.append(imp_slice)

final_imp = Concatenator.run(imp_crops)
final_imp.setRoi(final_multipoint_roi)
final_imp.show()