%matplotlib widget
import py4DSTEM

import numpy as np
import matplotlib.pyplot as plt

from IPython.display import display
import ipywidgets

file_path = 'data/'
file_data = file_path + 'parallax_apoferritin_simulation_100eA2.h5'

dataset = py4DSTEM.read(file_data)
dataset

from py4DSTEM.process.phase.utils import polar_aliases

energy = 300e3
object_padding_px = (8,8)
edge_blend = 4

parallax = py4DSTEM.process.phase.Parallax(
    datacube=dataset,
    energy = energy,
    object_padding_px=object_padding_px,
).preprocess(
    edge_blend=edge_blend,
    plot_average_bf=False,
)

defocus = 1.5e4
rotation_angle_deg = -15
plot_arrow_freq = 4

style = {'description_width': 'initial'}
output = ipywidgets.Output(
    layout=ipywidgets.Layout(width="675px",height="375px")
)
vbox_layout = ipywidgets.Layout(width='325px')
arrow_freq_slider = ipywidgets.IntSlider(min=1, max=16, step=1, value=4,style=style,description="plot arrow frequency",layout=vbox_layout,continuous_update=False)
upsampling_slider = ipywidgets.IntSlider(min=1, max=4, step=1, value=1,style=style,description="upsampling factor",layout=vbox_layout,continuous_update=False)
rotation_slider = ipywidgets.IntSlider(min=-90, max=90, step=1, value=-15, style = style, description="rotation angle [°]",layout=vbox_layout,continuous_update=False)
defocus_slider = ipywidgets.IntSlider(min=-2e4, max=2e4, step=5e2, value=1.5e4, style = style, description="defocus [Å]",layout=vbox_layout,continuous_update=False)
spherical_slider = ipywidgets.IntSlider(min=-2e9, max=2e9, step=5e5, value=0, style = style, description="spherical aberration [Å]",layout=vbox_layout,continuous_update=False)
astigmatism_slider = ipywidgets.IntSlider(min=0, max=1e4, step=5e2, value=0, style = style, description="astigmatism [Å]",layout=vbox_layout,continuous_update=False)
astigmatism_angle_slider = ipywidgets.IntSlider(min=-90, max=90, step=1, value=45, style = style, description="astigmatism angle [°]",layout=vbox_layout,continuous_update=False)
coma_slider = ipywidgets.IntSlider(min=0, max=1e7, step=5e3, value=0, style = style, description="coma [Å]",layout=vbox_layout,continuous_update=False)
coma_angle_slider = ipywidgets.IntSlider(min=-180, max=180, step=1, value=90, style = style, description="coma angle [°]",layout=vbox_layout,continuous_update=False)

def reset_aberrations(b):
    defocus_slider.value = 1.5e4
    spherical_slider.value = 0
    astigmatism_slider.value = 0
    coma_slider.value = 0
    astigmatism_angle_slider.value = 45
    coma_angle_slider.value = 90
    return None

reset_aberrations_button = ipywidgets.Button(description="reset aberrations")
reset_aberrations_button.on_click(reset_aberrations)

def reset_angle(b):
    rotation_slider.value = -15
    return None

reset_angle_button = ipywidgets.Button(description="reset angle")
reset_angle_button.on_click(reset_angle)

def widget_wrapper(
    plot_arrow_freq,
    upsampling_factor,
    rotation_angle_deg,
    defocus,
    spherical_aberration,
    astigmatism,
    astigmatism_angle_deg,
    coma,
    coma_angle_deg,
):
    """ """
    output.clear_output(wait=True)
    with output:
        parallax.guess_common_aberrations(
            defocus=defocus,
            Cs=spherical_aberration,
            astigmatism=astigmatism,
            astigmatism_angle=np.deg2rad(astigmatism_angle_deg),
            coma=coma,
            coma_angle=np.deg2rad(coma_angle_deg),
            rotation_angle_deg=rotation_angle_deg,
            kde_upsample_factor=upsampling_factor,
            plot_arrow_freq=plot_arrow_freq,
            figsize=(675/100,350/100),
        )
        fig = plt.gcf()
        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 = '375px'
        fig.canvas.toolbar_position = 'bottom'
        fig.tight_layout()
        plt.show()
    return None

ipywidgets.interactive_output(
    widget_wrapper,
    {
        'plot_arrow_freq': arrow_freq_slider,
        'upsampling_factor':upsampling_slider,
        'rotation_angle_deg': rotation_slider,
        'defocus': defocus_slider,
        'spherical_aberration': spherical_slider,
        'astigmatism': astigmatism_slider,
        'astigmatism_angle_deg': astigmatism_angle_slider,
        'coma': coma_slider,
        'coma_angle_deg': coma_angle_slider,
    }
)
None

display(
    ipywidgets.VBox(
        [
            ipywidgets.HBox(
                [
                    ipywidgets.VBox(
                        [
                            ipywidgets.HBox([ipywidgets.HTML("<b>Microscope Geometry:</b>",layout=ipywidgets.Layout(width="175px")),reset_angle_button]),
                            ipywidgets.HTML("<hr>"),
                            rotation_slider,
                            ipywidgets.HTML("<b>Plotting Options:</b>"),
                            ipywidgets.HTML("<hr>"),
                            arrow_freq_slider,
                            upsampling_slider,
                        ]
                    ),
                    ipywidgets.VBox(
                        [
                            ipywidgets.HBox([ipywidgets.HTML("<b>Aberrations:</b>",layout=ipywidgets.Layout(width="175px")),reset_aberrations_button]),
                            ipywidgets.HTML("<hr>"),
                            defocus_slider,
                            spherical_slider,
                            astigmatism_slider,
                            astigmatism_angle_slider,
                            coma_slider,
                            coma_angle_slider,
                        ],
                    )
                ],
                layout=ipywidgets.Layout(justify_content='center')
            ),
            output,
        ],
        layout=ipywidgets.Layout(align_content='center', width='675px')
    )
)