from track2p.ops.default import DefaultTrackOps
from types import SimpleNamespace

from track2p.io.s2p_loaders import load_all_imgs, check_nplanes, load_all_ds_stat_iscell, load_all_ds_mean_img, load_all_ds_centroids
from track2p.io.savers import save_track_ops, save_all_pl_match_mat

from track2p.register.loop import run_reg_loop, reg_all_ds_all_roi
from track2p.register.utils import get_all_ds_img_for_reg, get_all_ref_nonref_inters

from track2p.plot.progress import plot_all_planes
from track2p.plot.output import plot_reg_img_output, plot_thr_met_hist, plot_n_matched_roi, plot_roi_reg_output, plot_roi_match_multiplane, plot_allroi_match_multiplane

from track2p.match.loop import get_all_ds_assign, get_all_pl_match_mat 
import numpy as np
import os


def run_t2p(track_ops):

    # 1) initialise save paths for figures and matched neurons output
    track_ops.init_save_paths()

    # 2) Load data
    check_nplanes(track_ops)
    all_ds_avg_ch1, all_ds_avg_ch2 = load_all_imgs(track_ops)

    # 3) Plot available planes for registration
    plot_all_planes(all_ds_avg_ch1, track_ops)
    if track_ops.nchannels==2:
        plot_all_planes(all_ds_avg_ch2, track_ops, ch='anatomical')

    # 4) do the actual registration based on chosen channel
    all_ds_ref_img, all_ds_mov_img = get_all_ds_img_for_reg(all_ds_avg_ch1, all_ds_avg_ch2, track_ops)

    all_ds_mov_img_reg, all_ds_reg_params = run_reg_loop(all_ds_ref_img, all_ds_mov_img, track_ops) # TODO: save basic parameters for each registration as feedback (e. g. ammoung of shift, rotation, etc.) for later plotting

    plot_reg_img_output(track_ops)
    

    # 5) apply computed transorm to all ROIs
    all_ds_all_roi_ref, all_ds_all_roi_mov, all_ds_all_roi_reg, all_ds_roi_counter = reg_all_ds_all_roi(all_ds_reg_params, track_ops)
 

    # 6) optional: generate 'yellow intersection' plot (this is only needed for plotting below)
    all_ds_ref_reg_inters = get_all_ref_nonref_inters(all_ds_all_roi_ref, all_ds_all_roi_reg, track_ops)

    all_ds_ref_mov_inters = get_all_ref_nonref_inters(all_ds_all_roi_ref, all_ds_all_roi_mov, track_ops)


    track_ops.all_ds_ref_mov_inters = all_ds_ref_mov_inters
    track_ops.all_ds_ref_reg_inters = all_ds_ref_reg_inters

    # this line is very memory-intensive because of the ROIS (TODO: maybe instead of contours just plot RGB) (or somehow generate RGB image of contours (in the part before))
    if track_ops.show_roi_reg_output:
        plot_roi_reg_output(track_ops)
      

    # 7) get optimal assignments for all pairs of recordings (first to last)
    all_ds_assign, all_ds_assign_thr, all_ds_thr_met, all_ds_thr = get_all_ds_assign(track_ops, all_ds_all_roi_ref, all_ds_all_roi_reg)
    plot_thr_met_hist(all_ds_thr_met, all_ds_thr, track_ops)
    plot_n_matched_roi(all_ds_thr_met, all_ds_thr, track_ops)


    # 8) get match matrices for all pairs of recordings (first to last)
    all_pl_match_mat = get_all_pl_match_mat(all_ds_all_roi_ref, all_ds_assign_thr, track_ops)


    # 9) save results
    save_track_ops(track_ops)

    
    save_all_pl_match_mat(all_pl_match_mat, track_ops)


    # 10) save in suite2p format
    if track_ops.save_in_s2p_format:
        print('Saving in suite2p format...')
        save_in_s2p_format(track_ops)
        
    # 10) plot results
    print('Generating plots (this can take some time)...')
    all_ds_stat_iscell = load_all_ds_stat_iscell(track_ops)
    all_ds_centroids = load_all_ds_centroids(all_ds_stat_iscell, track_ops)
    all_ds_mean_img = load_all_ds_mean_img(track_ops)
    if track_ops.nchannels==2:
        all_ds_mean_img_ch2 = load_all_ds_mean_img(track_ops, ch=2)


    plot_roi_match_multiplane(all_ds_mean_img, all_ds_centroids, all_pl_match_mat, track_ops, win_size=track_ops.win_size) # TODO: match histogram to the first roi of first batch (not first roi of each batch)
    plot_allroi_match_multiplane(all_ds_mean_img, all_pl_match_mat, track_ops)
    if track_ops.nchannels==2:
        plot_roi_match_multiplane(all_ds_mean_img_ch2, all_ds_centroids, all_pl_match_mat, track_ops, win_size=track_ops.win_size, ch=2) # TODO: match histogram to the first roi of first batch (not first roi of each batch)
        plot_allroi_match_multiplane(all_ds_mean_img_ch2, all_pl_match_mat, track_ops, ch=2)


    
    print('\n\n\nDone!\n\n\n')
    

