import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.animation import PillowWriter

matplotlib.rcParams['font.sans-serif'] = "Arial"
# Then, "ALWAYS use sans-serif fonts"
matplotlib.rcParams['font.family'] = "sans-serif"


# Create the figure
fig = plt.figure()
ax1 = fig.add_subplot(1,1,1)

# Some axis variables
xlabel = 'E'
ylabel = 'N'

x1 = -3
x2 = 3

y1 = 0
y2 = 5

# Setting axis parameters
ax1.tick_params('both', labelsize=14, width=1,direction='in', pad=10);
ax1.set_xlabel(xlabel, fontsize=16,labelpad=12)
ax1.set_ylabel(ylabel,fontsize=16,labelpad=12)
ax1.set_xlim(x1,x2);
ax1.set_ylim(y1,y2)
ax1.spines['bottom'].set_linewidth(1.0)
ax1.spines['left'].set_linewidth(1.0)
ax1.spines['top'].set_linewidth(1.0)
ax1.spines['right'].set_linewidth(1.0)

plt.subplots_adjust(bottom=0.2)


# Function to plot
def dos(E,D):
    if np.abs(E)>D:
        return np.abs(E)/np.sqrt(E**2-D**2)
    else:
        return 0

# The parameter that varies with time
def DeltaT(T):
    if(T<1):
        return (1.0-T**3.4951)**0.54314	
    else:
        return 0
    
# Set up the plot
xlist = []
ylist = []
plot, = ax1.plot(xlist,ylist, linewidth=1, linestyle="-", color='red')

# Set up the gif
outfile = 'animation.gif'
metadata = dict(title='Movie',artist='user')
writer = PillowWriter(fps=15,metadata=metadata)


xlist = np.arange(-3,3,0.001)

# Loop from T = 0 to 1.2 in 50 steps
# calculate the function at each step
# and save each frame to the gif.
with writer.saving(fig, outfile, 100):
    for Tval in np.linspace(0, 1.2, 50):
        delta = DeltaT(Tval)     
        ylist=[]
        for xval in xlist:
                ylist.append(dos(xval,delta))

        plot.set_data(xlist,ylist)
        writer.grab_frame()