# contains code for generating plots based on track_ops object after the pipeline is run (this will be saved as an npy file)
# here all functions just take the track_ops object as input
import os

import numpy as np
import matplotlib.pyplot as plt

from skimage.exposure import match_histograms
from track2p.plot.utils import make_rgb_img, saturate_perc, get_all_wind_mean_img
from track2p.io.loaders import load_stat_ds_plane, get_all_roi_array_from_stat

def plot_reg_img_output(track_ops):
    # make a plot where on the top its all the images and the bottom is the overlays before and after registration
    nplanes = track_ops.nplanes
    n_row = nplanes + 2*nplanes # number of plays + 2 overlays per plane
    n_col = len(track_ops.all_ds_path) # number of datasets
    figsize = (10/3 * n_col, 10/3 * n_row)
    fig, axs = plt.subplots(n_row, n_col, figsize=figsize, dpi=300)

    # first populate first (n_planes) rows with images
    for i in range(nplanes):
        axs[i, 0].set_ylabel(f'plane{i}\nchan{track_ops.reg_chan}', rotation=0, size='large', ha='right', va='center', labelpad=20)
        for j in range(len(track_ops.all_ds_path)):
            img = track_ops.all_ds_avg_ch1[j][i] if track_ops.reg_chan==0 else track_ops.all_ds_avg_ch2[j][i]
            axs[i, j].imshow(img, cmap='gray', vmin=0, vmax=np.percentile(img, 99))


    for i in range(len(track_ops.all_ds_path)-1): # last one won't have overlay
        for j in range(nplanes):
            # get subplot indices
            row_nonreg = nplanes + 2*j # first shift for number of initial rows, then shift by 2 for each plane (before and after registration)
            row_reg = nplanes + 2*j + 1 # the row after the nonreg

            # get images
            ref_img = track_ops.all_ds_ref_img[i][j] # get ref image for this pair
            mov_img = track_ops.all_ds_mov_img[i][j] # get mov image for this pair
            mov_img_ref = track_ops.all_ds_mov_img_reg[i][j] # get mov image after registration for this pair

            # match histograms to reference
            mov_img = match_histograms(mov_img, ref_img)
            mov_img_ref = match_histograms(mov_img_ref, ref_img)

            # assemble and saturate the overlays
            img_rgb = make_rgb_img(ref_img, mov_img)
            img_rgb_reg = make_rgb_img(ref_img, mov_img_ref)
            img_rgb = saturate_perc(img_rgb, sat_perc=track_ops.sat_perc)
            img_rgb_reg = saturate_perc(img_rgb_reg, sat_perc=track_ops.sat_perc)

            # plot the overlays
            axs[row_nonreg, i].imshow(img_rgb)
            axs[row_reg, i].imshow(img_rgb_reg)

        if i == 0:
            for k in range(nplanes):
                axs[nplanes + 2*k,i].set_ylabel(f'plane{k}\nchan{track_ops.reg_chan}\nnon-reg', rotation=0, size='large', ha='right', va='center', labelpad=20)
                axs[nplanes + 2*k+1,i].set_ylabel(f'plane{k}\nchan{track_ops.reg_chan}\nreg', rotation=0, size='large', ha='right', va='center', labelpad=20)

            # loop through all subplots and remove ticks and labels and spines

    for ax in axs.flat:
        ax.set(xticks=[], yticks=[])
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['bottom'].set_visible(False)

    # add arrows and dashed lines etc.

    axs[2*nplanes, 0].annotate('', xy=(0, 1.1), xytext=(n_col + 0.5, 1.1), 
                xycoords='axes fraction', textcoords='axes fraction',
                arrowprops=dict(arrowstyle='-', linestyle='dashed', color='grey'), 
                annotation_clip=False)

    for i in range(len(track_ops.all_ds_path)-1):
        axs[nplanes-1,i].annotate('', xy=(0.5, -0.17), xytext=(0.5, -0.02), 
                        xycoords='axes fraction', textcoords='axes fraction',
                        arrowprops=dict(facecolor=(1,0,0), edgecolor=(1,0,0), shrink=0.05), 
                        annotation_clip=False)
        axs[nplanes-1,i+1].annotate('', xy=(-0.6, -0.17), xytext=(0.5, -0.02), 
                        xycoords='axes fraction', textcoords='axes fraction',
                        arrowprops=dict(facecolor=(0,1,0), edgecolor=(0,1,0), shrink=0.05), 
                        annotation_clip=False)
    

    # save figure into the output path
    fig.savefig(os.path.join(track_ops.save_path_fig, 'reg_img_output.png'), bbox_inches='tight', dpi=200)
    plt.close(fig)

