# ————————————————————————————————————————————————————————————————————————————————————————————————
# Import Libraries
# ————————————————————————————————————————————————————————————————————————————————————————————————
import numpy as np
from numpy.linalg import inv
from scipy.optimize import curve_fit
from scipy.ndimage import convolve, generate_binary_structure
from scipy import sparse
from scipy.sparse.linalg import eigsh
from scipy.sparse.linalg import eigs
import scipy.constants as spc
from scipy import signal
import matplotlib as mpl
import matplotlib.pyplot as plt
from ipywidgets import interactive, fixed
import matplotlib.animation
from matplotlib.animation import FuncAnimation
from IPython.core.display import HTML
from IPython.display import display, Math

%matplotlib inline
%config InlineBackend.figure_format = 'retina'
plt.rcParams.update({
    'mathtext.fontset':'cm', 
    'font.family':'serif', 
    'font.sans-serif':'Times New Roman', 
    'figure.figsize':[10,6], 
    'figure.titlesize':22,
    "figure.dpi":90, 
    'savefig.dpi':90,
    'axes.titlesize':20, 
    'axes.labelsize':18, 
    'axes.titley': 1.0, 
    'axes.titlepad': 5.0,
    'axes.edgecolor':'black', 
    'axes.grid': True, 
    'grid.alpha': .5,
    'xtick.labelsize':14,'ytick.labelsize':14,
    'xtick.major.size':6,'ytick.major.size':6,
    'xtick.major.width':1.25, 'ytick.major.width':1.25, 
    'xtick.direction':'inout','ytick.direction':'inout',
    'xtick.top':False, 'ytick.right':False,
    'legend.title_fontsize':14, 'legend.fontsize':14,
    'legend.borderaxespad': 1, 'legend.borderpad': 0.5,
    'legend.framealpha': 1,
    'legend.handleheight': 0.5, 'legend.handlelength': 2.0, 'legend.handletextpad':0.5,
    'legend.labelspacing': 0.25,
    'legend.fancybox':False,
    'legend.edgecolor': '0',
    'legend.frameon': True,
    'legend.markerscale': 1.25,
    'animation.embed_limit':2**128,
    'animation.html': 'jshtml'})
pi = spc.pi
μ0 = spc.mu_0
ε0 = spc.epsilon_0
c0 = 1/np.sqrt(μ0*ε0)
z0 = np.sqrt(μ0/ε0)

j_max = 500
j_source = 10
n_max = 4_000

ε = np.ones(j_max)*ε0
g = 10
ε[int(j_max/2):int(j_max/2 + 100)] = g*ε0

Ex = np.zeros(j_max)
Hz = np.zeros(j_max)

Ex_prev = np.zeros(j_max)
Hz_prev = np.zeros(j_max)

λ_min = 350e-9
dy = λ_min/20
dt = dy/c0

y = np.arange(j_max)*dy
t = np.arange(n_max)*dt
Y,T = np.meshgrid(y,t)

def source(T):
    λ0 = 550e-9
    w0 = 2*pi*c0/λ0
    tau = 100
    T0 = tau*3
    return np.exp(-((T-T0)/tau)**2) * np.sin(w0*T*dt)

E = np.zeros((n_max, j_max))
H = np.zeros((n_max, j_max))

for n in range(n_max):
    # update H boundaries
    Hz[j_max - 1] = Hz_prev[j_max - 2]
    
    # update H
    Hz[:j_max-1] = Hz_prev[:j_max-1] + dt/(dy*μ0) * (Ex[1:j_max] - Ex[:j_max-1])
    Hz_prev = Hz
        
    # include H source
    Hz[j_source-1] -= source(n)/z0
    Hz_prev[j_source-1] = Hz[j_source-1]
    
    # update E boundaries
    Ex[0] = Ex_prev[1]
    
    # update E
    Ex[1:] = Ex_prev[1:] + dt/(dy*ε[1:]) * (Hz[1:] - Hz[:j_max-1])
    Ex_prev = Ex
        
    # include E source
    Ex[j_source] += source(n+1)
    Ex_prev[j_source] = Ex[j_source]
    
    H[n] = Hz
    E[n] = Ex

S = -(E*H)

%matplotlib inline
%config InlineBackend.figure_format = 'retina'
fig, ax = plt.subplots(1,1,figsize = (7,6))
cax = ax.pcolormesh(Y, T, S, shading = 'auto', cmap = 'magma',vmin = -.01*S.max(), vmax = .01*S.max())
ax.set_xlabel(r'$y$')
ax.set_ylabel(r'$t$')
fig.tight_layout()
plt.colorbar(cax)
plt.show()

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

fig, ax = plt.subplots(1,1,figsize = (6,6), facecolor = 'white')

line, = ax.plot([], [], 'r-', lw = 2)
plt.axvspan(y[ε == g*ε0].min(),y[ε == g*ε0].max(), alpha=0.5, color='y')
plt.axvspan(0,y[ε == g*ε0].min(), alpha=0.5, color = 'white')
plt.axvspan(y[ε == g*ε0].max(),y.max(), alpha=0.5, color = 'white')
ax.set_xlim(y[0]+dy, y[-1]-dy)
ax.set_ylim(S.min(), S.max())
ax.set_xlabel(r'$y$')
ax.set_ylabel(r'$S_y$')
ax.grid(False)
ax.axis(False)

def animate(i):
    line.set_data(y, S[i])
    return line,
plt.close()

FPS = 24
N_FRAMES = 200
anim = FuncAnimation(fig, animate, interval = int(1e3/FPS), frames = np.arange(0, n_max, int(n_max/N_FRAMES)), repeat = False)

%matplotlib inline
%config InlineBackend.figure_format = 'retina'
display(anim)
# writergif = mpl.animation.PillowWriter(fps = 24) 
# anim.save('dielectric.gif', writer = writergif)

Once Loop Reflect