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

ex = int(sys.argv[1]) #Example 3.1, 3.2 or 3.4

#Parameters
N = 15
M = 55
T = 0.1 #Time horizon
K = 40   #Number of terms in the Fourier series

dt = float(T)/float(M)
dx = 1./float(N+1)
r = dt/(dx**2)
print 'T:', T,'K:', K,'r:', r, 'dx:', dx, 'dt:', dt

x = np.linspace(0,1,N+2)
t = np.linspace(0,T,M)
v = np.zeros([M,N+2]) #Explicit
w = np.zeros([M,N+2]) #Implicit

#Initial values
def f(x):
    if ex == 1:
        return 3.*np.sin(np.pi*x)+5.*np.sin(4.*np.pi*x)
    elif ex == 2:
        return np.ones(x.size)
    else:
        return x

v[0] = f(x)
w[0] = f(x)

#Computing solution, explicit scheme
for m in range(M-1):
    for j in range(1,N+1):
        v[m+1][j] = r*v[m][j-1] + (1-2.*r)*v[m][j] + r*v[m][j+1]

#Implicit scheme
A = 2.*np.eye(N)
for i in range(N-1):
    A[i][i+1] = -1.
    A[i+1][i] = -1.
B = np.eye(N) + r*A
B = np.linalg.inv(B)
for m in range(M-1):
    w[m+1][1:N+1] = np.dot(B,w[m][1:N+1])

#Fourier solution
def u(t,K):
    k = np.arange(1,K+1,1.)
    X, K = np.meshgrid(x,k)
    if ex == 1:
        return 3.*np.exp(-t*np.pi**2)*np.sin(np.pi*x) + 5.*np.exp(-t*16.*np.pi**2)*np.sin(4.*np.pi*x)
    elif ex == 2:
        return np.sum((4./np.pi)*1./(2.*K-1.)*np.exp(-t*((2.*K-1.)*np.pi)**2)*np.sin((2.*K-1)*np.pi*X),0)
    else:
        return  np.sum((2./np.pi)*((-1)**(K+1)/K)*np.exp(-t*(K*np.pi)**2)*np.sin(K*np.pi*X),0)



#Plot at time T.
plt.plot(x,v[M-1], label='Explicit scheme')
plt.plot(x,w[M-1], label='Implicit scheme')
plt.plot(x,u(T,K), label='Fourier sol.')
legend = plt.legend(loc='upper left')
plt.show()