def save_in_s2p_format(track_ops):
    
    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}')

    folderpath=track_ops.save_path
    track_ops_dict = np.load(os.path.join(folderpath,  "track_ops.npy"), allow_pickle=True).item()
    track_ops = SimpleNamespace(**track_ops_dict)
    iscell_thr = track_ops.iscell_thr


    for j in range(track_ops.nplanes):
        
        all_f_t2p= []
        all_ops = []
        all_stat_t2p = []
        all_iscell_t2p = []
        fneu_iscell_t2p= []
        spks_iscell_t2p= []
        
        fneu_chan2_iscell_t2p = []
        f_chan2_iscell_t2p = []
        redcell_iscell_t2p = []
            
        
        print(f'nplanes{j}')
        t2p_match_mat = np.load(os.path.join(folderpath,f'plane{str(j)}_match_mat.npy'), allow_pickle=True)
        t2p_match_mat_allday = t2p_match_mat[~np.any(t2p_match_mat == None, axis=1), :]
        for (i, ds_path) in enumerate(track_ops.all_ds_path):
                    ops = np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'ops.npy'), allow_pickle=True).item()
                    stat = np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'stat.npy'), allow_pickle=True)
                    f = np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'F.npy'), allow_pickle=True)
                    fneu= np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'Fneu.npy'), allow_pickle=True)
                    spks= np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'spks.npy'), allow_pickle=True)
                    iscell = np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'iscell.npy'), allow_pickle=True)
                    if track_ops.nchannels==2:
                        f_chan2=np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'F_chan2.npy'), allow_pickle=True)
                        fneu_chan2 = np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'Fneu_chan2.npy'), allow_pickle=True)
                        redcell=np.load(os.path.join(ds_path, 'suite2p', f'plane{str(j)}', 'redcell.npy'), allow_pickle=True)

                    if track_ops.iscell_thr==None:
                        stat_iscell = stat[iscell[:, 0] == 1]
                        f_iscell = f[iscell[:, 0] == 1, :]
                        fneu_iscell = fneu[iscell[:, 0] == 1, :]
                        spks_iscell = spks[iscell[:, 0] == 1, :]
                        is_cell = iscell[iscell[:, 0] == 1, :]
                        if track_ops.nchannels==2:
                            f_chan2_iscell = f_chan2[iscell[:, 0] == 1, :]
                            fneu_chan2_iscell = fneu_chan2[iscell[:, 0] == 1, :]
                            redcell_iscell = redcell[iscell[:, 0] == 1]
                    else:
                        stat_iscell = stat[iscell[:, 1] > iscell_thr]
                        f_iscell = f[iscell[:, 1] > iscell_thr, :]
                        fneu_iscell = fneu[iscell[:, 1] > iscell_thr, :]
                        spks_iscell = spks[iscell[:, 1] > iscell_thr, :]
                        is_cell = iscell[iscell[:, 1] > iscell_thr, :]
                        if track_ops.nchannels==2:
                            f_chan2_iscell = f_chan2[iscell[:, 1] > track_ops.iscell_thr, :]
                            fneu_chan2_iscell = fneu_chan2[iscell[:, 1] > track_ops.iscell_thr, :]
                            redcell_iscell = redcell[iscell[:, 1] > track_ops.iscell_thr]
                    
                    stat_t2p = stat_iscell[t2p_match_mat_allday[:, i].astype(int)]
                    f_t2p = f_iscell[t2p_match_mat_allday[:, i].astype(int), :]
                    fneu_t2p = fneu_iscell[t2p_match_mat_allday[:, i].astype(int), :]
                    spks_t2p = spks_iscell[t2p_match_mat_allday[:, i].astype(int), :]
                    iscell_t2p = is_cell[t2p_match_mat_allday[:, i].astype(int), :]
                    if track_ops.nchannels==2:
                        fneu_chan2_t2p = fneu_chan2_iscell[t2p_match_mat_allday[:, i].astype(int), :]
                        f_chan2_t2p = f_chan2_iscell[t2p_match_mat_allday[:, i].astype(int), :]
                        redcell_t2p = redcell_iscell[t2p_match_mat_allday[:, i].astype(int)]

                    all_stat_t2p.append(stat_t2p)
                    all_f_t2p.append(f_t2p)
                    all_ops.append(ops)
                    all_iscell_t2p.append(iscell_t2p)      
                    fneu_iscell_t2p.append(fneu_t2p)
                    spks_iscell_t2p.append(spks_t2p)  
                    if track_ops.nchannels==2:
                        fneu_chan2_iscell_t2p.append(fneu_chan2_t2p)
                        f_chan2_iscell_t2p.append(f_chan2_t2p)
                        redcell_iscell_t2p.append(redcell_t2p)  
        

        output_folderpath=os.path.join(folderpath, 'matched_suite2p')
        last_elements = [os.path.basename(path) for path in track_ops.all_ds_path]
        print(last_elements)
        
        # Save each element of each list to a .npy file
        for i in range(len(track_ops.all_ds_path)):
            stat_t2p, f_t2p, ops, iscell_t2p, fneu_t2p, spks_t2p = all_stat_t2p[i], all_f_t2p[i], all_ops[i], all_iscell_t2p[i], fneu_iscell_t2p[i], spks_iscell_t2p[i]
            subfolder_path = os.path.join(output_folderpath, last_elements[i])
            if not os.path.exists(subfolder_path):
                os.makedirs(subfolder_path)
            inner_folderpath = os.path.join(subfolder_path, 'suite2p')
            if not os.path.exists(inner_folderpath):
                os.makedirs(inner_folderpath)
            plane_folderpath = os.path.join(inner_folderpath, f'plane{j}')
            if not os.path.exists(plane_folderpath):
                os.makedirs(plane_folderpath)
            

            np.save(os.path.join(plane_folderpath,f"stat.npy"), stat_t2p)
            np.save(os.path.join(plane_folderpath,  f"F.npy"), f_t2p)
            np.save(os.path.join(plane_folderpath, f"ops.npy"), ops)
            np.save(os.path.join(plane_folderpath, f"iscell.npy"), iscell_t2p)
            np.save(os.path.join(plane_folderpath, f"Fneu.npy"), fneu_t2p)
            np.save(os.path.join(plane_folderpath,f"spks.npy"), spks_t2p)
            if track_ops.nchannels==2:
                for i, (redcell_t2p, f_chan2_t2p, fneu_chan2_t2p) in enumerate(zip(redcell_iscell_t2p, f_chan2_iscell_t2p, fneu_chan2_iscell_t2p)):
                    np.save(os.path.join(plane_folderpath,  f"F_chan2.npy"), f_chan2_t2p)
                    np.save(os.path.join(plane_folderpath,  f"Fneu_chan2.npy"), fneu_chan2_t2p)
                    np.save(os.path.join(plane_folderpath, f"redcell.npy"), redcell_t2p)