%matplotlib widget
import numpy as np
import matplotlib.pyplot as plt
import ipywidgets

def central_equation_entry(g_vectors,potential,k_vec,i,j,):
    """ returns the (i,j)th entry of the central matrix """
    gi = g_vectors[i]
    gj = g_vectors[j]

    if i==j:
        val = np.dot(k_vec-gi,k_vec-gi)
    else:
        val = potential[i,j]
    return val

def central_equation(
    g_vectors,
    potential,
    k_vec
):
    """ loops over potential indices to make central matrix """
    array = np.zeros_like(potential)
    n,m = array.shape
    for i in range(n):
        for j in range(m):
            array[i,j] = central_equation_entry(g_vectors,potential,k_vec,i,j)
    return array

def central_equation_eigenvalues(
    g_vectors,
    potential,
    k_vec
):
    """ diagonalizes central matrix """
    array = central_equation(g_vectors,potential,k_vec)
    return np.sort(np.linalg.eigvalsh(array).real)[:3]

potential_1d = np.zeros((5,5))
g_vectors_1d = np.arange(-2,3)[:,None]
k_vecs_1d = np.linspace(-0.5,0.5,48+1) # first BZ

bandstructure_1d = np.array([
    central_equation_eigenvalues(
        g_vectors_1d,
        potential_1d,
        k
    )
    for k in k_vecs_1d
])

bandstructure_1d -= bandstructure_1d[24,0]

dpi = 72
with plt.ioff():
    fig, ax = plt.subplots(figsize=(675/dpi,375/dpi),dpi=dpi)
ax.set_prop_cycle(color=['red', 'blue','black'])
lines = ax.plot(k_vecs_1d,bandstructure_1d)
ax.set(
    xlabel=r"wavevector, k ",
    ylabel="blochstate eigenvalue, E"
)
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 = '400px'
fig.canvas.toolbar_position = 'bottom'
fig.tight_layout()
None

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

def update_blochstates(change):
    """ """
    V0 = change["new"]
    vec = np.full(4,V0)
    potential_1d = np.diag(vec,1) + np.diag(vec,-1)
    
    bandstructure_1d = np.array([
        central_equation_eigenvalues(
            g_vectors_1d,
            potential_1d,
            k
        )
        for k in k_vecs_1d
    ])
    bandstructure_1d -= bandstructure_1d[24,0]

    lines[0].set_ydata(bandstructure_1d[:,0])
    lines[1].set_ydata(bandstructure_1d[:,1])
    lines[2].set_ydata(bandstructure_1d[:,2])
    fig.canvas.draw_idle()
    return None

V0_slider = ipywidgets.FloatSlider(
    min=0,
    max=0.5,
    step=0.025,
    layout=layout,
    style=style,
    description=r"potential strength, V$_0$"
)

V0_slider.observe(update_blochstates,"value")

ipywidgets.VBox(
    [
        V0_slider,
        fig.canvas
    ],
    layout=ipywidgets.Layout(
        align_items="center"
    )
)