import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
from mpl_toolkits.axes_grid1 import ImageGrid

import scipy as sp
from scipy import ndimage
from scipy import fftpack as fft

import hwt

dist = 501
hist = np.zeros((dist, dist), dtype=float)
midpoint = np.floor(dist/2)
hist[midpoint, midpoint] = 1
sigma = [40, 30]
xrot = 60
hist2 = hwt.smoothers.anisotropic_gauss(hist.copy(), sigma[0], sigma[1], xrot, 
                                        0, 0, 1, 10, True)
                                        
print(hist.sum(), hist2.sum())
fig = plt.figure(figsize=(8,8))
grid = ImageGrid(fig, 111, nrows_ncols=(1,1), direction="column",
                 axes_pad=1, add_all=True, label_mode="1",
                 share_all=True, cbar_location="bottom", cbar_mode="each",
                 cbar_size="5%", cbar_pad="5%")
ax0 = grid[0]; cax0 = grid.cbar_axes[0]
cbar0 = ax0.pcolormesh(hist2)
cax0.colorbar(cbar0)
ax0.set_xlim(0, dist)
ax0.set_ylim(0, dist)


dist = 501
midpoint = np.floor(dist/2)
print(midpoint)
hist = np.zeros((dist, dist), dtype=float)
hist[midpoint, midpoint] = 1
# x = np.random.randint(0, dist, (20,))
# y = np.random.randint(0, dist, (20,))
# hist[x, y] = 1

ihist = fft.fft2(hist.copy())
ihist2 = fft.fft2(hist2.copy())
ihist3 = ihist * ihist2
hist3 = fft.ifft2(ihist3)


fig = plt.figure(figsize=(8,8))
grid = ImageGrid(fig, 111, nrows_ncols=(1,1), direction="column",
                 axes_pad=1, add_all=True, label_mode="1",
                 share_all=True, cbar_location="bottom", cbar_mode="each",
                 cbar_size="5%", cbar_pad="5%")
ax0 = grid[0]; cax0 = grid.cbar_axes[0]
cbar0 = ax0.pcolormesh(hist3)
cax0.colorbar(cbar0)
ax0.set_xlim(0, dist)
ax0.set_ylim(0, dist)