def plot_roi_reg_output(track_ops):
    # make a plot where on the top its all the images and the bottom is the overlays before and after registration
    nplanes = track_ops.nplanes
    n_row = nplanes + 2*nplanes # number of plays + 2 overlays per plane
    n_col = len(track_ops.all_ds_path) # number of datasets
    figsize = (10/3 * n_col, 10/3 * n_row)
    # figsize = (10 * n_col, 10 * n_row)

    fig, axs = plt.subplots(n_row, n_col, figsize=figsize, dpi=300)

    # first populate first (n_planes) rows with images
    for i in range(nplanes):
        axs[i, 0].set_ylabel(f'plane{i}\nchan{track_ops.reg_chan}', rotation=0, size='large', ha='right', va='center', labelpad=20)
        for j in range(len(track_ops.all_ds_path)):
            img = track_ops.all_ds_avg_ch1[j][i]
            axs[i, j].imshow(img, cmap='gray', vmin=0, vmax=np.percentile(img, 99), alpha=0.5)
            # take the rois of reference unless it is last on, then take the rois of the mov
            all_roi = track_ops.all_ds_all_roi_array_ref[j][i] if j < len(track_ops.all_ds_path)-1 else track_ops.all_ds_all_roi_array_mov[j-1][i]
            for k in range(all_roi.shape[2]):
                axs[i, j].contour(all_roi[:,:,k], colors='C0', linewidths=0.3)

    # now populate the next (n_planes) rows with the overlays before registration
    for i in range(len(track_ops.all_ds_path)-1):
        print(f'Plotting contours for dataset {i}/{len(track_ops.all_ds_path)-1}')
        for j in range(track_ops.nplanes):

            row_nonreg = nplanes + 2*j # first shift for number of initial rows, then shift by 2 for each plane (before and after registration)
            row_reg = nplanes + 2*j + 1 # the row after the nonreg

            all_roi_array_ref = track_ops.all_ds_all_roi_array_ref[i][j]

            all_roi_array_mov = track_ops.all_ds_all_roi_array_mov[i][j]
            ref_mov_inters = track_ops.all_ds_ref_mov_inters[i][j]

            axs[row_nonreg,i].imshow(ref_mov_inters)
            for k in range(all_roi_array_ref.shape[2]):
                axs[row_nonreg,i].contour(all_roi_array_ref[:,:,k], colors='r', linewidths=0.3)
            for k in range(all_roi_array_mov.shape[2]):
                axs[row_nonreg,i].contour(all_roi_array_mov[:,:,k], colors='g', linewidths=0.3)

            # TODO: for reg row
            all_roi_array_reg = track_ops.all_ds_all_roi_array_reg[i][j]
            ref_reg_inters = track_ops.all_ds_ref_reg_inters[i][j]

            axs[row_reg,i].imshow(ref_reg_inters)
            for k in range(all_roi_array_ref.shape[2]):
                axs[row_reg,i].contour(all_roi_array_ref[:,:,k], colors='r', linewidths=0.3)
            for k in range(all_roi_array_reg.shape[2]):
                axs[row_reg,i].contour(all_roi_array_reg[:,:,k], colors='g', linewidths=0.3)

        if i == 0:
            for k in range(nplanes):
                axs[nplanes + 2*k,i].set_ylabel(f'plane{k}\nnon-reg', rotation=0, size='large', ha='right', va='center', labelpad=20)
                axs[nplanes + 2*k+1,i].set_ylabel(f'plane{k}\nreg', rotation=0, size='large', ha='right', va='center', labelpad=20)


    for ax in axs.flat:
        ax.set(xticks=[], yticks=[])
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['bottom'].set_visible(False)     

    # add arrows and dashed lines etc.

    axs[2*nplanes, 0].annotate('', xy=(0, 1.1), xytext=(n_col + 0.5, 1.1), 
                xycoords='axes fraction', textcoords='axes fraction',
                arrowprops=dict(arrowstyle='-', linestyle='dashed', color='grey'), 
                annotation_clip=False)

    for i in range(len(track_ops.all_ds_path)-1):
        axs[nplanes-1,i].annotate('', xy=(0.5, -0.17), xytext=(0.5, -0.02), 
                        xycoords='axes fraction', textcoords='axes fraction',
                        arrowprops=dict(facecolor='r', edgecolor='r', shrink=0.05), 
                        annotation_clip=False)
        axs[nplanes-1,i+1].annotate('', xy=(-0.6, -0.17), xytext=(0.5, -0.02), 
                        xycoords='axes fraction', textcoords='axes fraction',
                        arrowprops=dict(facecolor='g', edgecolor='g', shrink=0.05), 
                        annotation_clip=False)
    

    # save figure into the output path
    fig.savefig(track_ops.save_path_fig + 'reg_roi_output.png', bbox_inches='tight', dpi=200)     
    plt.close(fig)


