%matplotlib widget

import tcbf
import numpy as np
import matplotlib.pyplot as plt

from IPython.display import display
import ipywidgets


from temgymlite import components as comp
from temgymlite.model import Model
from temgymlite.run import show_matplotlib

file_name = "apoF-ice-embedded-potential-binned.npy"
binned_volume_zxy = np.load("data/"+file_name)

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

layout = ipywidgets.Layout(width="300px",height="30px")

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


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,
)

kx_slider = ipywidgets.FloatSlider(
    value = 0, min = -4, max = 4, 
    step = 0.125,
    description = r"k$_x$ [mrad]",
    style = style,
    #layout = layout,
)

ky_slider = ipywidgets.FloatSlider(
    value = 0, min = -4, max = 4, 
    step = 0.125,
    description = r"k$_y$ [mrad]",
    style = style,
    #layout = layout,
)

# constants
semiangle = 4  # mrad
wavelength = 0.0197  # A (300kV)
sigma = 0.00065  # 1/V (300kV)
rolloff = 0.125  # mrad

# PotentialArray

pixel_size = 2 / 3
bin_factor_xy = 2
bin_factor_z = 6

potential = tcbf.PotentialArray(
    binned_volume_zxy,
    slice_thickness=pixel_size * bin_factor_z,
    sampling=(pixel_size * bin_factor_xy, pixel_size * bin_factor_xy),
)

potential.slice_thickness = pixel_size * bin_factor_z + 1e4 * defocus_slider.value / binned_volume_zxy.shape[0]

# Tilted Plane Wave
tilted_plane_wave = tcbf.Waves(
    array=np.ones(potential.gpts, dtype=np.complex64),
    sampling=potential.sampling,
    wavelength=wavelength,
    sigma=sigma,
    tilt=(0, 0),
)

# CTF
ctf = tcbf.CTF(
    semiangle_cutoff=semiangle,
    rolloff=rolloff,
)

# Angles
alpha, phi = tilted_plane_wave.get_scattering_angles()
bright_field_disk = np.fft.fftshift(ctf.evaluate_aperture(alpha, phi))
kmax = alpha.max() * 1e3 / np.sqrt(2)

# Exit Waves

exit_wave = tilted_plane_wave.multislice(potential)
exit_wave = np.random.poisson(
    (
        np.abs(exit_wave) ** 2
        * np.prod(potential.sampling)
        * electrons_per_area_slider.value
    ).clip(0)
)

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

axis_view = 'x_axial'
model_ = Model(
    components,
    beam_z=1, 
    beam_type=axis_view,
    num_rays=11, 
    gun_beam_semi_angle=0.65,
#
)

# Static Figure

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

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_xticks([])

axs[0].set_title(
    "BF STEM ray diagram",
    fontsize=14,
)

sample_line = axs[0].lines[22]

tcbf.show(
    bright_field_disk,
    extent=[-kmax, kmax, kmax, -kmax],
    figax=(fig, axs[1]),
    ticks=False,
    cbar=False,
)

scatter = axs[1].scatter(0, 0, color="red")
axs[1].set_title("STEM bright-field disk", fontsize=14)

tcbf.add_scalebar(
    axs[1],
    color="white",
    sampling=1,
    length=3,
    units="mrad",
    size_vertical=0.25,
    pad=0.5,
)

tcbf.show(
    exit_wave,
    ticks=False,
    figax=(fig, axs[2]),
    cbar=False,
    cmap="gray",
)

tcbf.add_scalebar(
    axs[2],
    color="white",
    sampling=pixel_size * bin_factor_xy / 10,
    length=30,
    units="nm",
)

axs[2].set_title(
    "virtual bright-field image",
    fontsize=14,
)

im = axs[2].get_images()[0]
artists = [scatter,im]

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 = '275px'
fig.canvas.toolbar_position = 'bottom'
fig.tight_layout()

def update_figure(defocus, electrons_per_area, tilt_kx, tilt_ky):
    """ """
    potential.slice_thickness = pixel_size * bin_factor_z + 1e4 * defocus / binned_volume_zxy.shape[0]
    tilted_plane_wave = tcbf.Waves(
        array=np.ones(potential.gpts, dtype=np.complex64),
        sampling=potential.sampling,
        wavelength=wavelength,
        sigma=sigma,
        tilt=(tilt_kx, -tilt_ky),
    )
    exit_wave = tilted_plane_wave.multislice(potential)
    exit_wave = np.random.poisson(
        (
            np.abs(exit_wave) ** 2
            * np.prod(potential.sampling)
            * electrons_per_area
        ).clip(0)
    )
    
    exit_wave, vmin, vmax = tcbf.visualize.return_scaled_histogram(exit_wave)

    k = np.sqrt(tilt_kx**2 + tilt_ky**2)
    artists[1].set_data(exit_wave)
    artists[1].set_clim(vmin=vmin, vmax=vmax)
    artists[1].set_cmap("gray" if k <= 4 else "gray_r")

    artists[0].set_offsets([tilt_ky, -tilt_kx])

    sample_line.set_ydata([0.43 + defocus * 0.15/2]*2)
    fig.canvas.draw_idle()
    return None

ipywidgets.widgets.interactive_output(
    update_figure, 
    {
        'defocus':defocus_slider,
        'electrons_per_area':electrons_per_area_slider,
        'tilt_kx':kx_slider,
        'tilt_ky':ky_slider,
    },
)
None

display(
    ipywidgets.VBox([
        fig.canvas,
        ipywidgets.TwoByTwoLayout(
            top_left=defocus_slider,
            top_right=electrons_per_area_slider,
            bottom_left=kx_slider,
            bottom_right=ky_slider,
            layout=ipywidgets.Layout(justify_content="center",width="680px")
        )
    ])
)