from skued import diffread, bragg_peaks_persistence, brillouin_zones
import matplotlib.pyplot as plt
import numpy as np

im = diffread("data/ab_initio_monolayer_mos2.tif")
center = np.array([im.shape[0]//2, im.shape[1]//2])
peaks_persistence, _, _, _ = bragg_peaks_persistence(im, prominence=0.1)
BZ = brillouin_zones(im, mask=np.ones(im.shape), peaks=peaks_persistence.astype(int), center=center.astype(int))
BZ.getVisibleBZs(symmetry=6)
EQUIV = BZ.getEquivalentScatteringVectors(np.array([119, 174]), use_visible=True)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(6, 3))
for ax in [ax1, ax2]:
        ax.imshow(im, origin='lower', interpolation='bicubic', vmin=-1, vmax=1)
        BZ.renderVisibleBZs(ax, **dict(edgecolor='white'))

        ax.axis('off')

ax2.scatter(EQUIV[:, 0], EQUIV[:, 1], marker="x", color="r", s=20)

plt.tight_layout()
plt.show()