def plot_thr_met_hist(all_ds_thr_met, all_ds_thr, track_ops):
    fig, axs = plt.subplots(track_ops.nplanes, len(all_ds_thr_met), figsize=(6*len(all_ds_thr_met), 6*track_ops.nplanes), sharey=True, sharex=True)
    for i in range(len(all_ds_thr_met)):
        for j in range(track_ops.nplanes):
            axs = np.array([axs]) if type(axs) is not np.ndarray else axs
            this_ax = axs[i] if track_ops.nplanes==1 else axs[j][i]
            n_reg_roi = len(all_ds_thr_met[i][j])
            n_abovethr_roi = np.sum(all_ds_thr_met[i][j]>all_ds_thr[i][j])
            this_ax.hist(all_ds_thr_met[i][j], bins=20)
            this_ax.axvline(all_ds_thr[i][j], color='grey', linestyle='--')
            # label the line with 'otsu threshold'
            this_ax.text(all_ds_thr[i][j]+0.02, this_ax.get_ylim()[1]*0.9, f'{track_ops.thr_method} thr.: {all_ds_thr[i][j]:.2f}')
            this_ax.set_title(f'ds{i} (ref) to ds{i+1} (reg); matched {n_abovethr_roi}/{n_reg_roi} ({n_abovethr_roi/n_reg_roi*100:.1f}%)')
            
            if i==0:
                this_ax.set_ylabel(f'ROI count (plane{j})')
    
    # remove the top and right spines
    for a in axs.flatten():
        a.spines['top'].set_visible(False)
        a.spines['right'].set_visible(False)

    
    plt.tight_layout()
    plt.savefig(os.path.join(track_ops.save_path_fig, 'thr_met_hist.png'), dpi=200)
    plt.close(fig)



