import numpy as np
import scipy.optimize
from ipywidgets import widgets
from ipywidgets import Layout, interact
from matplotlib import pyplot as plt, cm as cm
from matplotlib.patches import Circle, Arrow, Rectangle
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.ndimage import median_filter
from skimage.filters import median
from skimage.morphology import disk

from piscat.Localization import gaussian_2D_fit
from piscat.Localization import directional_intensity
from piscat.Trajectory.data_handeling import protein_trajectories_list2dic
from piscat.InputOutput import read_status_line


class JupyterDisplay:

    def __init__(self, video, median_filter_flag=False, color='gray', title=None, xlabel=None, ylabel=None,
                 imgSizex=5, imgSizey=5, extent=None, IntSlider_width='500px', step=1):
        """
        This class displays the video in jupyter notebook.

        Parameters
        ----------
        video:(NDArray)
            Input video.

        median_filter_flag: bool
            In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
            It defines the colormap for visualization.

        title: list
            A list of string titles with a length equal to the number of frames in the video.

        xlabel: str
            The label text for x-axis.

        ylabel: str
            The label text for y-axis.

        imgSizex: int
            Image length size.

        imgSizey: int
            Image width size.

        extent:
            The extent kwarg determines the bounding box in data coordinates that the image will fill, which is
            specified in data coordinates as (left, right, bottom, top).

        IntSlider_width: str
            Size of slider

        step: int
            Stride between visualization frames.

        """
        self.color = color
        self.video = video
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.title = title
        self.extent = extent

        self.median_filter_flag = median_filter_flag

        self.xlabel = xlabel
        self.ylabel = ylabel

        interact(self.display, frame=widgets.IntSlider(min=0, max=self.video.shape[0] - 1, step=step, value=10,
                                                       layout=Layout(width=IntSlider_width),
                                                       readout_format='10', continuous_update=False,
                                                       description='Frame:'))

    def display(self, frame):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        ax = fig.add_axes([0, 0, 1, 1])

        if self.median_filter_flag:
            frame_v = median_filter(self.video[int(frame), :, :], 3)
        else:
            frame_v = self.video[int(frame), :, :]

        if self.extent is not None:
            myplot = ax.imshow(frame_v, cmap=self.color, extent=self.extent)
        else:
            myplot = ax.imshow(frame_v, cmap=self.color)

        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        plt.colorbar(myplot, cax=cax)

        if self.xlabel is not None:
            ax.set_xlabel(self.xlabel)

        if self.ylabel is not None:
            ax.set_ylabel(self.ylabel)

        if self.title is not None:
            ax.set_title(self.title[int(frame)])

        plt.show()


class JupyterDisplay_StatusLine:

    def __init__(self, video, median_filter_flag=False, color='gray', imgSizex=5, imgSizey=5, IntSlider_width='500px', step=1, value=0):
        """
        This class displays the video in the Jupyter notebook interactively while highlight status line

        Parameters
        ----------
        video: NDArray
            Input video.

        median_filter_flag: bool
            In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
            It defines the colormap for visualization.

        imgSizex: int
            Image length size.

        imgSizey: int
            Image width size.

        IntSlider_width: str
            Size of slider.

        step: int
            Stride between visualization frames.

        value: int
            Initial frame value for visualization
        """
        self.color = color
        self.video = video
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.median_filter_flag = median_filter_flag

        status_ = read_status_line.StatusLine(video)
        self.video_remove, status_information = status_.find_status_line()
        self.statusLine_position = status_information['status_line_position']

        interact(self.display, frame=widgets.IntSlider(min=0, max=self.video.shape[0] - 1, step=step, value=value,
                                                       layout=Layout(width=IntSlider_width),
                                                       readout_format='10', continuous_update=False,
                                                       description='Frame:'))

    def display(self, frame):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        ax = fig.add_axes([0, 0, 1, 1])

        if self.median_filter_flag:
            frame_v = median_filter(self.video[int(frame), :, :], 3)
            frame_v_rm = median_filter(self.video_remove[int(frame), :, :], 3)

        else:
            frame_v = self.video[int(frame), :, :]
            frame_v_rm = self.video_remove[int(frame), :, :]


        myplot_rm = ax.imshow(frame_v_rm, cmap=self.color)
        myplot = ax.imshow(frame_v, cmap=self.color)
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        plt.colorbar(myplot_rm, cax=cax)

        if self.statusLine_position == 'column':
            rect = Rectangle((0, self.video.shape[2]), self.video.shape[2], 1, linewidth=10, edgecolor='r',
                                  facecolor='none')
            ax.add_patch(rect)
        elif self.statusLine_position == 'row':
            rect = Rectangle((self.video.shape[1], 0), 1, self.video.shape[1], linewidth=10, edgecolor='r',
                                  facecolor='none')
            ax.add_patch(rect)
        plt.show()


