from skued import diffread, bragg_peaks
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Circle

im1 = diffread("data/graphite.tif")
mask1 = diffread("data/graphite_mask.tif").astype(bool)

im2 = diffread("data/ewald_walkoff.tif")
mask2 = diffread("data/ewald_walkoff_mask.tif").astype(bool)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6, 3))

for ax, im, mask, vmax, radius in zip(
        [ax1, ax2], [im1, im2], [mask1, mask2], [200, 4000], [15, 4]
):
        # Reduce size of images because of memory usage of ReadTheDocs
        im = im[::2, ::2]
        mask = mask[::2, ::2]

        peaks = bragg_peaks(im, mask)

        ax.imshow(im, vmin=0, vmax=vmax, cmap="inferno")

        for r, c in peaks:
                ax.add_patch(Circle(xy=(c, r), radius=radius, ec="r", fc="none"))

        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)

plt.tight_layout()
plt.show()