import numpy as np

def get_all_ds_img_for_reg(all_ds_avg_ch1, all_ds_avg_ch2, track_ops): # chooses which channel to use for registration and returns all_ref_img and all_mov_img (shifted by one day to always register to previous day)
    if track_ops.reg_chan==0:
        all_ds_avg = all_ds_avg_ch1
    elif track_ops.reg_chan==1:
        all_ds_avg = all_ds_avg_ch2
        print('WARNING: using anatomical channel for registration (this is not always available)')

    all_ds_ref_img = []
    all_ds_mov_img = []

    for i in range(len(track_ops.all_ds_path)-1):
        ds_ref_img = []
        ds_mov_img = []
        
        for j in range(track_ops.nplanes):
            ds_ref_img.append(all_ds_avg[i][j])
            ds_mov_img.append(all_ds_avg[i+1][j])

        all_ds_ref_img.append(ds_ref_img)
        all_ds_mov_img.append(ds_mov_img)

    track_ops.all_ds_ref_img = all_ds_ref_img
    track_ops.all_ds_mov_img = all_ds_mov_img

    return all_ds_ref_img, all_ds_mov_img


def get_ref_reg_inters(all_roi_array_ref, all_roi_array_nonref):
    # get the projection of all rois
    all_roi_array_ref_proj = np.sum(all_roi_array_ref, axis=2) > 0
    all_roi_array_nonref_proj = np.sum(all_roi_array_nonref, axis=2) > 0

    # now get the intersection of the reg and ref   
    all_roi_array_inters = np.logical_and(all_roi_array_ref_proj, all_roi_array_nonref_proj)

    # make rgb image of intersection
    ref_reg_inters = np.ones((all_roi_array_inters.shape[0], all_roi_array_inters.shape[1], 3))
    ref_reg_inters[:,:,0] 
    ref_reg_inters[:,:,1] -= all_roi_array_inters/6 # orange tint
    ref_reg_inters[:,:,2] -= all_roi_array_inters

    return ref_reg_inters

def get_all_ref_nonref_inters(all_ds_all_roi_array_ref, all_ds_all_roi_array_nonref, track_ops):
    all_ds_all_ref_nonref_inters = []
    for i in range(len(track_ops.all_ds_path)-1):
        ds_all_ref_nonref_inters = []
        for j in range(track_ops.nplanes):
            all_roi_array_ref = all_ds_all_roi_array_ref[i][j]
            all_roi_array_nonref = all_ds_all_roi_array_nonref[i][j]
            ref_nonref_inters = get_ref_reg_inters(all_roi_array_ref, all_roi_array_nonref)
            ds_all_ref_nonref_inters.append(ref_nonref_inters)
        all_ds_all_ref_nonref_inters.append(ds_all_ref_nonref_inters)
    return all_ds_all_ref_nonref_inters