class JupyterPSFs_localizationDisplay:

    def __init__(self, video, df_PSFs, median_filter_flag=False, color='gray', imgSizex=5, imgSizey=5, IntSlider_width='500px', step=1,  value=0):
        """
        This class displays the video in the Jupyter notebook interactively while highlight PSFs.

        Parameters
        ----------
        video: NDArray
           Input video.

        df_PSFs: panda data_frame
            Data Frames that contains the location of PSFs.

        median_filter_flag: bool
           In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
           It defines the colormap for visualization.

        imgSizex: int
           Image length size.

        imgSizey: int
           Image width size.

        IntSlider_width: str
           Size of slider.

        step: int
           Stride between visualization frames.

        value: int
           Initial frame value for visualization
       """
        self.video = video
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.df_PSFs = df_PSFs
        self.median_filter_flag = median_filter_flag
        self.color = color

        interact(self.show_psf, frame_number=widgets.IntSlider(min=0, max=self.video.shape[0] - 1, step=step, value=value,
                                                               readout_format='1', continuous_update=False, layout=Layout(width=IntSlider_width),
                                                               description='Frame:'))

    def show_psf(self, frame_number):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        ax = fig.add_axes([0, 0, 1, 1])
        if self.median_filter_flag:
            frame_v = median_filter(self.video[int(frame_number), :, :], 3)
        else:
            frame_v = self.video[int(frame_number), :, :]

        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        myplot = ax.imshow(frame_v, cmap=self.color)
        plt.colorbar(myplot, cax=cax)

        particle = self.df_PSFs.loc[self.df_PSFs['frame'] == frame_number]
        particle_X = particle['x'].tolist()
        particle_Y = particle['y'].tolist()
        particle_sigma = particle['sigma'].tolist()

        for j_ in range(len(particle_X)):
            y = int(particle_Y[j_])
            x = int(particle_X[j_])
            sigma = particle_sigma[j_]
            ax.add_patch(Circle((x, y), radius=np.sqrt(2) * sigma, edgecolor='r', facecolor='none', linewidth=2))
        plt.show()


class JupyterPSFs_localizationPreviewDisplay:

    def __init__(self, video, df_PSFs, title='', frame_num=None, median_filter_flag=False, color='gray', imgSizex=5, imgSizey=5, IntSlider_width='500px', step=1):
        """
        This class displays the video in the Jupyter notebook interactively while highlight PSFs.

        Parameters
        ----------
        video: NDArray
           Input video.

        df_PSFs: panda data_frame
            Data Frames that contains the location of PSFs.

        title: str
            It defines title of plot.

        frame_num: list
            list of frame that we want to see preview of localization.

        median_filter_flag: bool
           In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
           It defines the colormap for visualization.

        imgSizex: int
           Image length size.

        imgSizey: int
           Image width size.

        IntSlider_width: str
           Size of slider.

        step: int
           Stride between visualization frames.
       """
        self.video = video
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.df_PSFs = df_PSFs
        self.median_filter_flag = median_filter_flag
        self.color = color
        if self.df_PSFs is not None:
            self.frame_num = np.sort(np.unique(df_PSFs['frame'].tolist()))
        else:
            self.frame_num = frame_num
        self.currentFrame = 0
        self.IntSlider_width = IntSlider_width
        self.step = step
        self.title = title

    def display_run(self):
        interact(self.show_psf,
                 index=widgets.IntSlider(min=0, max=len(self.frame_num) - 1, step=self.step, value=self.currentFrame,
                                         readout_format='1', continuous_update=False,
                                         layout=Layout(width=self.IntSlider_width),
                                         description='index:'))

    def show_psf(self, index):
        frame_number = self.frame_num[index]
        self.currentFrame = frame_number
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        ax = fig.add_axes([0, 0, 1, 1])
        if self.median_filter_flag:
            frame_v = median_filter(self.video[int(frame_number), :, :], 3)
        else:
            frame_v = self.video[int(frame_number), :, :]

        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        myplot = ax.imshow(frame_v, cmap=self.color)
        ax.set_title(self.title)
        plt.colorbar(myplot, cax=cax)

        if self.df_PSFs is not None:
            particle = self.df_PSFs.loc[self.df_PSFs['frame'] == frame_number]
            particle_X = particle['x'].tolist()
            particle_Y = particle['y'].tolist()
            particle_sigma = particle['sigma'].tolist()

            for j_ in range(len(particle_X)):
                y = int(particle_Y[j_])
                x = int(particle_X[j_])
                sigma = particle_sigma[j_]
                ax.add_patch(Circle((x, y), radius=np.sqrt(2) * sigma, edgecolor='r', facecolor='none', linewidth=2))
        plt.show()

    def set_df(self, new_df):
        self.df_PSFs = new_df


