%matplotlib widget
import abtem
from py4DSTEM.visualize import return_scaled_histogram_ordering as histogram_ordering, add_scalebar, Complex2RGB

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

from IPython.display import display
import ipywidgets

abtem.config.set({"dask.lazy":False});

file_name = "apoF-ice-embedded-potential-binned.npy"
binned_volume_zxy = np.load("data/"+file_name)
projected_potential = np.sum(binned_volume_zxy,axis=0)
num_slices, nx, ny = binned_volume_zxy.shape

# constants
semiangle = 4  # mrad
energy = 300e3
pixel_size = 2 / 3
bin_factor_xy = 2
bin_factor_z = 6

# initial input values

defocus = 0 # Ang
electrons_per_area = 10 # e/Ang^2
zernike_radius = 0.005 # mrad
zernike_phase_shift = -np.pi/2 #rad
show_zernike = False

sampling = (pixel_size * bin_factor_xy,) * 2
thicknesses = np.tile(pixel_size * bin_factor_z,num_slices)
# thicknesses += defocus / num_slices

potential = abtem.PotentialArray(
    binned_volume_zxy,
    slice_thickness = thicknesses,
    sampling= sampling,
)

exit_wave_raw = abtem.PlaneWave(
    energy=energy,
).match_grid(
    potential
).multislice(
    potential,
    lazy=False
)

exit_wave = exit_wave_raw.apply_ctf(defocus=defocus*1e4)

angular_sampling = exit_wave.angular_sampling[0]

array_to_mutate = [exit_wave]

zernike_kernel = abtem.transfer.Zernike(
    semiangle_cutoff = semiangle,
    center_hole_cutoff = zernike_radius,
    phase_shift = zernike_phase_shift,
    energy=energy,
).match_grid(
    potential
)

alpha, phi = zernike_kernel._angular_grid(device='cpu')

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

scalebar_dict = {'pixelsize':sampling[1]/10,'pixelunits':'nm',"Nx":nx,"Ny":ny,"labelsize":10}

# potential
scaled_potential, _, _ = histogram_ordering(projected_potential)

axs[0].imshow(scaled_potential,cmap='magma')
axs[0].set_title("projected sample potential",fontsize=12)
axs[0].axis("off")
add_scalebar(axs[0],scalebar_dict)

# HRTEM
noisy_hrtem = exit_wave.intensity().poisson_noise(electrons_per_area).array
scaled_hrtem, _, _ = histogram_ordering(noisy_hrtem,normalize=True)

im_hrtem = axs[1].imshow(scaled_hrtem,cmap='gray')
axs[1].set_title("HRTEM intensity",fontsize=12)
axs[1].axis("off")

# Zernike
noisy_zernike = exit_wave.apply_ctf(zernike_kernel).intensity().poisson_noise(electrons_per_area).array
scaled_zernike, _, _ = histogram_ordering(noisy_zernike,normalize=True)

im_zernike = axs[2].imshow(scaled_zernike,cmap='gray')
axs[2].set_title("Zernike intensity",fontsize=12)
axs[2].axis("off")

fig.tight_layout()

## inset
inset_array = Complex2RGB(np.fft.fftshift(zernike_kernel._evaluate_from_angular_grid(alpha,phi)))
ax_inset = inset_axes(axs[2],width="30%",height="30%",loc="upper right",borderpad=0.25)
im_zernike_inset = ax_inset.imshow(inset_array)
ax_inset.axis("off")

if not show_zernike:
    axs[2].set_visible(False)
    ax_inset.set_visible(False)

fig.canvas.resizable = False
fig.canvas.header_visible = False
fig.canvas.footer_visible = False
fig.canvas.toolbar_visible = True
fig.canvas.layout.width = '680px'
fig.canvas.layout.height = "300px"
fig.canvas.toolbar_position = 'bottom'
None

def update_HRTEM(
    exit_wave,
    electrons_per_area,    
): 
    """ """
    noisy_hrtem = exit_wave.intensity().poisson_noise(electrons_per_area).array
    scaled_hrtem, _, _ = histogram_ordering(noisy_hrtem,normalize=True)
    
    im_hrtem.set_data(scaled_hrtem)
    fig.canvas.draw_idle()
    return None

