import numpy as np

def load_track_ops(track_ops_path):
    track_ops = np.load(track_ops_path + 'track_ops_postreg.npy', allow_pickle=True).item()
    return track_ops

def load_stat_ds_plane(track_ops_path, track_ops, plane_idx=0):
    stat = np.load(track_ops_path + f'/suite2p/plane{plane_idx}/stat.npy', allow_pickle=True)
    iscell = np.load(track_ops_path + f'/suite2p/plane{plane_idx}/iscell.npy', allow_pickle=True)
    # filter based on track_ops.iscell_thr
    len_stat_allcell = len(stat)
    
    if track_ops.iscell_thr==None:
        stat = stat[iscell[:,0] ==1]
    else:
        stat= stat[iscell[:,1]>track_ops.iscell_thr]
    len_stat_iscell = len(stat)
    print(f'Loading ROIs for plane{plane_idx} in dataset {track_ops_path.split("/")[-2]}')
    print(f'Chose {len_stat_iscell}/{len_stat_allcell} ROIs, based on s2p iscell threshold {track_ops.iscell_thr} (see track_ops.iscell_thr)')
    # make a stat_summary dictionary
    stat_summary = {
        'len_stat_allcell': len_stat_allcell,
        'len_stat_iscell': len_stat_iscell,
        'iscell_thr': track_ops.iscell_thr
    }
    return stat, stat_summary

def get_all_roi_array_from_stat(stat, track_ops):
    n_xpix = track_ops.all_ds_avg_ch1[0][0].shape[0]
    n_ypix = track_ops.all_ds_avg_ch1[0][0].shape[0]
    all_roi_array = np.zeros((n_xpix, n_ypix, len(stat)), bool)

    for i in range(len(stat)):

        roi_xpix = stat[i]['xpix']
        roi_ypix = stat[i]['ypix']
     
        # Convert the ROI coordinates into a grid of values
        roi_grid = np.zeros((n_xpix, n_ypix), dtype=np.float32)
        roi_grid[np.array(roi_ypix), np.array(roi_xpix)] = 1

        all_roi_array[:,:,i] = roi_grid
        
    return all_roi_array