class JupyterPSFs_subplotLocalizationDisplay:

    def __init__(self, list_videos, list_df_PSFs, numRows, numColumns, list_titles=None,
                 median_filter_flag=False, color='gray', imgSizex=5, imgSizey=5, IntSlider_width='500px', step=1, value=0):
        """
        This class shows several videos (with the same number of frames) at once in the Jupyter notebook interactively while highlight localize PSFs.

        Parameters
        ----------
        list_videos: list of NDArray
          List of videos.

        list_df_PSFs: list panda data_frame
           List Data Frames that contains the location of PSFs for each video.

        numRows: int
            It defines number of rows in sub-display.

        numColumns: int
            It defines number of columns in sub-display.

        list_titles: list str
            List of titles for each sub plot.

        median_filter_flag: bool
          In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
          It defines the colormap for visualization.

        imgSizex: int
          Image length size.

        imgSizey: int
          Image width size.

        IntSlider_width: str
          Size of slider

        step: int
          Stride between visualization frames.

        value: int
            Initial frame value for visualization
        """
        self.list_video = list_videos
        self.numVideos = len(list_videos)
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.list_df_PSFs = list_df_PSFs
        self.median_filter_flag = median_filter_flag
        self.color = color
        self.numRows = numRows
        self.numColumns = numColumns

        if list_titles is None:
            self.list_titles = [None for _ in range(self.numVideos)]
        else:
            self.list_titles = list_titles

        max_numberFrames = np.max([vid_.shape[0] for vid_ in list_videos])
        interact(self.show_psf, frame_number=widgets.IntSlider(min=0, max=max_numberFrames, step=step, value=value,
                                                               readout_format='1', continuous_update=False, layout=Layout(width=IntSlider_width),
                                                               description='Frame:'))

    def show_psf(self, frame_number):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        grid = plt.GridSpec(self.numRows, self.numColumns, hspace=0.3, wspace=0.3)

        imgGrid_list = []
        for i in range(self.numRows):
            for j in range(self.numColumns):
                imgGrid_list.append(fig.add_subplot(grid[i, j]))

        for img_, tit_, img_grid_, df_PSFs in zip(self.list_video, self.list_titles, imgGrid_list, self.list_df_PSFs):

            if self.median_filter_flag:

                frame_v = median_filter(img_[int(frame_number), :, :], 3)
            else:
                frame_v = img_[int(frame_number), :, :]

            divider = make_axes_locatable(img_grid_)
            cax = divider.append_axes("right", size="5%", pad=0.05)
            myplot = img_grid_.imshow(frame_v, cmap=self.color)
            img_grid_.axis('off')
            plt.colorbar(myplot, cax=cax)
            if tit_ is not None:
                img_grid_.set_title(tit_)

            particle = df_PSFs.loc[df_PSFs['frame'] == frame_number]
            particle_X = particle['x'].tolist()
            particle_Y = particle['y'].tolist()
            particle_sigma = particle['sigma'].tolist()

            for j_ in range(len(particle_X)):
                y = int(particle_Y[j_])
                x = int(particle_X[j_])
                sigma = particle_sigma[j_]
                img_grid_.add_patch(Circle((x, y), radius=np.sqrt(2) * sigma, edgecolor='r', facecolor='none', linewidth=2))
        plt.show()


