import os
import numpy as np

def check_nplanes(track_ops):
    all_nplanes = []

    for ds_path in track_ops.all_ds_path:
        # check how many subfolders starting with plane* in suite2p folder
        n_planes = len([name for name in os.listdir(ds_path + '/suite2p') if name.startswith('plane')])
        print(f'Found {n_planes} planes in {ds_path}')
        all_nplanes.append(n_planes)
    track_ops.all_nplanes = all_nplanes
    # if all elements in all_n_planes are the same, then save it in track_ops.n_planes
    if all_nplanes.count(all_nplanes[0]) == len(all_nplanes):
        track_ops.nplanes = all_nplanes[0]
        print(f'Found {track_ops.nplanes} planes in all datasets')
    else:
        print('Found different number of planes in different datasets')
        print('Please check your dataset paths')
        print('Exiting...')
        exit()


# loads mean images
def load_all_imgs(track_ops):
    all_ds_avg_ch1 = []
    all_ds_avg_ch2 = []
    all_ds_nchannels = []

    for ds_path in track_ops.all_ds_path:
        ds_nchannels = []
        ds_avg_ch1 = []
        ds_avg_ch2 = []

        for i in range(track_ops.nplanes):
            ops = np.load(ds_path + '/suite2p/plane' + str(i) + '/ops.npy', allow_pickle=True).item()
            nchannels = ops['nchannels']
            print('nchannels: ' + str(nchannels) + ' for plane ' + str(i) + ' in dataset ' + ds_path)
            ds_avg_ch1.append(ops['meanImg'])
            ds_avg_ch2.append(ops['meanImg_chan2']) if nchannels==2 else ds_avg_ch2.append(None)
            ds_nchannels.append(nchannels)

        all_ds_avg_ch1.append(ds_avg_ch1)
        all_ds_avg_ch2.append(ds_avg_ch2)
        all_ds_nchannels.append(ds_nchannels)

    track_ops.all_ds_avg_ch1 = all_ds_avg_ch1
    track_ops.all_ds_avg_ch2 = all_ds_avg_ch2
    track_ops.all_ds_nchannels = all_ds_nchannels

    # if all elements in all_ds_nchannels are the same, then print its fine otherwise exit
    if all_ds_nchannels.count(all_ds_nchannels[0]) == len(all_ds_nchannels):
        track_ops.nchannels = all_ds_nchannels[0][0]
        print(f'Found {track_ops.nchannels} channels in all datasets')
    else:    
        print('Found different number of channels in different datasets')
        print('Please check your dataset paths')
        print('Exiting...')
        exit()

    return all_ds_avg_ch1, all_ds_avg_ch2

def load_all_ds_stat_iscell(track_ops):
    all_ds_stat_iscell = []
    for (i, ds_path) in enumerate(track_ops.all_ds_path):
        ds_stat_iscell = []
        for j in range(track_ops.nplanes):
            stat = np.load(os.path.join(ds_path, 'suite2p', f'plane{j}', 'stat.npy'), allow_pickle=True)
            iscell = np.load(os.path.join(ds_path, 'suite2p', f'plane{j}', 'iscell.npy'), allow_pickle=True)
            if track_ops.iscell_thr==None:
                stat_iscell = stat[iscell[:,0]==1]
            else: 
                stat_iscell = stat[iscell[:,1]>track_ops.iscell_thr]
            ds_stat_iscell.append(stat_iscell)
        all_ds_stat_iscell.append(ds_stat_iscell)

    return all_ds_stat_iscell

def load_all_ds_ops(track_ops):
    all_ds_ops = []
    for ds_path in track_ops.all_ds_path:
        ds_ops = []
        for j in range(track_ops.nplanes):
            ops = np.load(os.path.join(ds_path, 'suite2p', f'plane{j}', 'ops.npy'), allow_pickle=True).item()
            ds_ops.append(ops)
        all_ds_ops.append(ds_ops)
    
    return all_ds_ops

def load_all_ds_mean_img(track_ops, ch=1):
    all_ds_ops = load_all_ds_ops(track_ops)
    all_ds_mean_img = [] 
    for ds_ops in all_ds_ops:
        ds_mean_img = []
        for ops in ds_ops:
            mean_img = ops['meanImg'] if ch==1 else ops['meanImg_chan2']
            ds_mean_img.append(mean_img)
        all_ds_mean_img.append(ds_mean_img)
        
    return all_ds_mean_img

def load_all_ds_centroids(all_ds_stat_iscell, track_ops):
    all_ds_centroids = []
    for i in range(len(track_ops.all_ds_path)):
        ds_centroids = []
        for stat_iscell in all_ds_stat_iscell[i]:
            centroids = []
            for roi_stat in stat_iscell:
                centroids.append(roi_stat['med'])
            ds_centroids.append(np.array(centroids))
        all_ds_centroids.append(ds_centroids)
        
    return all_ds_centroids
