# -*- coding: utf-8 -*-
"""
Created on Wed Mar 14 13:27:37 2018

@author: lbnc
"""

import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
import numpy as np

N = 500                         # number of points along string
dx = 1.0
dt = 0.5
nT = 1200                   # number of time steps to integrate
x0 = 80
w0 = 20
lbd = 16
k = 2*np.pi/lbd

ci = np.ones( N )
#ci[0:int(N/2)] *= 0.5
#ci = np.arange( N )/( N - 1 )*0.8 + 0.2

xs = np.arange( N )*dx
yi0 = np.zeros( N )                      # set up initial values
yi0[:] = np.sin( k*( xs - x0) )
yi0[:] *= np.exp( -( xs - x0)**2/( 2*w0**2 ) )
yim = np.zeros( N )                      # set up initial values
yim[:] = np.sin( k*( xs - x0 + dt*ci[x0]) )
yim[:] *= np.exp( -( xs - x0 + dt*ci[x0])**2/( 2*w0**2 ) )

yip = np.zeros( N )

pdf_pages = PdfPages('waveeq.pdf')   # open PDF file for printing

for t in range( 0, nT ):                       # loop over time

    yip[0]   = 2.*yi0[0]   - yim[0]   + \
      ( dt/dx )**2*ci[0]**2*( yi0[1] - yi0[0] )
    
    yip[N-1] = 2.*yi0[N-1] - yim[N-1] + \
      ( dt/dx )**2/ci[N-1]**2*( yi0[N-2] - yi0[N-1] )

    for i in range( 1, N-1 ):          # integrate along string
        yip[i] = 2.*yi0[i] - yim[i] + ( dt/dx )**2*ci[i]**2*( yi0[i+1] - 2*yi0[i] + yi0[i-1] )
    
    if t % 20 == 0:                 # plot every 20th time step
        fig = plt.figure(figsize=(11,8), dpi=300)
        plt.plot( yip, '-o', ms=.5 )
        plt.ylim( [ -1, 1 ] )
        plt.title('t = '+str(t) )
        pdf_pages.savefig(fig)
        plt.close(fig)
     
    yim = np.copy(yi0)                            # copy arrays
    yi0 = np.copy(yip)
    
pdf_pages.close()          # Write the PDF document to the disk