%matplotlib ipympl
import abtem
import ase
import matplotlib.pyplot as plt
import numpy as np

from IPython.display import display
import ipywidgets

srtio3 = ase.io.read("data/SrTiO3.cif")

# widget figure generation
with plt.ioff():
    dpi = 72
    fig, (ax1,ax2) = plt.subplots(1,2,figsize=(675/dpi, 275/dpi), dpi=dpi)

# 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 = '675px'
fig.canvas.layout.height = '315px'
fig.canvas.toolbar_position = 'bottom'

def return_sto_supercell(x, y, z, vac, axis):
    """ """
    if axis == [1,1,1]:
        layers = 3
        z = np.clip(z // 3,1,np.inf).astype("int")
    elif axis == [1,1,0]:
        layers = 2
        z =  np.clip(z // 2,1,np.inf).astype("int")
    elif axis == [1,0,0]:
        layers = 2
        z =  np.clip(z // 2,1,np.inf).astype("int")
    else:
        raise ValueError()
    
    cell = ase.build.surface(srtio3, indices=axis, layers=layers, periodic=True)
    cell = abtem.orthogonalize_cell(cell)
    cell.wrap()
    cell *= (x,y,z)
    cell.center(vac, axis=2)

    ax1.cla()
    ax2.cla()
    top = abtem.show_atoms(cell, ax=ax1,show_periodic = True)
    beam = abtem.show_atoms(cell, ax=ax2, plane='yz', show_periodic=True, legend=True)
    fig.canvas.draw_idle()
    
    return None

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

layout = ipywidgets.Layout(width='325px',height='30px')

dropdown = ipywidgets.Dropdown(
    options = [[1,0,0],[1,1,0],[1,1,1]],
    value=[1,1,0],
    layout=layout,
    style=style,
    description='Zone axis:     ',
)

sliderx = ipywidgets.IntSlider(
    orientation='horizontal',
    description='x repetitions: ',
    value=4,
    min=1,
    max=10,
    style = style,
    layout=layout,
)

slidery = ipywidgets.IntSlider(
    orientation='horizontal',
    description='y repetitions: ',
    value=4,
    min=1,
    max=10,
    style = style,
    layout=layout,
)

sliderz = ipywidgets.IntSlider(
    orientation='horizontal',
    description='z repetitions: ',
    value=4,
    min=1,
    max=10,
    style = style,
    layout=layout,
)

slidervac = ipywidgets.FloatSlider(
    orientation='horizontal',
    description='vacuum (A): ',
    value=2,
    min=1,
    max=10,
    step=1,
    style = style,
    layout=layout,
)

ipywidgets.interactive_output(
    return_sto_supercell, 
    {
        'x':sliderx,
        'y':slidery,
        'z':sliderz,
        'vac':slidervac,
        'axis':dropdown
    },
)
None

display(
    ipywidgets.VBox(
        [
            fig.canvas,
            ipywidgets.HBox([
                ipywidgets.VBox([
                    sliderx,
                    slidery,
                    sliderz,
                ]),
                ipywidgets.VBox([
                    slidervac,
                    dropdown,
                ])
            ]),
        ],
    )
)