import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib import rc
import ipywidgets as ipw

from ipywidgets import interact,interactive
import ipywidgets as widgets

from IPython.display import HTML
rc('text', usetex=True)

%matplotlib widget

global global_K1
global global_K2
global global_branche
global global_selected_k

def get_closest_point(clic_x,clic_y):
    # Get the data from the figure
    ax= plt.gca() # get axis handle
    acoustic_data = ax.lines[1]
    optic_data = ax.lines[0]

    # Find which line was closer
    #closet x data
    k_a = acoustic_data.get_xdata()
    k_o = optic_data.get_xdata()
    diff_k_a = np.abs(k_a-clic_x)
    diff_k_o = np.abs(k_o-clic_x)
    # get the x index corresponding to the closet k in the graph for both branches
    index_a = np.argmin(diff_k_a)
    index_o = np.argmin(diff_k_o)

    # compute the distance on the y axis
    y_diff_a = np.abs(acoustic_data.get_ydata()[index_a]-clic_y)
    y_diff_o = np.abs(optic_data.get_ydata()[index_o]-clic_y)

    ac_is_closer = y_diff_a < y_diff_o
    # print a x at the postion 
    global global_branche
    global global_selected_k
    global global_omega
    
    if(ac_is_closer):
        global_branche=1
        global_selected_k = acoustic_data.get_xdata()[index_a]
        global_omega = acoustic_data.get_ydata()[index_a]
        return [acoustic_data.get_xdata()[index_a], acoustic_data.get_ydata()[index_a]]
    else:
        global_branche=2
        global_selected_k = optic_data.get_xdata()[index_a]
        global_omega = optic_data.get_ydata()[index_a]

        return [optic_data.get_xdata()[index_a], optic_data.get_ydata()[index_a]]

def interactive_lattice_plot(G,K):
    '''
    :param G,K: coupling
    plot acoustic and optical phonons bands
    '''
    
    # save the current values of G and K    
    global global_K1
    global global_K2
    global_K1=G
    global_K2=K
    
    k = np.linspace(-np.pi/2,np.pi/2,100)
    
    omega_o = G+K-np.sqrt((G+K)**2-4*G*K*np.sin(k)**2)
    omega_a = G+K+np.sqrt((G+K)**2-4*G*K*np.sin(k)**2)
    
#     print(omega_o, omega_a)
    ax1.clear();
    ax1.plot(k,np.sqrt(omega_a))
    ax1.plot(k,np.sqrt(omega_o))
    ax1.set_xlabel("k")
    ax1.set_ylabel("$\omega$")
    ax1.set_xticks(np.linspace(-np.pi/2,np.pi/2,5))
    ax1.set_xticklabels(['$-\pi/a$','','0','','$\pi/a$'])
       
    def onclick(event):    
        # clear the previous point 
        for i in np.arange(3, len(ax1.lines)+1) :
            ax1.lines.pop()
        # Get the coordonate of the data closest from the click
        [x,y]=get_closest_point(event.xdata,event.ydata)
        ax1.plot(x, y, '.',color='red')
    
    cid = fig.canvas.mpl_connect('button_press_event', onclick)
    plt.show()

g = ipw.FloatSlider(min=1, max=6, step=0.05, value=6);
k = ipw.FloatSlider(min=1, max=6, step=0.05, value=1);
fig=plt.figure()
ax1=plt.axes()
plt.ion()
interactive(interactive_lattice_plot, G=g, K=k)

Canvas(toolbar=Toolbar(toolitems=[('Home', 'Reset original view', 'home', 'home'), ('Back', 'Back to previous …

interactive(children=(FloatSlider(value=6.0, description='G', max=6.0, min=1.0, step=0.05), FloatSlider(value=…

fig2 = plt.figure()
ax2 = plt.axes()
ax2.set_xlim([-2,10])
ax2.set_xticks([])
ax2.set_yticks([])
line_u, = ax2.plot([], [], "ok")
line_v, = ax2.plot([], [], "ob")
if global_branche == 1:
    ax2.title.set_text("Acoustic")
else : 
    ax2.title.set_text("Optic")

a = 2 # lattice parameter
d = 0.5 # distance between atoms 1 and 2 in basis 
N=5 # number of repeated basis
pos_u = np.arange(N)*a 
pos_v = np.arange(N)*a+d
ax2.plot(pos_u,np.zeros(N),"ok",mfc='none')
ax2.plot(pos_v,np.zeros(N),"ob",mfc='none')
# Initialization function: plot the background of each frame
def init():
    line_u.set_data([], [])
    line_v.set_data([], [])

    return line_u,

# Animation function which updates figure data.  This is called sequentially
def animate(iteration):
    t=0.05*iteration
    a = 2 # lattice parameter
    d = 0.5 # distance between atoms 1 and 2 in basis 
    N=5 # number of repeated basis 
    global global_branche # acoustic or optical selected in the previous graph
    global global_selected_k # wave number selected in the previous graph
    relative_phase = (-1)** (1+global_branche) # relative phase between atom 1 and 2 in the basis

    global global_K1
    global global_K2 # interactions 
    
    a_nu = d/3.45*(relative_phase*(global_K1+global_K2*np.exp(1j*global_selected_k*a))/np.abs(global_K1+global_K2*np.exp(1j*global_selected_k*a)))
    b_nu = d/3.45 
    
    ns = np.arange(N)
    u_n = np.real(a_nu*np.exp(1j*global_selected_k*ns*a-1j*global_omega*t))
    v_n = np.real(b_nu*np.exp(1j*global_selected_k*ns*a-1j*global_omega*t))
    
    pos_u = np.arange(N)*a + u_n
    pos_v = np.arange(N)*a+d+v_n
    
    line_u.set_data(np.asarray([pos_u,np.zeros(N)]))
    line_v.set_data(np.asarray([pos_v,np.zeros(N)]))
    return line_u,

# Call the animator
anim = animation.FuncAnimation(fig2, animate, init_func=init, frames=300, interval=30, blit=True)