def plot_n_matched_roi(all_ds_thr_met, all_ds_thr, track_ops):
    
    fig, axs = plt.subplots(track_ops.nplanes, 2, figsize=(8, 4*track_ops.nplanes), sharey='col', sharex=True)
                            
    for i in range(track_ops.nplanes):
        all_n_reg_roi = [len(all_ds_thr_met[j][i]) for j in range(len(all_ds_thr_met))]
        all_n_abovethr_roi = [np.sum(all_ds_thr_met[j][i]>all_ds_thr[j][i]) for j in range(len(all_ds_thr_met))]
        all_prop_abovethr = [n_abovethr_roi/n_reg_roi for n_abovethr_roi, n_reg_roi in zip(all_n_abovethr_roi, all_n_reg_roi)]

        ax0 = axs[i, 0] if track_ops.nplanes>1 else axs[0]
        
        ax0.bar(np.arange(len(all_n_reg_roi)), all_n_reg_roi, color='grey', label='total')
        ax0.bar(np.arange(len(all_n_abovethr_roi)), all_n_abovethr_roi, color='C0', label='above threshold')

        ax0.set_ylabel(f'ROI count (plane {i})')
        ax0.set_xticks(np.arange(len(all_n_reg_roi)))
        ax0.set_xticklabels([f'r{i}-{i+1}' for i in range(len(all_n_reg_roi))])
        ax0.legend(loc='lower right')
        ax0.spines['top'].set_visible(False)
        ax0.spines['right'].set_visible(False)

        ax1 = axs[i, 1] if track_ops.nplanes>1 else axs[1]

        ax1.plot(np.arange(len(all_n_reg_roi)), all_prop_abovethr, color='C0', marker='o')
        ax1.set_ylabel(f'ROI prop. (plane {i})')
        ax1.set_xticks(np.arange(len(all_n_reg_roi)))
        ax1.set_xticklabels([f'r{i}-{i+1}' for i in range(len(all_n_reg_roi))])
        ax1.set_ylim(0, 1)
        ax1.spines['top'].set_visible(False)
        ax1.spines['right'].set_visible(False)

        fig.savefig(os.path.join(track_ops.save_path_fig, 'n_prop_matched_roi.png'), dpi=200, bbox_inches='tight')
        plt.close(fig)

def plot_roi_match(all_ds_mean_img, all_ds_centroids, all_pl_match_mat, neuron_ids, track_ops, plane_idx=0, win_size=64, k=0, n=None, ch=1):

    if n is None:
        n = len(neuron_ids)

    nrows = len(neuron_ids)
    ncols = len(all_ds_mean_img)
    fig, axs = plt.subplots(nrows, ncols, figsize=(2*ncols, 2*nrows), dpi=50)

    for (i, nrn_id) in enumerate(neuron_ids):
        # check if neuron is not matched (if ith row of all_pl_match_mat is all None)
        if any(all_pl_match_mat[plane_idx][nrn_id,:]==None):
            continue
        
        # get all wind_mean_img (small window around centroid)
        all_wind_mean_img = get_all_wind_mean_img(all_ds_mean_img, all_ds_centroids, all_pl_match_mat, nrn_id, plane_idx=plane_idx, win_size=win_size)

        if i == 0:
            fig_ref = all_wind_mean_img[0] # reference image for matching histograms along whole image
        else:
            all_wind_mean_img[0] = match_histograms(all_wind_mean_img[0], fig_ref) # if not first ROI of whole image then match to first ROI of whole image

        for j in range(len(all_ds_mean_img)):
            wind_mean_img = all_wind_mean_img[j]
            ref_img = all_wind_mean_img[0]
            matched = match_histograms(wind_mean_img, ref_img)
            axs[i, j].imshow(matched, cmap='gray')
            # scatter middle pixel of the window
            axs[i, j].scatter(win_size/2, win_size/2, color='C0')
            if i==0:
                axs[i, j].set_title(f'ds {j} (pl {plane_idx})')
            if j==0:
                axs[i, j].set_ylabel(f'ROI {nrn_id} ({i+k}/{n})')

    # remove axes labels
    for ax in axs.flat:
        ax.set_xticks([])
        ax.set_yticks([])

    plt.tight_layout()
    if ch == 1:
        plt.savefig(os.path.join(track_ops.save_path_fig, f'roi_match_plane{plane_idx}_idx{k}-{k+len(neuron_ids)}.png'), dpi=50)
    elif ch == 2:
        plt.savefig(os.path.join(track_ops.save_path_fig, f'roi_match_chan2_plane{plane_idx}_idx{k}-{k+len(neuron_ids)}.png'), dpi=50)
    plt.close(fig)

