# -*- 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 = 2000                         # number of points along string
dx = 1.0
dt = 0.2
nT = 3000                   # number of time steps to integrate
x0 = 80
w0 = 20
lbd = 30
k = 2*np.pi/lbd
beta = 2.

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_disp.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] ) - \
      ( dt/dx**2 )**2*ci[0]**2*beta**2*( 3*yi0[0] - 4*yi0[1] + yi0[2] )
    yip[1]   = 2.*yi0[1]   - yim[1]   + \
      ( dt/dx )**2*ci[1]**2*( yi0[2] - 2*yi0[1] + yi0[0] ) - \
      ( dt/dx**2 )**2*ci[1]**2*beta**2*( yi0[3] - 4*yi0[2] + 6*yi0[1] - 3*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] ) - \
      ( dt/dx**2 )**2*ci[N-1]**2*beta**2*( yi0[N-3] - 4*yi0[N-2] + 3*yi0[N-1] )
    yip[N-2] = 2.*yi0[N-2] - yim[N-2] + \
      ( dt/dx )**2/ci[N-2]**2*( yi0[N-3] - 2*yi0[N-2] + yi0[N-1] ) - \
      ( dt/dx**2 )**2*ci[N-2]**2*beta**2*(  yi0[N-4] - 4*yi0[N-3] + 6*yi0[N-2] - 3*yi0[N-1] )

    for i in range( 2, N-2 ):          # 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] ) - \
          ( dt/dx**2 )**2*ci[i]**2*beta**2*( yi0[i+2] - 4*yi0[i+1] + 6*yi0[i] - 4*yi0[i-1] + yi0[i-2] )
    
    if t % 20 == 0:                 # plot every 20th time step
        fig = plt.figure(figsize=(11,8), dpi=300)
        plt.plot( yip, '-o', ms=0.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