from skued import diffread, align
from pathlib import Path

ref = diffread("data/Cr_1.tif")
im = diffread("data/Cr_2.tif")

mask = np.ones_like(ref, dtype=bool)
mask[0:1250, 975:1225] = False

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

aligned = align(image=im, reference=ref, mask=mask)

fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(nrows=2, ncols=3, figsize=(9, 6))
ax1.imshow(ref, vmin=0, vmax=200, cmap='inferno')
ax2.imshow(im, vmin=0, vmax=200, cmap='inferno')
ax3.imshow(ref - im, cmap="RdBu_r")
ax4.imshow(np.logical_not(mask) * ref, vmin=0, vmax=200, cmap="inferno")
ax5.imshow(aligned, vmin=0, vmax=200, cmap='inferno')
ax6.imshow(ref - aligned, cmap="RdBu_r")

for ax in (ax1, ax2, ax3, ax4, ax5, ax6):
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)

ax1.set_title("Reference")
ax2.set_title("Data")
ax3.set_title("Difference")
ax4.set_title("Mask")
ax5.set_title("Aligned data")
ax6.set_title("Difference after alignment")

plt.tight_layout()