import numpy as np
import matplotlib.pyplot as plt

from astropy.io import fits
from astropy.utils.data import get_pkg_data_filename
from astropy.convolution import Gaussian2DKernel, interpolate_replace_nans

# Load the data from data.astropy.org
filename = get_pkg_data_filename('galactic_center/gc_msx_e.fits')

hdu = fits.open(filename)[0]
img = hdu.data[50:90, 60:100] * 1e5

rng = np.random.default_rng(1379)
indices = rng.integers(low=0, high=img.size, size=300)

sampled_data = img.flat[indices]

# Build a new, sparsely sampled version of the original image
new_img = np.tile(np.nan, img.shape)
new_img.flat[indices] = sampled_data

# We smooth with a Gaussian kernel with x_stddev=1 (and y_stddev=1)
# It is a 9x9 array
kernel = Gaussian2DKernel(x_stddev=1)

# create a "reconstructed" image with NaNs replaced by interpolated values
reconstructed_image = interpolate_replace_nans(new_img, kernel)

# Now we do a bunch of plots.  In the first two plots, the originally masked
# values are marked with red X's
fig, axs = plt.subplots(ncols=3, figsize=(12, 6))

ax1 = axs[0]
im = ax1.imshow(img, vmin=-2., vmax=2.e1, origin='lower',
                interpolation='nearest', cmap='viridis')
y, x = np.where(np.isnan(img))
ax1.set(
    autoscale_on=False,
    title="Original",
    xticklabels=[],
    yticklabels=[],
)

ax2 = axs[1]
im = ax2.imshow(new_img, vmin=-2., vmax=2.e1, origin='lower',
                interpolation='nearest', cmap='viridis')
ax2.set(
    title="Sparsely Sampled",
    xticklabels=[],
    yticklabels=[],
)

ax3 = axs[2]
im = ax3.imshow(reconstructed_image, vmin=-2., vmax=2.e1, origin='lower',
                interpolation='nearest', cmap='viridis')
ax3.set(
    title="Reconstructed",
    xticklabels=[],
    yticklabels=[],
)