%matplotlib widget
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar

from IPython.display import display
import ipywidgets
import py4DSTEM

style = {'description_width': 'initial'}
layout = ipywidgets.Layout(width="340px",height="30px")

defocus_slider = ipywidgets.FloatSlider(
    value=150,
    min=-150,
    max=150,
    step=5,
    description = "defocus [Å]",
    style=style,
    layout=layout,
    # continuous_update = False,
)

convergence_slider = ipywidgets.FloatSlider(
    value=25,
    min=5,
    max=35,
    step=2.5,
    description = "semiangle [mrad]",
    style=style,
    layout=layout,
    # continuous_update = False,
)

energy_slider = ipywidgets.FloatSlider(
    value=80,
    min=20,
    max=300,
    step=20,
    description = "energy [kV]",
    style=style,
    layout=layout,
    # continuous_update = False,
)

multislice_checkbox = ipywidgets.Checkbox(
    value=False,
    description = "use multislice propagator?",
    indent=False,
    style=style,
    layout=layout,
)

# inputs
gpts = (244,242)
sampling = (0.1,0.1)
q_sampling = 1/gpts[1]/sampling[1]
dz = 20/7
dp_power = 0.25
dpi = 72

# arrays
potential = np.fromfile(
    "data/FCC-slab-potential-7x244x242-float32.npy", dtype=np.float32
).reshape((-1,) + gpts)
n_slices = potential.shape[0]

probe = (
    py4DSTEM.process.phase.utils.ComplexProbe(
        energy=energy_slider.value*1000,
        sampling=sampling,
        gpts=gpts,
        semiangle_cutoff=convergence_slider.value,
        defocus=defocus_slider.value,
    )
    .build()
    ._array
)

mutable_arrays = [
    potential, # potential
    np.exp(1j*potential), # cmplx potential
    potential.sum(0), # projected_potential
    np.exp(1j*potential.sum(0)), # cmplx projected potential
    probe, # probe
    np.fft.fftshift(probe), # shifted_probe
    np.array(gpts)/2 # xy pos
]

def update_probe_dp_panels(dummy=None):
    """ """
    dp = multislice_propagation() if multislice_checkbox.value else single_slice_propagation()
    scaled_probe = py4DSTEM.visualize.Complex2RGB(mutable_arrays[5],vmin=0,vmax=1)
    scaled_dp = py4DSTEM.visualize.return_scaled_histogram_ordering(
        np.fft.fftshift(dp),
        normalize=True,
        vmin=0,
        vmax=1
    )[0] 
    
    im_probe.set_data(scaled_probe)
    im_dp.set_data(scaled_dp)
    fig.canvas.draw_idle()
    return None

def update_probe(dummy):
    """ """
    mutable_arrays[4] = (
        py4DSTEM.process.phase.utils.ComplexProbe(
            energy=energy_slider.value*1000,
            sampling=sampling,
            gpts=gpts,
            semiangle_cutoff=convergence_slider.value,
            defocus=defocus_slider.value,
        )
        .build()
        ._array
    )
    mutable_arrays[5] = py4DSTEM.process.phase.utils.fft_shift(mutable_arrays[4],mutable_arrays[6])
    update_probe_dp_panels()
    return None

defocus_slider.observe(update_probe,names='value')
convergence_slider.observe(update_probe,names='value')
energy_slider.observe(update_probe,names='value')

def update_energy(change):
    """ """
    old = change['old']
    new = change['new']
    scaling_factor = py4DSTEM.process.utils.electron_interaction_parameter(new * 1e3) / py4DSTEM.process.utils.electron_interaction_parameter(old * 1e3)

    mutable_arrays[0] *= scaling_factor
    
    mutable_arrays[1] = np.exp(1j*mutable_arrays[0])
    mutable_arrays[2] = mutable_arrays[0].sum(0)
    mutable_arrays[3] = np.exp(1j*mutable_arrays[2])

    scaled_pot = py4DSTEM.visualize.return_scaled_histogram_ordering(mutable_arrays[2],normalize=True)[0]
    im_pot.set_data(scaled_pot)
    return None

