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

'''Simulate the movement of a string subject to different initial conditions'''

# the length of the string is always assumed to be equal to pi

# define the tension
T = 0.04
# define the line density
rho = (np.pi**2)

# sound speed
c = np.sqrt(T/rho)

# decide which of the functions defined below is used
testfun = 0
# 0...fundamental mode
# 1...2nd mode
# 2...3rd mode
# 3...4th mode
# 4...triangular wave
# 5...approximated sawtooth wave
# 6...quadratic function (Fourier series)
# 7...quadratic function (explicit definition)
# 8...???
# 9...approximated rectangular wave
# 10...rectangular wave

# switch through the different functions
if testfun == 0:
    # sinoidal wave
    def f(x):
        result = np.sin(x)
        return result
    scale = 1.5
elif testfun == 1:
    # sinoidal wave
    def f(x):
        result = np.sin(2*x)
        return result
    scale = 1.5
elif testfun == 2:
    # sinoidal wave
    def f(x):
        result = np.sin(3*x)
        return result
    scale = 1.5
elif testfun == 3:
    # sinoidal wave
    def f(x):
        result = np.sin(4*x)
        return result
    scale = 1.5
elif testfun == 4:
    # triangular wave (given by its Fourier coefficients)
    def f(x):
        result = 0*x
        for k in range(1,30):
            result += ((-1)**(k+1))*np.sin((2*k-1)*x)/((2*k-1)**2)
        return result
    scale = 2
elif testfun == 5:
    # approximation to a sawtooth wave
    def f(x):
        result = 0*x
        for k in range(1,15):
            result += np.sin(k*x)/k
        return result
    scale = 3
elif testfun == 6:
    # quadratic function (given by its Fourier coefficients)
    def f(x):
        result = 0*x
        for k in range(1,25):
            result += np.sin((2*k-1)*x)/((2*k-1)**3)
        return result
    scale = 1.5
elif testfun == 7:
    # the same quadratic function given explicitly
    # note that we need the odd, periodic extension
    def f(x):
        # shift x to the interval -pi,pi
        x_shift = (x+np.pi)%(2*np.pi)-np.pi
        result = x_shift*(np.pi-abs(x_shift))
        # scale f correctly - that is: same scaling as for function 3
        result = np.pi*result/8
        return result
    scale = 1.5
elif testfun == 8:
    # ???
    def f(x):
        result = 0*x
        for k in range(1,25):
            result += np.sin((2*k-1)*x)/((2*k-1)**2)
        return result
    scale = 1.5
elif testfun == 9:
    # approximation of rectangular wave
    def f(x):
        result = 0*x
        for k in range(1,100):
            result += np.cos(0.95*(2*k-1))*np.sin((2*k-1)*x)/((2*k-1))
        return result
    scale = 1.5
elif testfun == 10:
    # rectangular wave
    def f(x):
        x_shift = (x+np.pi)%(2*np.pi)-np.pi
        result = np.sign(x_shift)*(1+np.sign(np.pi/2-abs(abs(x_shift)-np.pi/2)))
        return result
    scale = 3.0

# space discretisation
x = np.arange(0, np.pi, 0.01)

# create an empty figure
fig = plt.figure()
# define the size of the figure
ax = plt.axes(xlim = (0,np.pi), ylim = (-scale,scale))
# plot an (empty) line - content comes later
line, = ax.plot([], [], lw=2)

# initialize the plot (background)
# the object that is updated in each frame has to be returned by this function
def init():
    line.set_data([], [])
    return line,

# define the frames of the animation
# again, the function has to return the object that change
def animate(i):
    # here 
    line.set_data(x,(f(x+c*i)+f(x-c*i))/2) 
    return line,


# define the routine that creates the animation
# np.arange...basically defines the length of the animation; replace 100 by some
#             larger number for longer animations
# interval...time interval between different frames
# blit...if true, only the parts of the animation that have changed are redrawn
#        may speed up the animation
# the animation will be created online - here we define only its 'constructor'
ani = animation.FuncAnimation(fig, animate, np.arange(1, 100), init_func=init,
    interval=50, blit=True)

# if you want to save the animations (as ogg with 30 frames per second), uncomment
# the following code - you need to have ffmpeg.exe installed, thus this code is
# not guaranteed to work on every computer
# ani.save('anim.ogg',fps=30)


# run the animation
plt.show()