%matplotlib widget
import abtem
import ase
import py4DSTEM

import numpy as np
import matplotlib.pyplot as plt

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

file_path = 'data/'
file_data_aberrated = file_path + 'dpc_STO_simulation_aberrated_1e5.h5'
dataset_aberrated = py4DSTEM.read(file_data_aberrated)

energy = 200e3
semiangle_cutoff = 20
rotation = -15
gpts = dataset_aberrated.Qshape
scan_gpts = dataset_aberrated.Rshape

ssb = py4DSTEM.process.phase.SSB(
    energy=energy,
    datacube=dataset_aberrated,
    semiangle_cutoff=semiangle_cutoff,
    verbose=False,
).preprocess(
    plot_center_of_mass=False,
    plot_rotation=False,
    plot_overlap_trotters=False,
    force_com_rotation=rotation,
    force_com_transpose=False,
    vectorized_com_calculation=False,
).reconstruct(
    progress_bar=False
)

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

angle_data, _, _ = py4DSTEM.visualize.return_scaled_histogram_ordering(np.angle(ssb.object),normalize=True)
ampl_data, _, _ = py4DSTEM.visualize.return_scaled_histogram_ordering(np.abs(ssb.object),normalize=True)
im_angle = axs[0].imshow(np.tile(np.fft.fftshift(angle_data),(2,2)),cmap='magma')
im_ampl = axs[1].imshow(np.tile(np.fft.fftshift(ampl_data),(2,2)),cmap='gray')

titles = [
    "reconstructed phase", "reconstructed amplitude"
]

scalebar_real = {'pixelsize':ssb.scan_sampling[1],'pixelunits':r'$\AA$',"Nx":scan_gpts[0]*2,"Ny":scan_gpts[1]*2,"labelsize":10}

for ax, title in zip(axs.flatten(),titles):
    ax.set(xticks=[],yticks=[],title=title)
    py4DSTEM.visualize.add_scalebar(ax,scalebar_real)

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

def update_plots(
    defocus,
    stig,
    stig_angle,
    coma,
    coma_angle,
    trefoil,
    trefoil_angle
):
    """ """
    ssb.reconstruct(
        polar_parameters={'C10':-defocus,'C12':stig,'phi12':np.deg2rad(stig_angle),'C21':coma,'phi21':np.deg2rad(coma_angle),'C23':trefoil,'phi23':np.deg2rad(trefoil_angle)},
        progress_bar=False,
    )

    angle_data, _, _ = py4DSTEM.visualize.return_scaled_histogram_ordering(np.angle(ssb.object),normalize=True)
    ampl_data, _, _ = py4DSTEM.visualize.return_scaled_histogram_ordering(np.abs(ssb.object),normalize=True)
    im_angle.set_data(np.tile(np.fft.fftshift(angle_data),(2,2)))
    im_ampl.set_data(np.tile(np.fft.fftshift(ampl_data),(2,2)))
    fig.canvas.draw_idle()
    return None


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

defocus_slider = ipywidgets.FloatSlider(value = 0,min = -100,max = 100, description = "defocus [Å]",**kwargs)
stig_slider = ipywidgets.FloatSlider(value = 0,min = 0,max = 100, description = "stig [Å]",**kwargs)
stig_angle_slider = ipywidgets.FloatSlider(value = 0,min = 0,max = 90, description = "stig angle [°]",**kwargs)
coma_slider = ipywidgets.FloatSlider(value = 0,min = 0,max = 10000, description = "coma [Å]",**kwargs)
coma_angle_slider = ipywidgets.FloatSlider(value = 0,min = 0,max = 180, description = "coma angle [°]",**kwargs)
trefoil_slider = ipywidgets.FloatSlider(value = 0,min = 0,max = 10000, description = "trefoil [Å]",**kwargs)
trefoil_angle_slider = ipywidgets.FloatSlider(value = 0,min = 0,max = 60, description = "trefoil angle [°]",**kwargs)

def set_solution(*args):
    """ """
    (
        defocus_slider.value,
        stig_slider.value,
        stig_angle_slider.value,
        coma_slider.value,
        coma_angle_slider.value,
        trefoil_slider.value,
        trefoil_angle_slider.value
    ) = (100,0,0,0,0,10000,27.5)
    
    return None

solutions_button = ipywidgets.Button(description='reveal solution!',**kwargs)
solutions_button.on_click(set_solution)

ipywidgets.interactive_output(
    update_plots,
    {
        'defocus':defocus_slider,
        'stig':stig_slider,
        'stig_angle':stig_angle_slider,
        'coma':coma_slider,
        'coma_angle':coma_angle_slider,
        'trefoil':trefoil_slider,
        'trefoil_angle':trefoil_angle_slider,
    }
);

ipywidgets.VBox(
    [
        ipywidgets.HBox([defocus_slider,solutions_button]),
        ipywidgets.HBox([stig_slider,stig_angle_slider]),
        ipywidgets.HBox([coma_slider,coma_angle_slider]),
        ipywidgets.HBox([trefoil_slider,trefoil_angle_slider]),
        fig.canvas
    ]
)