class JupyterPSFs_TrackingDisplay:

    def __init__(self, video, df_PSFs, median_filter_flag=False, step=1, color='gray', imgSizex=5, imgSizey=5, value=0):
        """
        This class displays video in the Jupyter notebook interactively while highlighting PSFs with trajectories.

        Parameters
        ----------
        video: NDArray
           Input video.

        df_PSFs: panda data_frame
            Data Frames that contains the location of PSFs.

        median_filter_flag: bool
           In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
           It defines the colormap for visualization.

        imgSizex: int
           Image length size.

        imgSizey: int
           Image width size.

        IntSlider_width: str
           Size of slider.

        step: int
           Stride between visualization frames.

        value: int
            Initial frame value for visualization
        """

        self.video = video
        self.df_PSFs = df_PSFs
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.df_PSFs = df_PSFs
        self.median_filter_flag = median_filter_flag
        self.color = color
        self.median_filter_flag = median_filter_flag
        if self.df_PSFs.particle.isnull().any():
            self.df_PSFs['particle'] = 0
            self.list_particles_idx = self.df_PSFs.particle.unique()
        else:
            self.list_particles_idx = self.df_PSFs.particle.unique()

        colors_ = cm.autumn(np.linspace(0, 1, len(self.list_particles_idx)))

        self.colors = colors_[0:len(self.list_particles_idx), :]
        interact(self.show_psf, frame_number=widgets.IntSlider(min=0, max=self.video.shape[0] - 1, step=step, value=value,
                    readout_format='1', continuous_update=False, description='Frame:'))

    def show_psf(self, frame_number):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        ax = fig.add_axes([0, 0, 1, 1])

        if self.median_filter_flag:
            frame_v = median_filter(self.video[int(frame_number), :, :], 3)
        else:
            frame_v = self.video[int(frame_number), :, :]

        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        myplot = ax.imshow(frame_v, cmap=self.color)
        plt.colorbar(myplot, cax=cax)

        particle = self.df_PSFs.loc[self.df_PSFs['frame'] == frame_number]
        particle_X = particle['x'].tolist()
        particle_Y = particle['y'].tolist()
        particle_sigma = particle['sigma'].tolist()
        particle_labels = particle['particle'].tolist()

        if len(particle_labels) != 0:

            for j_, p_l in zip(range(len(particle_X)), particle_labels):
                y = int(particle_Y[j_])
                x = int(particle_X[j_])
                sigma = particle_sigma[j_]
                ax = plt.gca()
                ax.add_patch(Circle((x, y), radius=np.sqrt(2) * sigma, edgecolor=self.colors[int(p_l)], facecolor='none', linewidth=2))

                all_particle_ = self.df_PSFs.loc[self.df_PSFs['particle'] == p_l]
                particle_f_ = all_particle_['frame'].tolist()
                particle_X_= all_particle_['x'].tolist()
                particle_Y_ = all_particle_['y'].tolist()
                particle_sigma_ = all_particle_['sigma'].tolist()
                for f_, x_, y_, sigma_ in zip(particle_f_, particle_X_, particle_Y_, particle_sigma_):
                    if f_ <= frame_number:
                        ax.add_patch(Circle((int(x_), int(y_)), radius=0.1, edgecolor=self.colors[int(p_l)], facecolor='none', linewidth=1))
                        # ax.add_collection((x_, y_), autolim=True, cmap=self.colors [p_l])
        plt.show()


