%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_in_focus = file_path + 'dpc_STO_simulation_in-focus_1e5.h5'
dataset_in_focus = py4DSTEM.read(file_data_in_focus)

dataset_real_space_in_focus = np.fft.ifftshift(dataset_in_focus.data,axes=(-1,-2))

energy = 200e3
semiangle = 20
wavelength = abtem.core.energy.energy2wavelength(energy)

scan_gpts = dataset_in_focus.Rshape
gpts = dataset_in_focus.Qshape

scan_sampling = (dataset_in_focus.calibration.R_pixel_size,)*2
angular_sampling = (dataset_in_focus.calibration.Q_pixel_size,)*2
sampling = (wavelength*1e3/angular_sampling[0]/gpts[0],)*2

kx, ky = abtem.core.grid.spatial_frequencies(gpts, sampling)
kxa, kya = np.meshgrid(kx,ky,indexing='ij')

qx, qy = abtem.core.grid.spatial_frequencies(scan_gpts, scan_sampling)
qxa, qya = np.meshgrid(qx,qy,indexing='ij')

q2 = qxa**2 + qya**2
q = np.sqrt(q2)
q2[0, 0] = np.inf

# iCoM operators
qx_op = -1.0j * qxa / q2
qy_op = -1.0j * qya / q2

def compute_com(
    dp,
    kxa,
    kya,
    rotation_angle_deg,
):
    """ """
    theta = np.deg2rad(rotation_angle_deg)
    ct = np.cos(theta)
    st = np.sin(theta)
    
    dp_sum = dp.sum()
    com_x = np.sum(dp*kxa) / dp_sum
    com_y = np.sum(dp*kya) / dp_sum

    com_x, com_y = (ct*com_x - st*com_y,st*com_x+ct*com_y)
    return com_x, com_y

def integrate_com(
    com_x,
    com_y,
    kx_op,
    ky_op,
):
    """ """

    icom_fft = np.fft.fft2(com_x)*kx_op + np.fft.fft2(com_y)*ky_op
    return np.real(np.fft.ifft2(icom_fft))

x, y = [np.arange(n) for n in scan_gpts]
positions = np.dstack(np.meshgrid(x,y,indexing='ij')).reshape(-1,2)

index_i, index_j = positions[0]
com_x_array = np.zeros(scan_gpts)
com_y_array = np.zeros(scan_gpts)
cmplx_com_array = np.zeros(scan_gpts,dtype=np.complex128)

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

# raw data
dp = dataset_real_space_in_focus[index_i,index_j].astype(np.float64)
scaled_dp, _, _ = py4DSTEM.visualize.return_scaled_histogram_ordering(np.fft.fftshift(dp),normalize=True)
im_dp = axs[0].imshow(scaled_dp,cmap='gray')

com_x, com_y = compute_com(
    dp,
    kxa,
    kya,
    rotation_angle_deg=-15,
)

com_x_array[index_i,index_j] = com_x
com_y_array[index_i,index_j] = com_y
cmplx_com_array[index_i,index_j] = com_x + 1j*com_y

im_com_x = axs[1].imshow(np.tile(com_x_array,(2,2)),cmap='RdBu',vmin=-0.125,vmax=0.125)
im_com_y = axs[2].imshow(np.tile(com_y_array,(2,2)),cmap='RdBu',vmin=-0.125,vmax=0.125)

scaled_cmplx = py4DSTEM.visualize.Complex2RGB(cmplx_com_array,vmin=0,vmax=1)
im_cmplx = axs[3].imshow(np.tile(scaled_cmplx,(2,2,1)),cmap='magma',vmin=0,vmax=1)

titles = [
    "diffraction pattern", "center of mass x", "center of mass y", "complex center of mass"
]

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

bars = [
    scalebar_fourier,scalebar_real,scalebar_real,scalebar_real
]

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

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

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

def reset_arrays(*args):
    """ """
    com_x_array[:,:] = 0
    com_y_array[:,:] = 0
    cmplx_com_array[:,:] = 0j
    
    slider.value=0
    update_plots({'new':slider.value})
    return None

def toggle_icom_button(
    change
):
    """ """
    icom = change["new"]
    if icom:
        axs[3].set_title("integrated center of mass")
    else:
        axs[3].set_title("complex center of mass")
    update_plots({'new':slider.value})
    return None

def update_plots(
    change
):
    """ """
    sorted_index = change["new"]
    index_i, index_j = positions[sorted_index]
    
    dp = dataset_real_space_in_focus[index_i,index_j].astype(np.float64)
    scaled_dp, _, _ = py4DSTEM.visualize.return_scaled_histogram_ordering(np.fft.fftshift(dp),normalize=True)
    im_dp.set_data(scaled_dp)
    
    com_x, com_y = compute_com(
        dp,
        kxa,
        kya,
        rotation_angle_deg=-15,
    )
    
    com_x_array[index_i,index_j] = com_x
    com_y_array[index_i,index_j] = com_y
    cmplx_com_array[index_i,index_j] = com_x + 1j*com_y
    
    im_com_x.set_data(np.tile(com_x_array,(2,2)))
    im_com_y.set_data(np.tile(com_y_array,(2,2)))

    if icom_toggle.value:
        icom = integrate_com(
            com_x_array,
            com_y_array,
            qx_op,
            qy_op
        )
        scaled_cmplx, _, _ = py4DSTEM.visualize.return_scaled_histogram_ordering(icom,normalize=True)
        scaled_cmplx = np.tile(scaled_cmplx,(2,2))
    else:
        scaled_cmplx = py4DSTEM.visualize.Complex2RGB(cmplx_com_array,vmin=0,vmax=1)
        scaled_cmplx = np.tile(scaled_cmplx,(2,2,1))
        
    im_cmplx.set_data(scaled_cmplx)

    fig.canvas.draw_idle()
    return None

play = ipywidgets.Play(
    value=0,
    min=0,
    max=len(positions)-1,
    step=1,
    interval=50,
    show_repeat=False,
    style=style,
    layout=layout,
)

slider = ipywidgets.IntSlider(
    min=0,
    max=len(positions)-1,
    step=1,
    layout=layout,
    style=style,
    description="scan position index"
)

icom_toggle = ipywidgets.ToggleButton(
    value=False,
    layout=layout,
    style=style,
    description="integrated CoM"
)

reset_button = ipywidgets.Button(
    description="reset reconstructions",
    style=style,
    layout=layout,
)

ipywidgets.jslink((play, 'value'), (slider, 'value'))
slider.observe(update_plots,"value")
icom_toggle.observe(toggle_icom_button,"value")
reset_button.on_click(reset_arrays)

ipywidgets.VBox(
    [
        ipywidgets.HBox([play,icom_toggle,reset_button]),
        fig.canvas
    ],
    layout=ipywidgets.Layout(
        align_items="center"
    )
)