# import some packages
from __future__ import print_function
import numpy as np

########################
# define some functions:
########################

# right deriv
def Dp(u, dx, du):
    for i in range(0, len(u)-1):
        du[i] = (u[i+1] - u[i])/dx 

# left deriv
def Dm(u, dx, du):
    for i in range(1, len(u)):
        du[i] = (u[i] - u[i-1])/dx 

# centered deriv
def D0(u, dx, du):
    for i in range(1, len(u)-1):
        du[i] = (u[i+1] - u[i-1])/(2.0*dx) 

# RHS of $\partial_t u$
def eval_rhs(u, dx, du):
    Dm(u, dx, du)
    return -u * du

# u(t+dt) = u(t) + \partial_t u * dt
def calc_unew(u, rhs, dt):
    return u + rhs * dt


# discontinuiy
def set_discont_u0(x, u):
    for i in range(0, len(u)):
        if x[i] < 1.0:
            u[i] = 1.5
        else:
            u[i] = 0.5

# print colums with data at time t
def pr_timeframe(t, x, u):
    print("# time =", t)
    for i in range(0, len(u)):
        print(x[i], u[i])
    print()


###############
# main program:
###############

# grid:
x = np.linspace(0, 16, 161)
dx = x[1]-x[0]  # grid spacing
dtfac = 0.5
dt = dtfac*dx     # time step

# du will contain deriv of u, initialze to 0
du = np.zeros(len(x))

# initial u
t = 0.0
u = np.zeros(len(x))
set_discont_u0(x,u)
pr_timeframe(t, x, u)

timesteps = int(17.0/dt)

# loop over time steps, and print results
for n in range(0, timesteps):
    # make time step, using simple Euler method
    t = t + dt
    rhs = eval_rhs(u, dx, du)
    u = calc_unew(u, rhs, dt)
    # print colums with data
    pr_timeframe(t, x, u)