#!/usr/bin/python3.5

#if __name__ == "__main__":

import sys
import scipy
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.01
time = 0.0

#time step
dt = 0.5 * dx ** 2 * 0.8

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()
plt.show(block=False)
plt.ioff() 
    
while (time < tmax):

    print("time:",time)
    dudx_left =  (np.roll (u, 0) - np.roll (u,1)) / dx
    dudx_rght =  (np.roll (u,-1) - np.roll (u,0)) / dx
# boundary conditions: no flux 
#    dudx_left[0 ] = 0.0
#    dudx_rght[-1] = 0.0
    d2udx2 = (dudx_rght - dudx_left) / dx
# for Dirichlet boundary conditions:
    d2udx2 [ 0]  = 0.0
    d2udx2 [-1]  = 0.0

    u = u + dt * d2udx2
    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)

        
        
