import abtem
import ase
import matplotlib.pyplot as plt
import numpy as np

abtem.config.set({"visualize.cmap": "viridis"})
abtem.config.set({"visualize.continuous_update": True})
abtem.config.set({"visualize.autoscale": True})
#abtem.config.set({"visualize.reciprocal_space_units": "mrad"})
abtem.config.set({"device": "cpu"})
abtem.config.set({"fft": "fftw"});

from matplotlib.patches import Circle
from ipywidgets import HBox, VBox, widgets, interact, Dropdown, Label, Layout

unit_cell = ase.build.bulk("Au", cubic=True)

atoms = unit_cell * (1, 1, 30)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
abtem.show_atoms(atoms, ax=ax1, title="Beam view")
abtem.show_atoms(atoms, ax=ax2, plane="xz", title="Side view", linewidth=0);

plane_wave = abtem.PlaneWave(energy=200e3)

potential = abtem.Potential(atoms, slice_thickness=4.08 / 4, sampling=0.04, exit_planes=1, projection='finite')

exit_wave = plane_wave.multislice(potential).compute()

Cs = -20e-6 * 1e10

ctf = abtem.CTF(Cs=Cs, energy=200e3)

ctf.scherzer_defocus

ctf.defocus = ctf.scherzer_defocus
ctf.semiangle_cutoff = ctf.crossover_angle
ctf.focal_spread = 10
ctf.angular_spread = 1

image = exit_wave.apply_ctf(ctf).intensity().compute()

%matplotlib ipympl

with plt.ioff():
    dpi = 72
    fig, axs = plt.subplots(1,3, figsize=(675/dpi, 400/dpi), dpi=dpi)

    
im0_vmax = potential[:].array.sum(0).max()/2
im0_vmin = 0 

im1_vmax = (np.abs(exit_wave[1:][:].array)**2).max()/2
im1_vmin = 0 

im2_vmax = image[1:][:].array.max()
im2_vmin = image[1:][:].array.min()

potential.build().compute()
exit_wave.compute()
image.compute()

pots = potential.to_images()[0].show(ax=axs[0], vmin = im0_vmin, vmax = im0_vmax)
wavs = exit_wave[1:][0].show(ax=axs[1], vmin = im1_vmin, vmax = im1_vmax)
imgs = image[1:][0].show(ax=axs[2], vmin = im2_vmin, vmax = im2_vmax)

#pots = axs[0].imshow(potential[0].array[0])
#wavs = axs[1].imshow(exit_wave[1:][0].intensity().array)
#imgs = axs[2].imshow(image[1:][0].array)

axs[0].set_title('Cumulative potential')
axs[1].set_title('Exit wave')
axs[2].set_title('Image')

axs[1].set_yticks([])
axs[2].set_yticks([])
axs[1].set_ylabel("")
axs[2].set_ylabel("")

plt.tight_layout()  

im0 = axs[0].get_images()[0]
im1 = axs[1].get_images()[0]
im2 = axs[2].get_images()[0]

# interact
def update_images(plane):    
    # axs[0].cla()
    # axs[1].cla()
    # axs[2].cla()

    # pots._artists[0,0].set_data(potential[plane].array)
    # wavs._artists[0,0].set_data(exit_wave[1:][plane].array)
    # imgs._artists[0,0].set_data(image[1:][plane].array)

    #pots.set_data(potential[plane].array[0])
    #wavs.set_data(exit_wave[1:][plane].intensity().array)
    #imgs.set_data(image[1:][plane].array)
    
    im0.set_data(potential[0:plane].array.sum(0))
    im1.set_data(np.abs(exit_wave[:-1][plane].array)**2)
    im2.set_data(image[:-1][plane].array)
    
    im0.set_clim(vmax = im0_vmax, vmin = im0_vmin)
    im1.set_clim(vmax = im1_vmax, vmin = im1_vmin)
    im2.set_clim(vmax = im2_vmax, vmin = im1_vmin)
    
    fig.canvas.draw_idle()

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

layout = Layout(width='250px',height='30px')

plane = widgets.IntSlider(
    value = 0, max = len(potential),
    step = 1,
    description = "Slice",
    style = style,
    layout = layout,
)

widgets.interactive_output(
    update_images, 
    {
        'plane':plane,
    },
)

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 = '400px'
fig.canvas.toolbar_position = 'bottom'

widget = widgets.VBox(
    [
        fig.canvas,
            HBox(
                [plane]
            )
    ],
)

display(widget);

diffraction = exit_wave.diffraction_patterns()

indexed_spots = diffraction[1:].crop(120).index_diffraction_spots(cell=unit_cell)

%matplotlib ipympl

diffraction.compute()
indexed_spots.compute()

with plt.ioff():
    dpi = 72
    fig, axs = plt.subplots(1,2, figsize=(675/dpi, 375/dpi), dpi=dpi)

diff = diffraction[1:].block_direct().crop(120).show(ax=axs[0])#, cbar=True)
spts = indexed_spots[1:].block_direct().crop(120).show(ax=axs[1],
        scale=0.15,
        #cbar=True,
        annotation_kwargs={"threshold": 0.01, "fontsize": 7},
        vmin = 0, 
        vmax = 0.01,                                         
    )

axs[0].set_title('Diffraction pattern')
axs[1].set_title('Indexed reflections')

#axs[1].set_yticks([])
axs[1].set_ylabel("")

plt.tight_layout()  

# ax0_loc = fig.axes[0].get_position()
# ax1_loc = fig.axes[1].get_position()

# interact
def update_images(plane):    
    axs[0].cla()
    axs[1].cla()
    
    #for ax in list(fig.axes):  
    #    if ax.get_label() == '<colorbar>':
    #        fig.delaxes(ax)
     
#     fig.axes[0].set_position(ax0_loc)
#     fig.axes[1].set_position(ax1_loc)
    
    diff = diffraction[(plane+1):].block_direct().crop(120).show(ax=axs[0])
    spts = indexed_spots[plane:].block_direct().crop(120).show(
        ax=axs[1], 
        scale=0.15,
        #cbar=True,
        annotation_kwargs={"threshold": 0.001, "fontsize": 7},
        vmin = 0, 
        vmax = 0.01,
        figsize = (8,8)
    )

    axs[1].set_yticks([])
    axs[1].set_ylabel("")
    
    axs[0].set_title('Diffraction pattern')
    axs[1].set_title('Indexed reflections')


    fig.canvas.draw_idle()
    return None

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

layout = Layout(width='250px',height='30px')

plane = widgets.IntSlider(
    value = 0, max = len(potential)-1,
    step = 4,
    description = "Slice",
    style = style,
    layout = layout,
)

widgets.interactive_output(
    update_images, 
    {
        'plane':plane,
    },
)

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'

widget2 = widgets.VBox(
    [
        fig.canvas,
            HBox(
                [plane]
            )
    ],
)

display(widget2);