%matplotlib widget
import numpy as np
import matplotlib.pyplot as plt
from colorspacious import cspace_convert
from IPython.display import display
from ipywidgets import HBox, VBox, widgets, Dropdown, Label, Layout, IntSlider, FloatSlider,FloatLogSlider
import abtem

# Copied from py4DSTEM directly
def Complex2RGB(complex_data, vmin=None, vmax=None, power=None, chroma_boost=1):
    """
    complex_data (array): complex array to plot
    vmin (float)        : minimum absolute value
    vmax (float)        : maximum absolute value
    power (float)       : power to raise amplitude to
    chroma_boost (float): boosts chroma for higher-contrast (~1-2.5)
    """
    amp = np.abs(complex_data)
    phase = np.angle(complex_data)

    if power is not None:
        amp = amp**power

    if np.isclose(np.max(amp), np.min(amp)):
        if vmin is None:
            vmin = 0
        if vmax is None:
            vmax = np.max(amp)
    else:
        if vmin is None:
            vmin = 0.02
        if vmax is None:
            vmax = 0.98
        vals = np.sort(amp[~np.isnan(amp)])
        ind_vmin = np.round((vals.shape[0] - 1) * vmin).astype("int")
        ind_vmax = np.round((vals.shape[0] - 1) * vmax).astype("int")
        ind_vmin = np.max([0, ind_vmin])
        ind_vmax = np.min([len(vals) - 1, ind_vmax])
        vmin = vals[ind_vmin]
        vmax = vals[ind_vmax]

    amp = np.where(amp < vmin, vmin, amp)
    amp = np.where(amp > vmax, vmax, amp)
    amp = ((amp - vmin) / vmax).clip(1e-16, 1)

    J = amp * 61.5  # Note we restrict luminance to the monotonic chroma cutoff
    C = np.minimum(chroma_boost * 98 * J / 123, 110)
    h = np.rad2deg(phase) + 180

    JCh = np.stack((J, C, h), axis=-1)
    rgb = cspace_convert(JCh, "JCh", "sRGB1").clip(0, 1)

    return rgb


# widget figure generation
with plt.ioff():
    dpi = 72
    fig = plt.figure(figsize=(675/dpi, 300/dpi), dpi=dpi)

b = (0.005,0.01)
ax0 = fig.add_axes([0/3+b[0],  0+b[1],  1/3-2*b[0], 1-2*b[1]])
ax1 = fig.add_axes([1/3+b[0],  0+b[1],  1/3-2*b[0], 1-2*b[1]])
ax2 = fig.add_axes([2/3+b[0],  0+b[1],  1/3-2*b[0], 1-2*b[1]])

sampling = 0.2
extent = 100

semiangle_cutoff = 10
energy = 60e3

ctf = abtem.transfer.CTF(
    energy=energy, 
    semiangle_cutoff = semiangle_cutoff,
    C30= 1e7,
)
probe = abtem.Probe(
    energy = energy,
    semiangle_cutoff=semiangle_cutoff,
    extent=(extent,extent),
    sampling=sampling,
).build().apply_ctf(ctf)

probe_Q = Complex2RGB(
    np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(probe.array.compute()))), 
    vmax = 1, vmin = 0
)
probe_R_vals = probe.array.compute()
probe_R = Complex2RGB(probe_R_vals, vmax = 1, vmin = 0)
probe_R_int = np.abs(probe_R_vals)**2
probe_R_int /= np.max(probe_R_int)

im0 = ax0.imshow(probe_Q)
im1 = ax1.imshow(probe_R)
im2 = ax2.imshow(
    probe_R_int,
    cmap = 'gray',
)

ax0.set_xticks([])  
ax0.set_yticks([]) 
ax1.set_xticks([])  
ax1.set_yticks([]) 
ax2.set_xticks([])  
ax2.set_yticks([]) 

ax0.set_title('Reciprocal Space')
ax1.set_title('Real Space Probe (complex)')
ax2.set_title('Real Space Probe (intensity)')

def update_probe(energy, defocus, C30, semiangle):

    ctf = abtem.transfer.CTF(
        energy=energy*1e3, 
        semiangle_cutoff = semiangle,
        defocus =  defocus,
        C30 =  C30*1e7,
    )

    probe = abtem.Probe(
        energy = energy*1e3,
        semiangle_cutoff = semiangle,
        extent=(extent,extent),
        sampling=sampling,
    ).build().apply_ctf(ctf)

    probe_Q = Complex2RGB(
        np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(probe.array.compute()))), 
        vmax = 1, vmin = 0
    )
    probe_R_vals = probe.array.compute()
    probe_R = Complex2RGB(probe_R_vals, vmax = 1, vmin = 0)
    probe_R_int = np.abs(probe_R_vals)**2
    probe_R_int /= np.max(probe_R_int)
    
    im0.set_data(probe_Q)
    im1.set_data(probe_R)
    im2.set_data(probe_R_int)
    fig.canvas.draw_idle()
    
# List of options
option_list = (
    'select an option',
    'uncorrected STEM',
    'uncorrected STEM at Scherzer',
    'aberration corrected STEM',
    'SEM',
    'boundary artifacts', 
)

# update the plots with a pre-selected function
def select_preset_eventhandler(change):
    if change.new == option_list[1]: #uncorrected STEM
        energy.value = 300
        defocus.value = 0
        C30.value = 1.3
        semiangle.value = 10
    
    if change.new == option_list[2]: #uncorrected STEM @ Scherzer
        #calculate scherzer 
        lambda_300kv = abtem.core.energy.energy2wavelength(300e3)
        C1_Scherzer = 0.87 * (1.3e7*lambda_300kv) ** 0.5 
        
        energy.value = 300
        defocus.value = C1_Scherzer
        C30.value = 1.3
        semiangle.value = 10
        
    if change.new == option_list[3]: #corrected STEM
        energy.value = 300
        defocus.value = 0
        C30.value = 0.001
        semiangle.value = 10
    
    if change.new == option_list[4]: #SEM
        energy.value = 20
        defocus.vaue = 0
        C30.value = 5
        semiangle.value = 2.5
    
    if change.new == option_list[5]: #artifacts
        energy.value = 10
        defocus.value = 2000
        C30.value = 0
        semiangle.value = 20     
            
# Widgets
dropdown = Dropdown(
    options = option_list,
    layout = Layout(width='200px',height='30px'),
)
dropdown.observe(select_preset_eventhandler, names='value')

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

energy = IntSlider(
    value=60, min=10, max=300, 
    step = 2,
    description = "energy (kV)",
    style = style,
)

defocus = IntSlider(
    value = 0, min = -2000, max = 2000, 
    step = 20,
    description = "defocus / -1*C1 (A)",
    style = style
)

C30 = FloatSlider(
    value = 0, min = 0, max =5, 
    step = 0.1,
    description = "C3 (mm)",
    style = style
)

semiangle = FloatLogSlider(
    value=20,
    base=10,
    min=0, # min exponent of base
    max=1.6021, # max exponent of base
    step=0.05, # exponent step
    description = "semiangle (mrad)",
    style = style,
)

widgets.interactive_output(
    update_probe, 
    {
        'energy':energy,
        'defocus':defocus,
        'semiangle':semiangle,
        'C30':C30,
    },
)

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

widget = VBox(
    [
        fig.canvas,
        HBox([
            VBox([
                energy,
                defocus,
                semiangle,
                C30,  
            ]),
            VBox([
                Label('Preset probes',layout=Layout(width='100px',height='30px')), 
                dropdown,
            ])
        ]),
    ],
)

display(widget);