def plot_roi_match_multiplane(all_ds_mean_img, all_ds_centroids, all_pl_match_mat, track_ops, win_size=48, ch=1):

    for i in range(track_ops.nplanes):
        pl_neuron_ids = np.arange(all_pl_match_mat[i].shape[0])
        for j in range(len(track_ops.save_path)):
            neuron_ids = pl_neuron_ids[~np.any(all_pl_match_mat[i]==None, axis=1)]

        if len(neuron_ids) < 100:
            plot_roi_match(all_ds_mean_img, all_ds_centroids, all_pl_match_mat, neuron_ids, track_ops, plane_idx=i, win_size=win_size, ch=ch)
        else:
            # plot in batches of 100
            for k in range(0, len(neuron_ids), 100):
                plot_roi_match(all_ds_mean_img, all_ds_centroids, all_pl_match_mat, neuron_ids[k:k+100], track_ops, plane_idx=i, win_size=win_size, k=k, n=len(neuron_ids), ch=ch)


def plot_allroi_match_multiplane(all_ds_mean_img, all_pl_match_mat, track_ops, ch=1):
    

    fig, axs = plt.subplots(track_ops.nplanes, len(track_ops.all_ds_path), figsize=(4*len(track_ops.all_ds_path), 4*track_ops.nplanes), dpi=300)

    for ds_idx, ds_path in enumerate(track_ops.all_ds_path):
        for plane_idx in range(track_ops.nplanes):
            # saturating the reference image
            ref_mean_img = all_ds_mean_img[0][plane_idx]
            perc = np.percentile(ref_mean_img, track_ops.sat_perc)
            ref_mean_img[ref_mean_img>perc] = perc

            ax = axs[ds_idx] if track_ops.nplanes==1 else axs[plane_idx, ds_idx]
            mean_img = all_ds_mean_img[ds_idx][plane_idx]
            mean_img_matched = match_histograms(mean_img, ref_mean_img)
            ax.imshow(mean_img_matched, cmap='gray')
            ax.set_title(f'ds{ds_idx} (plane{plane_idx})')

    # add contours
    for plane_idx in range(track_ops.nplanes):
        pl_neuron_ids = np.arange(all_pl_match_mat[plane_idx].shape[0])
        neuron_ids = pl_neuron_ids[~np.any(all_pl_match_mat[plane_idx]==None, axis=1)]
        neuron_colors = [np.random.rand(3) for i in range(len(neuron_ids))]

        for ds_idx, _ in enumerate(track_ops.all_ds_path):
            
            stat_ds_plane, _ = load_stat_ds_plane(track_ops.all_ds_path[ds_idx], track_ops, plane_idx=plane_idx)
            all_roi_array = get_all_roi_array_from_stat(stat_ds_plane, track_ops)
            
            ax = axs[ds_idx] if track_ops.nplanes==1 else axs[plane_idx, ds_idx]
        
            for (i, neuron_idx) in enumerate(neuron_ids):
                idx_orig = all_pl_match_mat[plane_idx][neuron_idx, ds_idx] # neurons index in the original recording
                cont = stat_ds_plane

                cont_plot = ax.contour(all_roi_array[:,:,idx_orig], linewidths=0.5)

                for collection in cont_plot.collections:
                    collection.set_edgecolor(neuron_colors[i])  # RGB value for red
        
        left_axs = axs[0] if track_ops.nplanes==1 else axs[plane_idx,0]
        left_axs.set_ylabel(f'plane{plane_idx} (n={len(neuron_ids)})')

    # remove axes elements
    for ax in axs.flat:
        ax.set_xticks([])
        ax.set_yticks([])
        ax.set_xticklabels([])
        ax.set_yticklabels([])

    plt.tight_layout()
    if ch == 1:
        plt.savefig(os.path.join(track_ops.save_path_fig, f'all_roi_match.png'), dpi=300)
    else:
        plt.savefig(os.path.join(track_ops.save_path_fig, f'all_roi_match_chan2.png'), dpi=300)
    plt.close(fig)
