# for now the only algorithm is elastix, TODO: add other algorithms and run them in the same way within this loop

from track2p.register.elastix import reg_img_elastix
from track2p.io.loaders import load_stat_ds_plane, get_all_roi_array_from_stat
from track2p.register.elastix import itk_reg_all_roi

def run_reg_loop(all_ds_ref_img, all_ds_mov_img, track_ops):
    all_ds_mov_img_reg = []
    all_ds_reg_params = []

    for (i, ds_ref_img) in enumerate(all_ds_ref_img):
        ds_mov_img = all_ds_mov_img[i]
        ds_mov_img_reg = []
        ds_reg_params = []

        for j in range(track_ops.nplanes):
            ref_img = ds_ref_img[j]
            mov_img = ds_mov_img[j]
            mov_img_reg,reg_params = reg_img_elastix(ref_img, mov_img, track_ops)
            ds_mov_img_reg.append(mov_img_reg)
            ds_reg_params.append(reg_params)

        all_ds_mov_img_reg.append(ds_mov_img_reg)
        all_ds_reg_params.append(ds_reg_params)
        
    track_ops.all_ds_mov_img_reg = all_ds_mov_img_reg
    # track_ops.all_ds_reg_params = all_ds_reg_params # for now not saving elastix transform object (TODO: transforme it to serializable to be able to pickle (for example by saving params to a dictionary))
    
    return all_ds_mov_img_reg, all_ds_reg_params


def reg_all_ds_all_roi(all_ds_reg_params, track_ops):
    all_ds_all_roi_array_ref = []
    all_ds_all_roi_array_mov = []
    all_ds_all_roi_array_reg = []
    all_ds_roi_counter = [] # this will keep track of how many cells make it thorugh the registration

    for i in range(len(track_ops.all_ds_path)-1):
        print(f'...\nTransforming ROIs for registration {i}/{len(track_ops.all_ds_path)-1}')

        ds_all_roi_array_ref = [] # for one dataset (all planes)
        ds_all_roi_array_mov = []
        ds_all_roi_array_reg = []
        ds_roi_counter_ref = []
        ds_roi_counter_mov = []

        for j in range(track_ops.nplanes):

            # 1) Set paths and transformation
            ref_ds_path = track_ops.all_ds_path[i]
            reg_ds_path = track_ops.all_ds_path[i+1]

            reg_params = all_ds_reg_params[i][j]

            # 2) Load ROIs # TODO: add (non-s2p dependent) Cellpose ROIs compatibility
            stat_ref, roi_counter_ref = load_stat_ds_plane(ref_ds_path, track_ops, plane_idx=j) # loading rois for one dataset one plane
            stat_mov, roi_counter_mov = load_stat_ds_plane(reg_ds_path, track_ops, plane_idx=j) # loading rois for one dataset one plane

            all_roi_array_ref = get_all_roi_array_from_stat(stat_ref, track_ops) # for dataset i, plane j
            all_roi_array_mov = get_all_roi_array_from_stat(stat_mov, track_ops) # for dataset i, plane j

            # 3) Apply transformation
            all_roi_array_reg = itk_reg_all_roi(all_roi_array_mov, reg_params)

            # 4) append
            ds_all_roi_array_ref.append(all_roi_array_ref)
            ds_all_roi_array_mov.append(all_roi_array_mov)
            ds_all_roi_array_reg.append(all_roi_array_reg)

            ds_roi_counter_ref.append(roi_counter_ref)
            ds_roi_counter_mov.append(roi_counter_mov) # only keep this one if it is the last pair
        
        print('Done with dataset...')

        all_ds_all_roi_array_ref.append(ds_all_roi_array_ref)
        all_ds_all_roi_array_mov.append(ds_all_roi_array_mov)
        all_ds_all_roi_array_reg.append(ds_all_roi_array_reg)
        
        all_ds_roi_counter.append(ds_roi_counter_ref) # keep both if its last pair (last recording does't appear as a ref)
        if i == len(track_ops.all_ds_path)-2:
            all_ds_roi_counter.append(ds_roi_counter_mov)

        # save all to track ops
        track_ops.all_ds_all_roi_array_ref = all_ds_all_roi_array_ref
        track_ops.all_ds_all_roi_array_mov = all_ds_all_roi_array_mov
        track_ops.all_ds_all_roi_array_reg = all_ds_all_roi_array_reg
        track_ops.all_ds_roi_counter = all_ds_roi_counter

    return all_ds_all_roi_array_ref, all_ds_all_roi_array_mov, all_ds_all_roi_array_reg, all_ds_roi_counter