# %matplotlib inline
%matplotlib notebook

from ipywidgets import interact, widgets, Layout
import matplotlib.pyplot as plt
from matplotlib.image import AxesImage
from IPython.display import display
from numpy.fft import fftshift, fft2, ifft2, ifftshift
import time
from skimage import transform

def makeGaussian(size, fwhm = 3, center=None):
    if isinstance(size, int):
        size = (size, size)
    if center:
        assert isinstance(center, tuple)
    """
    if size is tuple - (y_size, x_size)
    if center is tuple - (y_size, x_size)

    """

    y, x = np.meshgrid(np.arange(size[1]), np.arange(size[0]))

    if center is None:
        x0 = size[1] // 2
        y0 = size[0] // 2
    else:
        x0 = center[1]
        y0 = center[0]
    notnorm_gaus = np.exp(-((x-x0)**2 + (y-y0)**2) / fwhm**2)
    return notnorm_gaus/notnorm_gaus.max()

class Image_analizer(object):
    '''
    Allow to make fast analysis of image among ferquences and dynamic range.
    Image must be grayscale, nunmpy array, 2D or 3D(time domain)
    '''
    def __init__(self, figsize=(6,6), imgsize=(512,512)):
        self.dynamic_range = None
        self.seq_frame = None
        self.hight_freq_pass = None
        self.low_freq_pass = None
        self.figsize = figsize
        self.imgsize = imgsize

    def process_frame(self, frame):
        '''
        produce 2D image processed according with selected parameters
        '''
        if not self.hight_freq_pass == 0:
            gaussian = makeGaussian(frame.shape, self.hight_freq_pass)
            fft_data = fftshift(fft2(ifftshift(frame)))
            filtered_fft = fft_data*gaussian
            frame = np.abs(fftshift(ifft2(ifftshift(filtered_fft))))
        if not self.low_freq_pass == 0:
            gaussian = makeGaussian(frame.shape, self.low_freq_pass)
            fft_data = fftshift(fft2(ifftshift(frame)))
            filtered_fft = fft_data*(1-gaussian)
            frame = np.abs(fftshift(ifft2(ifftshift(filtered_fft))))
        frame = np.clip(frame,*bounds)
        return frame

    def show(self, data):
        '''
        call iteractive tool, for frequnces and dynamic range selection
        '''
        if len(data.shape) == 2:
            data = data[np.newaxis, ...]
        if self.dynamic_range is None:
            self.dynamic_range = (0, data.max())
        else:
            self.dynamic_range = tuple(np.clip(self.dynamic_range, 0, data.max()))
        if self.seq_frame is None:
            self.seq_frame = 0
        else:
            self.seq_frame = np.clip(self.seq_frame, 0, len(data))
        if self.hight_freq_pass is None:
            self.hight_freq_pass = 0
        if self.low_freq_pass is None:
            self.low_freq_pass = 0
        print(data.min(), data.max())
        print(self.dynamic_range)
        data = transform.resize(data, (data.shape[0],*self.imgsize), preserve_range=True)

        def update(bounds, step, hp, lp):
            '''
            called each widget position update
            '''
            self.hight_freq_pass = hp
            self.low_freq_pass = lp
            self.seq_frame = step
            self.dynamic_range = bounds
            tic = time.time()
            d = data[step]
            if not hp==0:
                gaussian = makeGaussian(data.shape[1:], hp)
                fft_data = fftshift(fft2(ifftshift(d)))
                filtered = fft_data*gaussian
                d = np.abs(fftshift(ifft2(ifftshift(filtered))))
            if not lp==0:
                gaussian = makeGaussian(data.shape[1:], lp)
                fft_data = fftshift(fft2(ifftshift(d)))
                filtered = fft_data*(1-gaussian)
                d = np.abs(fftshift(ifft2(ifftshift(filtered))))
                d = d - d.min()
            img_obj.set_data(d)
            img_obj_init.set_data(data[step])
            img_obj.set_clim(*bounds)
            fig.canvas.draw()
    #         display(fig)
            print('Processing time took {}'.format(time.time() - tic), end='\r')

        slider_r = widgets.FloatRangeSlider(
            value=[self.dynamic_range[0], self.dynamic_range[1]],
            min=0,#data.min(),
            max=data.max(),
            step=(data.max()-data.min())/1000,
            description='range:',
            disabled=False,
            continuous_update=False,
            orientation='horizontal',
            readout=True,
            readout_format='.1f',
            layout=Layout(width='80%')
        )

        slider_s = widgets.IntSlider(
            value=self.seq_frame,
            min=0,
            max=len(data)-1,
            step=1,
            description='step',
            disabled=False,
            continuous_update=False,
            orientation='horizontal',
            readout=True,
            readout_format='d',
            layout=Layout(width='80%')
        )

        slider_hp = widgets.FloatSlider(
            value=self.hight_freq_pass,
            min=0,
            max=data.shape[1],
            step=data.shape[1]/1000,
            description='hight_pass',
            disabled=False,
            continuous_update=False,
            orientation='horizontal',
            readout=True,
            readout_format='f',
            layout=Layout(width='80%')
        )

        slider_lp = widgets.FloatSlider(
            value=self.low_freq_pass,
            min=0,
            max=data.shape[1],
            step=data.shape[1]/1000,
            description='low_pass',
            disabled=False,
            continuous_update=False,
            orientation='horizontal',
            readout=True,
            readout_format='f',
            layout=Layout(width='80%')
        )
        fig, ax = plt.subplots(1,2, figsize=self.figsize)
        img_obj = ax[0].imshow(data[0], cmap='gray')
        img_obj_init = ax[1].imshow(data[0], cmap='gray')
    #     plt.close()
        interact(update, bounds=slider_r, step=slider_s, hp=slider_hp, lp=slider_lp)