energy_slider.observe(update_energy,names='value')
multislice_checkbox.observe(update_probe_dp_panels,names='value')

def return_propagator_array(gpts,sampling, energy, dz):
    """ """
    prefactor = py4DSTEM.process.utils.electron_wavelength_angstrom(energy) * np.pi * dz

    kx = np.fft.fftfreq(gpts[0],sampling[0])
    ky = np.fft.fftfreq(gpts[1],sampling[1])
    KX, KY = np.meshgrid(kx,ky,indexing='ij')

    chi = (KX**2 + KY**2) * prefactor
    propagator = np.exp(1j*chi)
    return propagator

def propagate_wavefunction(array,prop_array):
    """ """
    array_fourier = np.fft.fft2(array)
    return np.fft.ifft2(array_fourier * prop_array)

def multislice_propagation():
    """ """
    wavefunction = mutable_arrays[5].copy()
    for s, cmplx_pot in enumerate(mutable_arrays[1]):
        wavefunction *= cmplx_pot
        if s + 1 < n_slices:
            wavefunction = propagate_wavefunction(wavefunction,propagator_array)
    dp = np.abs(np.fft.fft2(wavefunction))**(2*dp_power)
    return dp
    
def single_slice_propagation():
    """ """
    wavefunction = mutable_arrays[3] * mutable_arrays[5]
    dp = np.abs(np.fft.fft2(wavefunction))**(2*dp_power)
    return dp

# arrays to visualize
propagator_array = return_propagator_array(gpts,sampling,energy_slider.value*1000,dz)

scaled_pot = py4DSTEM.visualize.return_scaled_histogram_ordering(mutable_arrays[2],normalize=True)[0]
scaled_probe = py4DSTEM.visualize.Complex2RGB(mutable_arrays[5],vmin=0,vmax=1)
scaled_dp = py4DSTEM.visualize.return_scaled_histogram_ordering(
    np.fft.fftshift(
        single_slice_propagation()
    ),
    normalize=True,
    vmin=0,
    vmax=1
)[0]

def add_scalebar(ax, length, sampling, units, color="white", size_vertical=1, pad=0.5):
    """ """
    bar = AnchoredSizeBar(
        ax.transData,
        length,
        f"{sampling*length:.2f} {units}",
        "lower right",
        pad=pad,
        color=color,
        frameon=False,
        label_top=True,
        size_vertical=size_vertical,
    )
    ax.add_artist(bar)
    return ax, bar

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

im_pot = axs[0].imshow(scaled_pot,cmap='magma')
im_probe = axs[1].imshow(scaled_probe)
im_dp = axs[2].imshow(scaled_dp,cmap='magma')

titles = [
    "projected sample potential",
    "converged electron probe",
    "diffracted probe intensity",
]

scalebars = [
    {'sampling':sampling[1],'length':50,'units':'Å'},
    {'sampling':sampling[1],'length':50,'units':'Å'},
    {'sampling':q_sampling,'length':48.4,'units':r'Å$^{-1}$'},
]

for ax, title, bar in zip(
    axs,
    titles,
    scalebars
):
    add_scalebar(ax,**bar)
    ax.set_title(title)
    ax.set_xticks([])
    ax.set_yticks([])

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'

def onmove(event):
    """ """
    pos = np.array([event.ydata,event.xdata])
    
    if pos[0] is not None:
        mutable_arrays[6] = pos
        mutable_arrays[5] = py4DSTEM.process.phase.utils.fft_shift(mutable_arrays[4],mutable_arrays[6])
        update_probe_dp_panels()

cid = fig.canvas.mpl_connect('motion_notify_event',onmove)

ipywidgets.VBox(
    [
        ipywidgets.HBox(
            [
                energy_slider,
                multislice_checkbox
            ]
        ),
        ipywidgets.HBox(
            [
                defocus_slider,
                convergence_slider
            ]
        ),
        fig.canvas,
    ]
)