#!/usr/bin/python3.5


import sys
import scipy
import scipy.linalg
import numpy as np
import matplotlib
matplotlib.use('TkAgg') 
import matplotlib.pyplot as plt
import pylab
import time as tim

    #grid and array
nx = 100
u  = np.zeros (nx)
total_length = 1
width = 0.1
dx = total_length / float (nx)
x  = np.arange (0, total_length, dx)

tmax = 0.02
time = 0.0

#time step
dt = 0.5 * dx ** 2  * 1
print('Time step:',dt)

    # initial conditions
u = np.exp (- (x-0.5) ** 2 / (2 * width ** 2))
u0 = u
u_exact = u

plt.plot (x,u)
plt.ylim ( (-2,2) )
plt.xlabel ('x')
plt.ylabel ('u')
plt.draw()
plt.show(block=False)
plt.ioff() 
    
while (time < tmax):

    print("time:",time)

# set up tridiagonal system for Dirichlet boundary conditions (u=0)
    coeff = np.zeros([3,len(x)])
    coeff [1,:] =   2.0 / dx ** 2 + 1.0 / dt
    coeff [0,:] =  -1.0 / dx ** 2
    coeff [2,:] =  -1.0 / dx ** 2

#    du/dt + (u_(i+1)-2u(i)-u(i-1))/dx**2 = 0 
#    u (i) *(1/dt - 2/dx**2) + 

# Equation to solve:
# du / dt +  A . (u + du)=0
# (A + 1 / dt) . du = - A . u
# Note that coeff is the matrix A + 1/dt, not A alone
    b = - 2 * u / dx ** 2 
    b [1:] = b [1:] + u[0:-1] / dx ** 2 
    b [0:-1] = b [0:-1] + u[1:] / dx ** 2 
    du = scipy.linalg.solve_banded ( (1, 1), coeff, b, overwrite_ab = True, overwrite_b = True)
#    print(du)
    u = u + du
    time = time + dt

    # analytic solution
    u_exact = np.exp (- (x-0.5) ** 2 / (2 * (width **2 + 2*time))) * np.sqrt(width**2/(width**2+2*time))

    plt.clf()
    plt.ylim ( (-2,2) )
    plt.xlabel ('x')
    plt.ylabel ('u')
    c1, = plt.plot(x,u0,label='initial conditions')
    c2, = plt.plot(x,u,label='numerical solution')
    c3, = plt.plot(x,u_exact,ls='dashed',label='analytic solution')
    plt.legend(handles=[c1,c2,c3])
    plt.draw()
    plt.show(block = False)
    plt.pause(0.05)

plt.ion()
plt.clf()
plt.ylim ( (-2,2) )
plt.xlabel ('x')
plt.ylabel ('u')
c1, = plt.plot(x,u0,label='initial conditions')
c2, = plt.plot(x,u,label='numerical solution')
c3, = plt.plot(x,u_exact,ls='dashed',label='analytic solution')
plt.legend(handles=[c1,c2,c3])
plt.draw()
plt.show(block = True)

        
        
