import itk
import numpy as np

def reg_img_elastix(ref_img, mov_img, track_ops):
    # convert to itk images
    ref_img_itk = itk.GetImageFromArray(ref_img)
    mov_img_itk = itk.GetImageFromArray(mov_img)

    # import default parameter map
    parameter_object = itk.ParameterObject.New()
    parameter_map = parameter_object.GetDefaultParameterMap(track_ops.transform_type)
    parameter_object.AddParameterMap(parameter_map)

    # call registration function
    mov_img_reg_itk, reg_params = itk.elastix_registration_method(
        ref_img_itk,
        mov_img_itk,
        parameter_object = parameter_object
    )

    # convert back to numpy array
    mov_img_reg = itk.GetArrayFromImage(mov_img_reg_itk)
    reg_params.SetParameter("FinalBSplineInterpolationOrder", "0")

    return mov_img_reg, reg_params

def itk_reg_roi(roi, reg_params):
    roi_itk = itk.GetImageFromArray(roi.astype(np.uint8))
    roi_itk_trans = itk.transformix_filter(roi_itk, reg_params)
    roi_trans = itk.GetArrayFromImage(roi_itk_trans)
    return roi_trans

def itk_reg_all_roi(all_roi, reg_params):
    all_roi_array_reg = np.zeros_like(all_roi)
    for i in range(all_roi.shape[2]):
        roi_array = all_roi[:,:,i]
        roi_array_reg = itk_reg_roi(roi_array, reg_params)
        all_roi_array_reg[:,:,i] = roi_array_reg
    return all_roi_array_reg