class JupyterSelectedPSFs_localizationDisplay:

    def __init__(self, video, particles, particles_num='#0', frame_extend=0, median_filter_flag=False, flag_fit2D=False,
                 color='gray', imgSizex=10, imgSizey=10, IntSlider_width='500px', step=1, value=0):
        """
        This class interactively shows video in a Jupyter notebook while highlighting PSFs based on ID.

        Parameters
        ----------
        video: NDArray
           Input video.

        particles: dic
             Dictionary similar to the following structures.:

            | {"#0": {'intensity_horizontal': ..., 'intensity_vertical': ..., ..., 'particle_ID': ...},
                "#1": {}, ...}

        particles_num: str
            Choose the corresponding key in the particles dictionary.

        frame_extend: int
            Display particle for ``frame_extend`` before and after segmented ones. In case there are not enough frames before/after
            , it shows only for the number of existing frames.

        flag_fit2D: bool
            It activate 2D-Gaussian fit to extract fitting information of selected PSF.

        median_filter_flag: bool
           In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
           It defines the colormap for visualization.

        imgSizex: int
           Image length size.

        imgSizey: int
           Image width size.

        IntSlider_width: str
           Size of slider.

        step: int
           Stride between visualization frames.

        value: int
            Initial frame value for visualization
        """

        if type(particles) is list:
            particles = protein_trajectories_list2dic(particles)

        if type(particles) is dict:
            for key, item in particles.items():
                if key == particles_num:
                    intensity_horizontal = item['intensity_horizontal']
                    intensity_vertical = item['intensity_vertical']
                    center_int = item['center_int']
                    center_int_flow = item['center_int_flow']
                    self.frame_number_ = item['frame_number']
                    self.sigma = item['sigma']
                    self.x_center = item['x_center']
                    self.y_center = item['y_center']
                    self.particle_ID = item['particle_ID']

        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.video = video
        self.color = color
        self.flag_fit2D = flag_fit2D
        self.median_filter_flag = median_filter_flag

        video_frameNum = video.shape[0]
        max_particle_frame = np.max(self.frame_number_)
        min_particle_frame = np.min(self.frame_number_)

        max_extend_particle_frame = np.min([max_particle_frame+frame_extend, video_frameNum])
        min_extend_particle_frame = np.max([min_particle_frame-frame_extend, 0])

        self.frame_number = list(range(min_extend_particle_frame, max_extend_particle_frame, 1))
        interact(self.show_psf,
                 f_num=widgets.IntSlider(min=self.frame_number[0], max=self.frame_number[-1] - 1, step=step, value=value,
                                         readout_format='1', continuous_update=False, layout=Layout(width=IntSlider_width),
                                         description='Frame:'))

    def show_psf(self, f_num):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        ax = fig.add_axes([0, 0, 1, 1])
        if self.flag_fit2D:
            fit_params = self.fit_2D_gussian(frame_num=f_num, scale=5)

        if self.median_filter_flag:
            frame_v = median_filter(self.video[int(f_num), :, :], 3)
        else:
            frame_v = self.video[int(f_num), :, :]

        myplot = ax.imshow(frame_v, cmap=self.color)
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        plt.colorbar(myplot, cax=cax)

        if f_num <= np.max(self.frame_number_) and f_num >= np.min(self.frame_number_):
            idx_ = np.where(self.frame_number_ == f_num)

            particle_X = self.x_center[idx_]
            particle_Y = self.y_center[idx_]
            particle_sigma = self.sigma[idx_]
            particle_labels = self.particle_ID[idx_]

            y = int(particle_Y)
            x = int(particle_X)
            sigma = particle_sigma
            ax.add_patch(Circle((x, y), radius=np.sqrt(2) * sigma, edgecolor='r', facecolor='none', linewidth=2))

        plt.show()

    def fit_2D_gussian(self, frame_num, scale=5):

        particle = self.list_psf.loc[self.list_psf['frame'] == frame_num]
        particle_X = particle['x'].tolist()
        particle_Y = particle['y'].tolist()
        particle_sigma = particle['sigma'].tolist()
        particle_center_intensity = particle['center_intensity'].tolist()

        index_list = [index for index in particle.index]

        list_one_frame_fit = []
        if particle.shape[0] > 0:
            for p_x, p_y, sigma_0, c_0, i_ in zip(particle_X, particle_Y, particle_sigma, particle_center_intensity,
                                                  index_list):
                window_size = scale * np.sqrt(2) * sigma_0
                start_sigma = sigma_0
                if p_x > window_size and p_y > window_size:
                    window_frame = self.video[int(frame_num), int(p_y - window_size) + 1:int(p_y + window_size),
                                   int(p_x - window_size) + 1:int(p_x + window_size)]
                    w_s = window_size
                else:
                    window_size_temp = window_size
                    while p_x < window_size_temp or p_y < window_size_temp:
                        window_size_temp = window_size_temp - 2
                    window_frame = self.video[int(frame_num),
                                   int(p_y - window_size_temp) + 1:int(p_y + window_size_temp),
                                   int(p_x - window_size_temp) + 1:int(p_x + window_size_temp)]
                    w_s = window_size_temp

                fit_params_ = gaussian_2D_fit.fit_2D_Gaussian_varAmp(window_frame, sigma_x=start_sigma,
                                                                     sigma_y=start_sigma,
                                                                     display_flag=False)
                fit_params = fit_params_[1]
                fit_errors = fit_params_[2]

                print('Fit Amplitude:', fit_params[0], '\u00b1', fit_errors[0])
                print('Fit X-Center: ', fit_params[1], '\u00b1', fit_errors[1])
                print('Fit Y-Center: ', fit_params[2], '\u00b1', fit_errors[2])
                print('Fit X-Sigma:  ', fit_params[3], '\u00b1', fit_errors[3])
                print('Fit Y-Sigma:  ', fit_params[4], '\u00b1', fit_errors[4])
                print('Fit Bias:  ', fit_params[5], '\u00b1', fit_errors[5])

                fit_params = [frame_num, i_, fit_params_, p_x, p_y, w_s]
                return fit_params


