import matplotlib.pyplot as plt
from skued import nfold, diffread, autocenter
import numpy as np

image = diffread('data/graphite.tif')

mask = np.ones_like(image, dtype = bool)
mask[1100::, 442:480] = False # Artifact line
mask[0:1260, 900:1140] = False # beamblock

# Reduce size of images because of memory usage of ReadTheDocs
image = image[::3, ::3]
mask = mask[::3, ::3]

yc, xc = autocenter(image, mask=mask)
av = nfold(image, mod = 6, center = (xc, yc), mask = mask)

fig , (ax1, ax2, ax3) = plt.subplots(1,3, figsize = (9,3))
ax1.imshow(image, vmin = 0, vmax = 150, cmap='inferno')
ax2.imshow(np.logical_not(mask) * image, vmin = 0, vmax = 150, cmap='inferno')
ax3.imshow(av, vmin = 0, vmax = 150, cmap='inferno')

for ax in (ax1, ax2, ax3):
        ax.xaxis.set_visible(False)
        ax.yaxis.set_visible(False)

ax1.set_title('Graphite')
ax2.set_title('Mask')
ax3.set_title('Averaged')

plt.tight_layout()
plt.show()