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

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

# exact flux
def f(u):
    return 0.5*u*u

# flux at interface
def F_interface(ul, ur):
    return f(0.5*(ul+ur))

# flux F_{m+1/2}
def Fp(u, m):
    ul = u[m]
    if m < len(u)-1:
        ur = u[m+1]
    else:
        ur = ul
    return F_interface(ul, ur)

# flux F_{m-1/2}
def Fm(u, m):
    ur = u[m]
    if m > 0:
        ul = u[m-1]
    else:
        ul = ur
    return F_interface(ul, ur)
    
# RHS of $\partial_t u$
def eval_rhs(u, dx, dF):
    for i in range(0, len(u)):
        # print(i, Fp(u, i), Fm(u, i))
        dF[i] = (Fp(u, i) - Fm(u, i))/dx
    return -dF

# 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

# dF will contain deriv of num. flux, initialze to 0
dF = 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, dF)
    u = calc_unew(u, rhs, dt)
    # print colums with data
    pr_timeframe(t, x, u)