# contains code for generating plots while the pipeline is running

import numpy as np
import matplotlib.pyplot as plt

from track2p.plot.utils import match_hist_all

def plot_all_planes(all_ds_avg_ch, track_ops, sat_perc=99, ch='funcional'):
    nplanes = track_ops.nplanes
    fig, axs = plt.subplots(nplanes, len(track_ops.all_ds_path), figsize=(3 * len(track_ops.all_ds_path), 3 * nplanes), dpi=300)
    # add dummy dimension to axs if only one plane
    if nplanes==1:
        axs = np.expand_dims(axs, axis=0)
    all_ds_avg_ch_matched = match_hist_all(all_ds_avg_ch)


    for i in range(nplanes):
        for j in range(len(track_ops.all_ds_path)):
            img = all_ds_avg_ch_matched[j][i]
            axs[i, j].imshow(img, cmap='gray', vmin=0, vmax=np.percentile(img, sat_perc))
            axs[i, j].set_title('Plane ' + str(i) + ' in dataset ' + str(j))
            axs[i, j].axis('off')

    fig.savefig(track_ops.save_path_fig + f'mean_fov_{ch}.png', bbox_inches='tight', dpi=200)     

    plt.close(fig)

    