class JupyterSubplotDisplay:

    def __init__(self, list_videos, numRows, numColumns, list_titles=None, imgSizex=20, imgSizey=20, IntSlider_width='500px',
                 median_filter_flag=False, color='gray', step=1, value=0, vmin=None, vmax=None):

        """
        This class interactively displays several videos (with the same number of frames) in a Jupyter notebook.

        Parameters
        ----------
        list_videos: list of NDArray
          List of videos

        numRows: int
            It defines number of rows in sub-display

        numColumns: int
            It defines number of columns in sub-display

        list_titles: list str
            List of titles for each sub plot

        median_filter_flag: bool
          In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
            It defines the colormap for visualization.

        imgSizex: int
            Image length size.

        imgSizey: int
            Image width size.

        IntSlider_width: str
            Size of slider

        step: int
            Stride between visualization frames.

        value: int
            Initial frame value for visualization
        """

        self.color = color
        self.vmin = vmin
        self.vmax = vmax

        self.video = list_videos
        self.median_filter_flag = median_filter_flag
        self.numRows = numRows
        self.numColumns = numColumns
        self.numVideos = len(list_videos)
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey

        if list_titles is None:
            self.list_titles = [None for _ in range(self.numVideos)]
        else:
            self.list_titles = list_titles

        max_numberFrames = np.max([vid_.shape[0] for vid_ in list_videos])

        interact(self.display, frame=widgets.IntSlider(min=0, max=max_numberFrames-1, step=step, value=value, layout=Layout(width=IntSlider_width),
                                                       readout_format='100', continuous_update=False, description='Frame:'))

    def display(self, frame):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))

        grid = plt.GridSpec(self.numRows, self.numColumns, hspace=0.3, wspace=0.3)

        imgGrid_list = []
        for i in range(self.numRows):
            for j in range(self.numColumns):
                imgGrid_list.append(fig.add_subplot(grid[i, j]))

        for img_, tit_, img_grid_ in zip(self.video, self.list_titles, imgGrid_list):

            if self.median_filter_flag:
                frame_v = median_filter(img_[int(frame), :, :], 3)
            else:
                frame_v = img_[int(frame), :, :]

            # img_grid_.imshow(frame_v, cmap=self.color)
            # img_grid_.axis('off')

            divider = make_axes_locatable(img_grid_)
            cax = divider.append_axes("right", size="5%", pad=0.05)
            if self.vmin is None or self.vmax is None:
                myplot = img_grid_.imshow(frame_v, cmap=self.color)
            else:
                myplot = img_grid_.imshow(frame_v, cmap=self.color, vmin=self.vmin, vmax=self.vmax)
            img_grid_.axis('off')
            plt.colorbar(myplot, cax=cax)
            if tit_ is not None:
                img_grid_.set_title(tit_)

        plt.show()





