import numpy as np
import scipy
import scipy.fftpack as sfft
import scipy.signal
import tqdm
from cache import timeit


class SpinImage_Two_Phase:
    resolution = 0.05
    offset = 40

    @timeit
    def __init__(self, x_pos_low, y_pos_low, x_pos_high, y_pos_high):
        assert x_pos_low.shape == y_pos_low.shape
        assert x_pos_high.shape == y_pos_high.shape
        assert x_pos_low.shape == x_pos_high.shape

        offset_shift = self.offset * self.resolution

        x_pos_low = x_pos_low - np.min(x_pos_low) + offset_shift
        y_pos_low = y_pos_low - np.min(y_pos_low) + offset_shift
        x_pos_high = x_pos_high - np.min(x_pos_high) + offset_shift
        y_pos_high = y_pos_high - np.min(y_pos_high) + offset_shift
        max_len_x = np.maximum(np.max(x_pos_low), np.max(x_pos_high))
        max_len_y = np.maximum(np.max(y_pos_low), np.max(y_pos_high))
        self.length_x = max_len_x + (self.offset + 1) * self.resolution
        self.length_y = max_len_y + (self.offset + 1) * self.resolution
        self.x_ind = np.arange(0, self.length_x, self.resolution)  # angstrom
        self.y_ind = np.arange(0, self.length_y, self.resolution)  # angstrom
        self.x_low, self.y_low = self._stuff(x_pos_low, y_pos_low)
        self.x_high, self.y_high = self._stuff(x_pos_high, y_pos_high)

        X, Y = np.meshgrid(self.x_ind, self.y_ind, indexing="ij")
        print(X.shape, self.x_low.flatten().max(), self.y_low.flatten().max())
        sigma = .1
        self.low_list = []
        for x, y, x_ind, y_ind in zip(
            x_pos_low.flatten(),
            y_pos_low.flatten(),
            self.x_low.flatten(),
            self.y_low.flatten(),
        ):
            xl, yl, xu, yu = self.clean_bounds(
                x_ind - self.offset,
                y_ind - self.offset,
                x_ind + self.offset,
                y_ind + self.offset,
                X.shape,
            )
            self.low_list.append(np.exp(-0.5 * ((X[xl:xu, yl:yu] - x) ** 2 +
                                                (Y[xl:xu, yl:yu] - y) ** 2) / sigma**2))
        self.low_list = np.array(
            self.low_list, dtype=object)
        #print("BEFOR:", self.low_list.shape)
        #self.low_list = self.low_list.reshape((*self.x_high.shape,80,80))
        #print("HIER:", self.low_list.shape)
        self.high_list = []
        for x, y, x_ind, y_ind in zip(
            x_pos_high.flatten(),
            y_pos_high.flatten(),
            self.x_high.flatten(),
            self.y_high.flatten(),
        ):
            xl, yl, xu, yu = self.clean_bounds(
                x_ind - self.offset,
                y_ind - self.offset,
                x_ind + self.offset,
                y_ind + self.offset,
                X.shape,
            )
            self.high_list.append(np.exp(-0.5 * ((X[xl:xu, yl:yu] - x) ** 2 +
                                                 (Y[xl:xu, yl:yu] - y) ** 2) / sigma**2))
        self.high_list = np.array(self.high_list, dtype=object)

    def clean_bounds(self, xl, yl, xu, yu, shape):
        if xl < 0:
            xl = 0
        if yl < 0:
            yl = 0

        if xu > shape[0]:
            xu = shape[0]
        if yu > shape[1]:
            yu = shape[1]

        return xl, yl, xu, yu

    def _stuff(self, pos_x, pos_y):
        xind = np.searchsorted(self.x_ind, pos_x).astype(int)
        yind = np.searchsorted(self.y_ind, pos_y).astype(int)
        return xind, yind

    @timeit
    def apply_mask(self, maske):
        mask = np.empty_like(self.x_high, dtype=bool)
        mask[0::2, 0::2] = maske
        mask[1::2, 0::2] = maske
        mask[0::2, 1::2] = maske
        mask[1::2, 1::2] = maske
        mask = mask.flatten()
        self.img = np.zeros((self.x_ind.size, self.y_ind.size))
        print(self.img.shape)
        for x, y, dat in zip(
            self.x_high.flatten()[mask],
            self.y_high.flatten()[mask],
            self.high_list[mask],
        ):
            xl, yl, xu, yu = self.clean_bounds(
                x - self.offset,
                y - self.offset,
                x + self.offset,
                y + self.offset,
                self.img.shape,
            )
            self.img[x - self.offset: x + self.offset,
                     y - self.offset: y + self.offset] = dat
            self.img[x, y] += 1
        mask = np.invert(mask)
        for x, y, dat in zip(
            self.x_high.flatten()[mask],
            self.y_high.flatten()[mask],
            self.high_list[mask],
        ):
            xl, yl, xu, yu = self.clean_bounds(
                x - self.offset,
                y - self.offset,
                x + self.offset,
                y + self.offset,
                self.img.shape,
            )
            if self.img[xl:xu, yl:yu].shape == dat.shape:
                self.img[xl:xu, yl:yu] = dat
            else:
                self.img[x, y] += 1

    @timeit
    def fft(self):
        Z_fft = sfft.fft2(self.img)
        Z_shift = sfft.fftshift(Z_fft)
        fft_freqx = sfft.fftfreq(self.img.shape[0], self.resolution)
        fft_freqy = sfft.fftfreq(self.img.shape[1], self.resolution)
        fft_freqx_clean = sfft.fftshift(fft_freqx)
        fft_freqy_clean = sfft.fftshift(fft_freqy)
        return fft_freqx_clean, fft_freqy_clean, np.abs(Z_shift) ** 2

    @timeit
    def plot(self, ax, scale=None):
        if scale is None:
            ax.imshow(self.img)
        else:
            quad = np.ones((int(scale / self.resolution),
                           int(scale / self.resolution)))
            img = scipy.signal.convolve2d(self.img, quad)
            ax.imshow(img)

    @timeit
    def pad_it_square(self, additional_pad=0, size=None):
        h = self.img.shape[0]
        w = self.img.shape[1]
        xx = np.maximum(h, w) + 2 * additional_pad + 1
        if size is not None:
            xx = np.maximum(xx, size)
        yy = xx

        a = (xx - h) // 2
        aa = xx - a - h

        b = (yy - w) // 2
        bb = yy - b - w

        self.img = np.pad(self.img, pad_width=(
            (a, aa), (b, bb)), mode="constant")

    def gaussian(self, sigma):
        x = np.arange(-self.length_x / 2, self.length_x / 2, self.resolution)
        y = np.arange(-self.length_y / 2, self.length_y / 2, self.resolution)
        X, Y = np.meshgrid(x, y)
        z = 1 / (2 * np.pi * sigma * sigma) * \
            np.exp(-(X**2 / (2 * sigma**2) + Y**2 / (2 * sigma**2)))
        self.img = np.multiply(self.img, z.T)


