import numpy as np
import scipy.fftpack as sfft


class SpinImage:
    resolution = 0.1

    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 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):
        h = self.img.shape[0]
        w = self.img.shape[1]
        print(h, w)
        xx = np.maximum(h, w) + 2 * additional_pad
        yy = xx
        self.length_x = xx * self.resolution
        self.length_y = yy * self.resolution
        print("Pad to: ", xx, yy)

        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)

    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)