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

import numpy as np
import matplotlib.pyplot as plt

from IPython.display import display
import temgymlite
import ipywidgets

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

file_name = "apoF-ice-embedded-potential-binned.npy"
binned_volume_zxy = np.load("data/"+file_name)
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 # um
electrons_per_area = 10 # e/Ang^2
tilt_px = (nx/2,ny/2)
tilt_y = (nx/2-tilt_px[0])*0.077 # mrad
tilt_x = (ny/2-tilt_px[1])*0.077 # mrad

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

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

exit_wave = abtem.PlaneWave(
    energy=energy,
    tilt=(tilt_x,tilt_y),
).match_grid(
    potential
).multislice(
    potential,
    lazy=False
)

angular_sampling = exit_wave.angular_sampling[0]

aperture = abtem.Aperture(
    semiangle_cutoff=semiangle,
    energy=energy
).match_grid(
    potential
)

alpha, phi = aperture._angular_grid(device='cpu')
bright_field_disk = np.fft.fftshift(aperture._evaluate_from_angular_grid(alpha,phi))
kmax = exit_wave.full_cutoff_angles[0]

arrays_to_mutate = [exit_wave,tilt_px]

components = [
    temgymlite.components.Lens(name = ' ', z = 0.8, f = -0.13, radius=0.25),
    temgymlite.components.Sample(name = 'Sample', z = 0.43)
]

axis_view = 'x_axial'
model = temgymlite.Model(
    components,
    beam_z=1, 
    beam_type=axis_view,
    num_rays=11, 
    gun_beam_semi_angle=0.65,
)

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


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

# ray diagram

temgymlite.show_matplotlib(
    model,
    figax=(fig,axs[0]),
    label_fontsize=12,
    plot_rays=True,
    ray_color="#01F100",
    fill_between=False,
    highlight_edges=False,
    show_labels=False,
    ray_lw=2,
)
axs[0].set_title("ray diagram",fontsize=12)
sample_line = axs[0].lines[22]

# detector pixel
axs[1].imshow(
    bright_field_disk,
    cmap='gray'
)
axs[1].set(xticks=[],yticks=[])
add_scalebar(
    axs[1],
    scalebar_dict_bf
)
scatter = axs[1].scatter(tilt_px[0],tilt_px[1],color="red")
axs[1].set_title("bright field disk pixel",fontsize=12)

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

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


fig.tight_layout()
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 = "275px"
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_BF_disk(
    tilt_px,
): 
    """ """
    scatter.set_offsets(tilt_px)
    fig.canvas.draw_idle()
    return None

def update_ray_diagram(
    defocus,
):
    """ """
    sample_line.set_ydata([0.43 + defocus * 0.15/2]*2)
    fig.canvas.draw_idle()
    return None

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

layout_top = ipywidgets.Layout(width="335px",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']
    thicknesses = np.tile(pixel_size * bin_factor_z,num_slices)
    thicknesses += defocus*1e4 / num_slices
    
    potential = abtem.PotentialArray(
        binned_volume_zxy,
        slice_thickness = thicknesses,
        sampling= sampling,
    )

    tilt_y = (nx/2-arrays_to_mutate[1][0])*angular_sampling
    tilt_x = (ny/2-arrays_to_mutate[1][1])*angular_sampling
    
    exit_wave = abtem.PlaneWave(
        energy=energy,
        tilt=(tilt_x,tilt_y),
    ).match_grid(
        potential
    ).multislice(
        potential,
        lazy=False
    )
    arrays_to_mutate[0] = exit_wave
    
    update_HRTEM(
        arrays_to_mutate[0],
        electrons_per_area_slider.value
    )
    update_ray_diagram(
        defocus
    )
    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(
        arrays_to_mutate[0],
        dose
    )
    return None
    
electrons_per_area_slider.observe(change_dose,names='value')

def onclick(event):
    """ """
    pos = np.array([event.xdata,event.ydata])
    
    if pos[0] is not None:
        arrays_to_mutate[1] = pos
        update_BF_disk(arrays_to_mutate[1])
        change_defocus({'new':defocus_slider.value})

cid = fig.canvas.mpl_connect('button_press_event',onclick)

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