%matplotlib widget

import numpy as np
import matplotlib.pyplot as plt

from IPython.display import display
from ipywidgets import VBox, interact,IntSlider, Layout, widgets
import abtem
from ase.io import read
from scipy.ndimage import gaussian_filter

#build cell
C = read('data/C.cif')
C = abtem.orthogonalize_cell(C)
C = C * (28, 16, 1)
del C[C.positions[:,2 ]>3]
C.center(vacuum = 0)
C.center(vacuum = 3, axis = 2)

#simulate
potential = abtem.Potential(
    C, sampling=0.2, projection = "finite"
)
probe = abtem.Probe(
    energy=80e3, 
    semiangle_cutoff=2
)
probe.grid.match(potential)
point_scan = abtem.CustomScan((15,15))
measurement = probe.scan(potential, detectors=abtem.PixelatedDetector(max_angle="valid"), scan=point_scan).compute();
probe_show = np.abs(probe.build().array)**2

#show
with plt.ioff():
    dpi = 72
    fig = plt.figure(figsize=(675/dpi, 375/dpi), dpi=dpi)
    
ax0 = fig.add_axes([0.05,  0.15,  0.35, 0.75])
ax1 = fig.add_axes([0.55,  0.15,  0.35, 0.75])
im0 = ax0.imshow(
    measurement.array ** 0.5, 
    cmap = "turbo"
)

im1 = ax1.imshow(probe_show, cmap = "gray")
ax1.set_xticks([]);
ax1.set_yticks([]);

ax0.set_xticks([]);
ax0.set_yticks([]);

ax0.set_title('graphene diffraction pattern');
ax1.set_title('probe real space intensity');



def update_ims(convergence_angle):
    probe = abtem.Probe(
        energy=80e3, 
        semiangle_cutoff=convergence_angle
    )
    probe.grid.match(potential)
    point_scan = abtem.CustomScan((15,15))
    measurement = probe.scan(potential, detectors=abtem.PixelatedDetector(max_angle="valid"), scan=point_scan).compute();
    probe_show = np.abs(probe.build().array)**2

    
    im0.set_data(
        measurement.array ** 0.5,
    )
    im0.set_clim(
        vmax = (measurement.array** 0.5).max(), 
        vmin = (measurement.array**0.5).min()
    )
    im1.set_data(probe_show)
    im1.set_clim(
        vmax = probe_show.max(), 
        vmin = probe_show.min()
    )

    fig.canvas.draw_idle()

style = {
    'description_width': 'initial',
}

convergence_angle = IntSlider(
    value = 2, min = 2, max = 30, 
    step = 1,
    description = "convergence angle (mrad)",
    style = style
)


widgets.interactive_output(
    update_ims, 
    {'convergence_angle':convergence_angle},
)

fig.canvas.resizable = False
fig.canvas.header_visible = False
fig.canvas.footer_visible = False
fig.canvas.toolbar_visible = True
fig.canvas.layout.width = '675px'
fig.canvas.layout.height = '375px'
fig.canvas.toolbar_position = 'bottom'

widget = VBox(
    [
        fig.canvas,
        convergence_angle
    ],
)

display(widget);