class SpinImage:
    resolution = 0.05

    def __init__(self, x_pos, y_pos):
        x_pos = x_pos - np.min(x_pos)
        y_pos = y_pos - np.min(y_pos)
        self.length_x = np.max(x_pos) + self.resolution
        self.length_y = np.max(y_pos) + self.resolution
        self.img = self.image_from_pos(x_pos, y_pos)

    def image_from_pos(self, pos_x, pos_y):
        x_ind = np.arange(0, self.length_x, self.resolution)  # angstrom
        y_ind = np.arange(0, self.length_y, self.resolution)  # angstrom
        img = np.zeros((x_ind.size, y_ind.size))
        xind = np.searchsorted(x_ind, pos_x)
        yind = np.searchsorted(y_ind, pos_y)
        img[xind, yind] = 1
        return img

    def image_from_pos_with_gauss(self, pos_x, pos_y, sigma=1):
        x_ind = np.arange(0, self.length_x, self.resolution)  # angstrom
        y_ind = np.arange(0, self.length_y, self.resolution)  # angstrom
        img = np.zeros((x_ind.size, y_ind.size))
        X, Y = np.meshgrid(y_ind, x_ind)
        for px, py in tqdm.tqdm(zip(pos_x.flatten(), pos_y.flatten())):
            img += np.exp(-0.5 * ((X - px) ** 2 + (Y - py) ** 2) / sigma**2)
        return img

    def fft(self):
        Z_fft = sfft.fft2(self.img)
        Z_shift = sfft.fftshift(Z_fft)
        fft_freqx = sfft.fftfreq(self.img.shape[0], self.resolution)
        fft_freqy = sfft.fftfreq(self.img.shape[1], self.resolution)
        fft_freqx_clean = sfft.fftshift(fft_freqx)
        fft_freqy_clean = sfft.fftshift(fft_freqy)
        return fft_freqx_clean, fft_freqy_clean, np.abs(Z_shift) ** 2

    def pad_it_square(self, additional_pad=0, size=None):
        h = self.img.shape[0]
        w = self.img.shape[1]
        xx = np.maximum(h, w) + 2 * additional_pad + 1
        if size is not None:
            xx = np.maximum(xx, size)
        yy = xx
        self.length_x = xx * self.resolution
        self.length_y = yy * self.resolution

        a = (xx - h) // 2
        aa = xx - a - h

        b = (yy - w) // 2
        bb = yy - b - w

        self.img = np.pad(self.img, pad_width=(
            (a, aa), (b, bb)), mode="constant")

    def pad_it(self, x_size, y_size):
        h = self.img.shape[0]
        w = self.img.shape[1]
        xx = x_size
        yy = y_size
        self.length_x = xx * self.resolution
        self.length_y = yy * self.resolution

        a = (xx - h) // 2
        aa = xx - a - h

        b = (yy - w) // 2
        bb = yy - b - w

        self.img = np.pad(self.img, pad_width=(
            (a, aa), (b, bb)), mode="constant")

    def percentage_gaussian(self, mask, sigma):
        x = np.linspace(-self.length_x / 2, self.length_x / 2, mask.shape[0])
        y = np.linspace(-self.length_y / 2, self.length_y / 2, mask.shape[1])
        X, Y = np.meshgrid(x, y)
        z = 1 / (2 * np.pi * sigma * sigma) * \
            np.exp(-(X**2 / (2 * sigma**2) + Y**2 / (2 * sigma**2)))
        return np.multiply(mask, z.T)

    def gaussian(self, sigma):
        x = np.arange(-self.length_x / 2, self.length_x / 2, self.resolution)
        y = np.arange(-self.length_y / 2, self.length_y / 2, self.resolution)
        X, Y = np.meshgrid(x, y)
        z = 1 / (2 * np.pi * sigma * sigma) * \
            np.exp(-(X**2 / (2 * sigma**2) + Y**2 / (2 * sigma**2)))
        self.img = np.multiply(self.img, z.T)

    def plot(self, ax, scale=None):
        if scale is None:
            ax.imshow(self.img)
        else:
            quad = np.ones((int(scale / self.resolution),
                           int(scale / self.resolution)))
            img = scipy.signal.convolve2d(self.img, quad)
            ax.imshow(img)

    def blur(self, sigma):
        self.img = scipy.ndimage.gaussian_filter(self.img, sigma)