class JupyterPSFs_2_modality_subplotLocalizationDisplay:

    def __init__(self, list_videos, list_df_PSFs_1, list_df_PSFs_2, numRows, numColumns, list_titles=None,
                 median_filter_flag=False, color='gray', imgSizex=5, imgSizey=5, IntSlider_width='500px', step=1, value=0,
                 edgecolor_1='r', edgecolor_2='g'):
        """
        This class will interactively sub-display multiple videos in a Jupyter notebook while highlighting PSFs determined in two modalities.

        Parameters
        ----------
        list_videos: list of NDArray
          List of videos.

        list_df_PSFs_1: list
            list of all data Frames that contains the location of PSFs in first modality.

        list_df_PSFs_2: list
            list of all data Frames that contains the location of PSFs in second modality.

        numRows: int
            It defines number of rows in sub-display.

        numColumns: int
            It defines number of columns in sub-display.

        list_titles: list str
            List of titles for each sub plot.

        median_filter_flag: bool
          In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
          It defines the colormap for visualization.

        imgSizex: int
          Image length size.

        imgSizey: int
          Image width size.

        IntSlider_width: str
          Size of slider

        step: int
          Stride between visualization frames.

        value: int
            Initial frame value for visualization

        edgecolor_1: str
            The color of the circle that annotates the first modality on the display.

        edgecolor_2: str
            The color of the circle that annotates the second modality on the display.

        """
        self.list_video = list_videos
        self.numVideos = len(list_videos)
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.list_df_PSFs_1 = list_df_PSFs_1
        self.list_df_PSFs_2 = list_df_PSFs_2
        self.edgecolor_1 = edgecolor_1
        self.edgecolor_2 = edgecolor_2
        self.median_filter_flag = median_filter_flag
        self.color = color
        self.numRows = numRows
        self.numColumns = numColumns

        if list_titles is None:
            self.list_titles = [None for _ in range(self.numVideos)]
        else:
            self.list_titles = list_titles

        max_numberFrames = np.max([vid_.shape[0] for vid_ in list_videos])
        interact(self.show_psf, frame_number=widgets.IntSlider(min=0, max=max_numberFrames-1, step=step, value=value,
                                                               readout_format='1', continuous_update=False, layout=Layout(width=IntSlider_width),
                                                               description='Frame:'))

    def show_psf(self, frame_number):
        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))
        grid = plt.GridSpec(self.numRows, self.numColumns, hspace=0.3, wspace=0.3)

        imgGrid_list = []
        for i in range(self.numRows):
            for j in range(self.numColumns):
                imgGrid_list.append(fig.add_subplot(grid[i, j]))

        for img_, tit_, img_grid_, df_PSFs_1,  df_PSFs_2 in zip(self.list_video, self.list_titles, imgGrid_list, self.list_df_PSFs_1, self.list_df_PSFs_2):

            if self.median_filter_flag:

                frame_v = median_filter(img_[int(frame_number), :, :], 3)
            else:
                frame_v = img_[int(frame_number), :, :]

            divider = make_axes_locatable(img_grid_)
            cax = divider.append_axes("right", size="5%", pad=0.05)
            myplot = img_grid_.imshow(frame_v, cmap=self.color)
            img_grid_.axis('off')
            plt.colorbar(myplot, cax=cax)
            if tit_ is not None:
                img_grid_.set_title(tit_)

            particle_1 = df_PSFs_1.loc[df_PSFs_1['frame'] == frame_number]
            particle_X_1 = particle_1['x'].tolist()
            particle_Y_1 = particle_1['y'].tolist()
            particle_sigma_1 = particle_1['sigma'].tolist()

            particle_2 = df_PSFs_2.loc[df_PSFs_2['frame'] == frame_number]
            particle_X_2 = particle_2['x'].tolist()
            particle_Y_2 = particle_2['y'].tolist()
            particle_sigma_2 = particle_2['sigma'].tolist()

            for j_ in range(len(particle_X_1)):
                y = int(particle_Y_1[j_])
                x = int(particle_X_1[j_])
                sigma = particle_sigma_1[j_]
                img_grid_.add_patch(Circle((x, y), radius=np.sqrt(2) * sigma, edgecolor=self.edgecolor_1, facecolor='none',
                                    linewidth=2))

            for j_ in range(len(particle_X_2)):
                y = int(particle_Y_2[j_])
                x = int(particle_X_2[j_])
                sigma = particle_sigma_2[j_]
                img_grid_.add_patch(Circle((x, y), radius=np.sqrt(2) * sigma, edgecolor=self.edgecolor_2, facecolor='none',
                                    linewidth=2))

        plt.show()