def update_Zernike(
    exit_wave,
    electrons_per_area,
    zernike_radius,
    zernike_phase_shift,
): 
    """ """
    zernike_kernel = abtem.transfer.Zernike(
        semiangle_cutoff = semiangle,
        center_hole_cutoff = zernike_radius,
        phase_shift = zernike_phase_shift,
        energy=energy,
    ).match_grid(
        potential
    )
    noisy_zernike = exit_wave.apply_ctf(zernike_kernel).intensity().poisson_noise(electrons_per_area).array
    scaled_zernike, _, _ = histogram_ordering(noisy_zernike,normalize=True)
    inset_array = Complex2RGB(np.fft.fftshift(zernike_kernel._evaluate_from_angular_grid(alpha,phi)))
    
    im_zernike.set_data(scaled_zernike)
    im_zernike_inset.set_data(inset_array)
    
    fig.canvas.draw_idle()
    return None

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

layout_top = ipywidgets.Layout(width="335px",height="30px")
layout_bottom = ipywidgets.Layout(width="225px",height="30px")

defocus_slider = ipywidgets.FloatSlider(
    value = 0, 
    min = -2, 
    max = 2, 
    step = 0.05,
    description = r"defocus [$\mu$m]",
    style = style,
    layout = layout_top,
)

def change_defocus(change):
    defocus = change['new']
    array_to_mutate[0] = exit_wave_raw.apply_ctf(defocus=defocus*1e4)
    
    update_HRTEM(
        array_to_mutate[0],
        electrons_per_area_slider.value
    )
    update_Zernike(
        array_to_mutate[0],
        electrons_per_area_slider.value,
        zernike_radius_slider.value,
        zernike_phase_shift_slider.value
    )
    return None
    
defocus_slider.observe(change_defocus,names='value')

electrons_per_area_slider = ipywidgets.FloatLogSlider(
    value=10,
    base=10,
    min=1, # min exponent of base
    max=3, # max exponent of base
    step=0.05, # exponent step
    description = r"dose [e/A$^2$]",
    style = style,
    layout = layout_top,
)

def change_dose(change):
    dose = change['new']
    update_HRTEM(
        array_to_mutate[0],
        dose
    )
    update_Zernike(
        array_to_mutate[0],
        dose,
        zernike_radius_slider.value,
        zernike_phase_shift_slider.value
    )
    return None
    
electrons_per_area_slider.observe(change_dose,names='value')

zernike_radius_slider = ipywidgets.FloatSlider(
    value = 0.005,
    min = 0.005,
    max = semiangle, 
    step = 0.005,
    description = "radius [mrad]",
    style = style,
    layout = ipywidgets.Layout(width="250px",height="30px"),
)

def change_radius(change):
    radius = change['new']
    update_Zernike(
        array_to_mutate[0],
        electrons_per_area_slider.value,
        radius,
        zernike_phase_shift_slider.value
    )
    return None
    
zernike_radius_slider.observe(change_radius,names='value')

zernike_phase_shift_slider = ipywidgets.FloatSlider(
    value = -np.pi/2,
    min = -np.pi,
    max = 0, 
    step = np.pi/128,
    description = "phase shift [rad]",
    style = style,
    layout = ipywidgets.Layout(width="250px",height="30px"),
)

def change_phase_shift(change):
    phase_shift = change['new']
    update_Zernike(
        array_to_mutate[0],
        electrons_per_area_slider.value,
        zernike_radius_slider.value,
        phase_shift
    )
    return None
    
zernike_phase_shift_slider.observe(change_phase_shift,names='value')

show_zernike_switch = ipywidgets.ToggleButton(
    value=False,
    description="use Zernike phase plate",
    indent=False,
    style=style,
    layout = ipywidgets.Layout(width="175px",height="30px"),
)

def toggle_zernike(change):
    show_zernike = change['new']
    axs[2].set_visible(show_zernike)
    ax_inset.set_visible(show_zernike)
    fig.canvas.draw_idle()
    return None

show_zernike_switch.observe(toggle_zernike,names='value')

display(
    ipywidgets.VBox([
        ipywidgets.HBox(
            [
                defocus_slider,
                electrons_per_area_slider,
            ],
            layout=ipywidgets.Layout(justify_content="center",width="680px")
        ),
        ipywidgets.HBox(
            [
                show_zernike_switch,
                zernike_radius_slider,
                zernike_phase_shift_slider,
            ],
            layout=ipywidgets.Layout(justify_content="center",width="680px")
        ),
        fig.canvas,
    ]),
)