class JupyterFPNcDisplay:

    def __init__(self, list_videos, list_titles=None, correction_axis=0, numRows=1, numColumns=2,
                 imgSizex=20, imgSizey=20, IntSlider_width='500px',
                 median_filter_flag=False, color='gray', step=1, value=0):
        """
        This class can sub-display several FPNc video in the Jupyter notebook interactively while 1D projection on
        the direction of correction axis illustrates below on each subplot.

        Parameters
        ----------
        list_videos: list of NDArray
          List of videos

        correction_axis: (0/1)
            selecting the axis that FPNc correction should visualize for it.

        numRows: int
            It defines number of rows in sub-display.

        numColumns: int
            It defines number of columns in sub-display.

        list_titles: list str
            List of titles for each sub plot.

        median_filter_flag: bool
            In case it defines as True, a median filter is applied with size 3 to remove hot pixel effect.

        color: str
            It defines the colormap for visualization.

        imgSizex: int
            Image length size.

        imgSizey: int
            Image width size.

        IntSlider_width: str
            Size of slider

        step: int
            Stride between visualization frames.

        value: int
            Initial frame value for visualization
        """
        self.color = color
        self.video = list_videos
        self.median_filter_flag = median_filter_flag
        self.numRows = numRows
        self.numColumns = numColumns
        self.numVideos = len(list_videos)
        self.imgSizex = imgSizex
        self.imgSizey = imgSizey
        self.correction_axis = correction_axis

        if list_titles is None:
            self.list_titles = [None for _ in range(self.numVideos)]
        else:
            self.list_titles = list_titles

        max_numberFrames = np.max([vid_.shape[0] for vid_ in list_videos])

        interact(self.display, frame=widgets.IntSlider(min=0, max=max_numberFrames, step=step, value=value, layout=Layout(width=IntSlider_width),
                                                       readout_format='100', continuous_update=False,
                                                       description='Frame:'))

    def display(self, frame):
        self.fig = plt.figure(figsize=(self.imgSizex, self.imgSizey))

        grid = self.fig.add_gridspec(nrows=self.numRows, ncols=self.numColumns, left=0.05, right=0.48, wspace=0.05)
        axs = grid.subplots(sharex='col', sharey='row')

        for img_, tit_, img_grid_ in zip(self.video, self.list_titles, axs):

            if self.median_filter_flag:
                frame_v = median_filter(img_[int(frame), :, :], 3)
            else:
                frame_v = img_[int(frame), :, :]
            if self.correction_axis == 0:
                img_grid_.imshow(frame_v, cmap=self.color)
            elif self.correction_axis == 1:
                img_grid_.imshow(np.transpose(frame_v), cmap=self.color)
            if tit_ is not None:
                img_grid_.set_title(tit_)

        for ax in axs.flat:
            ax.label_outer()

        plt.show()

        fig = plt.figure(figsize=(self.imgSizex, self.imgSizey/10))
        grid = fig.add_gridspec(nrows=self.numRows, ncols=self.numColumns, left=0.05, right=0.48, wspace=0.05)
        axs = grid.subplots(sharex='col', sharey='row')

        for img_, img_grid_ in zip(self.video, axs):
            frame_v = img_[int(frame), :, :]
            gainMap1D_median = np.mean(frame_v, axis=self.correction_axis)

            x_grid = list(range(0, frame_v.shape[1]))
            y_grid = list(range(0, frame_v.shape[0]))

            x, y = np.meshgrid(x_grid, y_grid)
            if self.correction_axis == 0:
                img_grid_.plot(np.ravel(x), np.ravel(frame_v, order='A'), 'b.', linewidth=1, markersize=0.5)
            elif self.correction_axis == 1:
                img_grid_.plot(np.ravel(x), np.ravel(frame_v, order='F'), 'b.', linewidth=1, markersize=0.5)

            img_grid_.plot(gainMap1D_median, 'r-', linewidth=1, markersize=0.5)
            min_y = frame_v.mean() - 4*frame_v.std()
            max_y = frame_v.mean() + 4*frame_v.std()
            img_grid_.set_ylim(min_y, max_y)
            img_grid_.ticklabel_format(style='sci', axis='y', scilimits=(0, 0))

        # Hide x labels and tick labels for all but bottom plot.
        for ax in axs.flat:
            ax.label_outer()

        plt.show()


