#!/usr/bin/env python
# Christoph Federrath 2020

import os
import sys
import argparse
from time import sleep
import numpy as np
import matplotlib
from matplotlib import rcParams
import matplotlib.pyplot as plt
import matplotlib.animation as animation


# define globals
ndim    = 1 # the dimensionality (1,2)
nguard  = 2 # number of guard cells
nx      = None # number of cells x
ny      = None # number of cells y
dx      = None # cell size x
dy      = None # cell size y
x       = None # cell-centred coordinates x
y       = None # cell-centred coordinates y
rho     = None # mass density
rhoux   = None # momentum density x
rhouy   = None # momentum density y
rhoetot = None # energy density
uxi     = None # velocity (face-centred) x
uyi     = None # velocity (face-centred) y
gamma   = None # adiabatic index


# === START init_grid ===
def init_grid(n=100, xr=[0.0,1.0], yr=[0.0,1.0], boundary='periodic'):
    # globals
    global nguard, nx, ny, dx, dy, x, y, rho, rhoux, rhouy, rhoetot, uxi, uyi
    # total number of cell centers (including guard cells)
    if type(n) is not list:
        nx = n + 2*nguard
        if ndim == 2:
            ny = n + 2*nguard
    else:
        nx = n[0] + 2*nguard
        if ndim == 2:
            ny = n[1] + 2*nguard
    # cell width
    dx = (xr[1]-xr[0]) / (nx-2*nguard)
    if ndim == 2:
        dy = (yr[1]-yr[0]) / (ny-2*nguard)
    # coordinates at cell centers
    x = (np.arange(nx)-nguard)*dx + dx/2.0 + xr[0]
    if ndim == 2:
        y = (np.arange(ny)-nguard)*dy + dy/2.0 + yr[0]
    # define containers to store density, momentum, energy, etc.
    zeros = None
    if ndim == 1:
        zeros = np.zeros(nx)
    if ndim == 2:
        zeros = np.zeros([nx,ny])
    rho = np.copy(zeros) # mass density
    rhoux = np.copy(zeros) # momentum density x
    if ndim == 2:
        rhouy = np.copy(zeros) # momentum density y
    rhoetot = np.copy(zeros) # energy density
    uxi = np.copy(zeros) # velocity at interfaces x
    if ndim == 2:
        uyi = np.copy(zeros) # velocity at interfaces y
# === END init_grid ===

# ================== START pres ===================
def pres():
    # compute and return pressure
    return (gamma-1.0) * rho * eint()
# ================== END pres ===================

# ================== START eint ===================
def eint():
    # compute and return the specific internal energy
    ei = rhoetot/rho - 0.5*(rhoux/rho)**2
    if ndim == 2:
        ei -= 0.5*(rhouy/rho)**2
    return ei
# ================== END eint ===================

# ================== START compute_dt ===================
def compute_dt(cfl):
    # compute sound speed
    cs = np.sqrt(gamma*pres()/rho)
    # compute maximum signal speed
    max_speed = np.max(abs(rhoux/rho) + cs)
    # compute timestep based on CFL criterion
    dt = cfl * dx / max_speed
    if ndim == 2:
        dt = np.min( [ dt, cfl * dy / np.max(abs(rhouy/rho) + cs) ] )
    return dt
# ================== END compute_dt ===================

# === START advect ===
def advect(q, dt, sweep='x', scheme='upwind', fluxlim='mc'):

    if ndim == 2: # no flux limiters implemented in 2D (yet)
        # sweep direction
        if sweep == 'x':
            ui = uxi
            axis = 0
        if sweep == 'y':
            ui = uyi
            axis = 1
        # positive velocity indices
        ind_up = ui >= 0
        ind_um = ui <  0
        # now construct the flux
        flux = np.zeros([nx,ny])
        if scheme == 'upwind':
            flux[ind_up] = ui[ind_up] * np.roll(q,1,axis=axis)[ind_up]
            flux[ind_um] = ui[ind_um] * np.roll(q,0,axis=axis)[ind_um]
        if scheme == 'centred':
            flux = ui * (np.roll(q,0,axis=axis)+np.roll(q,1,axis=axis))/2.0
        if scheme == 'downwind':
            flux[ind_up] = ui[ind_up] * np.roll(q,0,axis=axis)[ind_up]
            flux[ind_um] = ui[ind_um] * np.roll(q,1,axis=axis)[ind_um]
        # finally, update q
        q[:] = q - dt * (np.roll(flux,-1,axis=axis)-flux) / dx
        return

    # determine the r_{i-1/2} for the flux limiter
    r = np.zeros(nx)
    # derivative of q
    dq = q - np.roll(q,1)
    # positive derivative of q indices
    ind_qp = abs(dq) > 0
    # positive velocity indices
    ind_up = uxi >= 0
    # negative velocity indices
    ind_um = uxi < 0
    # assign to r in preparation of slope limiters
    r[ind_qp & ind_up] = ((np.roll(q, 1)-np.roll(q,2)))[ind_qp & ind_up] / dq[ind_qp & ind_up]
    r[ind_qp & ind_um] = ((np.roll(q,-1)-np.roll(q,0)))[ind_qp & ind_um] / dq[ind_qp & ind_um]
    # phi, flux
    phi  = np.zeros(nx)
    flux = np.zeros(nx)
    # Determine the flux limiter (many other flux limiters can be implemented here!)
    if fluxlim == 'donor-cell':
        phi[:] = 0.0
    if fluxlim == 'lax-wendroff':
        phi[:] = 1.0
    if fluxlim == 'beam-warming':
        phi = r
    if fluxlim == 'fromm':
        phi = 0.5*(r+1.0)
    if fluxlim == 'minmod':
        ind1 = (r > 0) & (abs(r) > 1)
        ind2 = (r > 0) & (abs(r) < 1)
        phi[ind1] = 1.0
        phi[ind2] = r[ind2]
    if fluxlim == 'superbee':
        for i in np.arange(nx):
            a = np.min([1.0, 2.0*r[i]])
            b = np.min([2.0, r[i]])
            phi[i] = np.max([0.0, a, b])
    if fluxlim == 'mc':
        for i in np.arange(nx):
            a = np.min( [ (1.0+r[i])/2.0, 2.0, 2.0*r[i] ] )
            phi[i] = max([0.0, a])
    if fluxlim == 'vanLeer':
        for i in np.arange(nx):
            phi[i] = ( r[i] + abs(r[i]) ) / ( 1.0 + abs(r[i]) )

    # now construct the flux
    if scheme == 'upwind':
        flux[ind_up] = uxi[ind_up] * np.roll(q,1)[ind_up] # note: uxi is the left face of the cell
        flux[ind_um] = uxi[ind_um] * np.roll(q,0)[ind_um] # note: uxi is the left face of the cell
    if scheme == 'centred':
        flux = uxi * (np.roll(q,0)+np.roll(q,1))/2.0
    if scheme == 'downwind':
        flux[ind_up] = uxi[ind_up] * np.roll(q,0)[ind_up]
        flux[ind_um] = uxi[ind_um] * np.roll(q,1)[ind_um]

    # compute flux correction based on flux limiter
    fluxlimiter_correction = 0.5 * abs(uxi) * (1.0 - abs(uxi*dt/dx)) * phi * dq
    flux += fluxlimiter_correction

    # finally, update q
    q[:] = q - dt * (np.roll(flux,-1)-flux) / dx

# === END advect ===


# === START hydro_step_2d ===
def hydro_step_2d(dt, scheme='upwind', fluxlim='mc', boundary='periodic'):
    # Strang splitting
    hydro_step(0.5*dt, sweep='x', scheme=scheme, fluxlim=fluxlim, boundary=boundary)
    hydro_step(0.5*dt, sweep='y', scheme=scheme, fluxlim=fluxlim, boundary=boundary)
    hydro_step(0.5*dt, sweep='y', scheme=scheme, fluxlim=fluxlim, boundary=boundary)
    hydro_step(0.5*dt, sweep='x', scheme=scheme, fluxlim=fluxlim, boundary=boundary)

# === START hydro_step ===
def hydro_step(dt, sweep='x', scheme='upwind', fluxlim='mc', boundary='periodic'):

    # globals
    global uxi, uyi, rho, rhoux, rhouy, rhoetot

    # impose boundary conditions
    update_boundary(boundary=boundary)

    # compute the velocity at the cell interfaces (by averaging)
    if ndim == 1:
        uxi = 0.5 * ( rhoux/rho + np.roll(rhoux,1)/np.roll(rho,1) )
    if ndim == 2:
        uxi = 0.5 * ( rhoux/rho + np.roll(rhoux,1,axis=0)/np.roll(rho,1,axis=0) )
        uyi = 0.5 * ( rhouy/rho + np.roll(rhouy,1,axis=1)/np.roll(rho,1,axis=1) )

    # advect rho
    advect(rho, dt, sweep=sweep, scheme=scheme, fluxlim=fluxlim)

    # advect rho*ux
    if sweep == 'x':
        advect(rhoux, dt, sweep=sweep, scheme=scheme, fluxlim=fluxlim)
    # advect rho*uy
    if sweep == 'y':
        advect(rhouy, dt, sweep=sweep, scheme=scheme, fluxlim=fluxlim)

    # advect rho*etot
    advect(rhoetot, dt, sweep=sweep, scheme=scheme, fluxlim=fluxlim)

    # re-impose boundary conditions after advection
    update_boundary(boundary=boundary)

    # compute pressure (calls equation of state)
    p = pres()

    # Now add the pressure force and energy, for all cells except the guard cells using operator-splitting
    if ndim == 1:
        rhoux   = rhoux   - dt/(2*dx) * (np.roll(p,-1)                -np.roll(p,1)               )
        rhoetot = rhoetot - dt/(2*dx) * (np.roll(p,-1)*np.roll(uxi,-1)-np.roll(p,1)*np.roll(uxi,1))
    if ndim == 2:
        # sweep direction
        if sweep == 'x':
            rhoux   = rhoux   - dt/(2*dx) * (np.roll(p,-1,axis=0)                       -np.roll(p,1,axis=0)                      )
            rhoetot = rhoetot - dt/(2*dx) * (np.roll(p,-1,axis=0)*np.roll(uxi,-1,axis=0)-np.roll(p,1,axis=0)*np.roll(uxi,1,axis=0))
        if sweep == 'y':
            rhouy   = rhouy   - dt/(2*dy) * (np.roll(p,-1,axis=1)                       -np.roll(p,1,axis=1)                      )
            rhoetot = rhoetot - dt/(2*dy) * (np.roll(p,-1,axis=1)*np.roll(uyi,-1,axis=1)-np.roll(p,1,axis=1)*np.roll(uyi,1,axis=1))

    # re-impose boundary conditions after source terms
    update_boundary(boundary=boundary)

# === END hydro_step ===


# === START update_boundary (at the moment only works for 2 guard cells) ===
def update_boundary(boundary="periodic"):
    global rho, rhoux, rhouy, rhoetot
    if nguard != 2:
        print("Error: update_boundary assumes nguard=2.")
        exit()
    if boundary == 'periodic':
        if ndim == 1:
            rho    [0]    = rho    [nx-4]
            rho    [1]    = rho    [nx-3]
            rho    [nx-2] = rho    [2]
            rho    [nx-1] = rho    [3]
            rhoux  [0]    = rhoux  [nx-4]
            rhoux  [1]    = rhoux  [nx-3]
            rhoux  [nx-2] = rhoux  [2]
            rhoux  [nx-1] = rhoux  [3]
            rhoetot[0]    = rhoetot[nx-4]
            rhoetot[1]    = rhoetot[nx-3]
            rhoetot[nx-2] = rhoetot[2]
            rhoetot[nx-1] = rhoetot[3]
        if ndim == 2:
            # x boundaries
            rho    [0   ,:] = rho    [nx-4,:]
            rho    [1   ,:] = rho    [nx-3,:]
            rho    [nx-2,:] = rho    [2   ,:]
            rho    [nx-1,:] = rho    [3   ,:]
            rhoux  [0   ,:] = rhoux  [nx-4,:]
            rhoux  [1   ,:] = rhoux  [nx-3,:]
            rhoux  [nx-2,:] = rhoux  [2   ,:]
            rhoux  [nx-1,:] = rhoux  [3   ,:]
            rhouy  [0   ,:] = rhouy  [nx-4,:]
            rhouy  [1   ,:] = rhouy  [nx-3,:]
            rhouy  [nx-2,:] = rhouy  [2   ,:]
            rhouy  [nx-1,:] = rhouy  [3   ,:]
            rhoetot[0   ,:] = rhoetot[nx-4,:]
            rhoetot[1   ,:] = rhoetot[nx-3,:]
            rhoetot[nx-2,:] = rhoetot[2   ,:]
            rhoetot[nx-1,:] = rhoetot[3   ,:]
            # y boundaries
            rho    [:,   0] = rho    [:,ny-4]
            rho    [:,   1] = rho    [:,ny-3]
            rho    [:,ny-2] = rho    [:,   2]
            rho    [:,ny-1] = rho    [:,   3]
            rhoux  [:,   0] = rhoux  [:,ny-4]
            rhoux  [:,   1] = rhoux  [:,ny-3]
            rhoux  [:,ny-2] = rhoux  [:,   2]
            rhoux  [:,ny-1] = rhoux  [:,   3]
            rhouy  [:,   0] = rhouy  [:,ny-4]
            rhouy  [:,   1] = rhouy  [:,ny-3]
            rhouy  [:,ny-2] = rhouy  [:,   2]
            rhouy  [:,ny-1] = rhouy  [:,   3]
            rhoetot[:,   0] = rhoetot[:,ny-4]
            rhoetot[:,   1] = rhoetot[:,ny-3]
            rhoetot[:,ny-2] = rhoetot[:,   2]
            rhoetot[:,ny-1] = rhoetot[:,   3]
    if boundary == 'outflow':
        if ndim == 1:
            rho    [0]    =  rho    [2]
            rho    [1]    =  rho    [2]
            rho    [nx-2] =  rho    [nx-3]
            rho    [nx-1] =  rho    [nx-3]
            rhoux  [0]    =  rhoux  [2]
            rhoux  [1]    =  rhoux  [2]
            rhoux  [nx-2] =  rhoux  [nx-3]
            rhoux  [nx-1] =  rhoux  [nx-3]
            rhoetot[0]    =  rhoetot[2]
            rhoetot[1]    =  rhoetot[2]
            rhoetot[nx-2] =  rhoetot[nx-3]
            rhoetot[nx-1] =  rhoetot[nx-3]
        if ndim == 2:
            # x boundaries
            rho    [0   ,:] =  rho    [2   ,:]
            rho    [1   ,:] =  rho    [2   ,:]
            rho    [nx-2,:] =  rho    [nx-3,:]
            rho    [nx-1,:] =  rho    [nx-3,:]
            rhoux  [0   ,:] =  rhoux  [2   ,:]
            rhoux  [1   ,:] =  rhoux  [2   ,:]
            rhoux  [nx-2,:] =  rhoux  [nx-3,:]
            rhoux  [nx-1,:] =  rhoux  [nx-3,:]
            rhouy  [0   ,:] =  rhouy  [2   ,:]
            rhouy  [1   ,:] =  rhouy  [2   ,:]
            rhouy  [nx-2,:] =  rhouy  [nx-3,:]
            rhouy  [nx-1,:] =  rhouy  [nx-3,:]
            rhoetot[0   ,:] =  rhoetot[2   ,:]
            rhoetot[1   ,:] =  rhoetot[2   ,:]
            rhoetot[nx-2,:] =  rhoetot[nx-3,:]
            rhoetot[nx-1,:] =  rhoetot[nx-3,:]
            # y boundaries
            rho    [:,   0] =  rho    [:,   2]
            rho    [:,   1] =  rho    [:,   2]
            rho    [:,ny-2] =  rho    [:,ny-3]
            rho    [:,ny-1] =  rho    [:,ny-3]
            rhoux  [:,   0] =  rhoux  [:,   2]
            rhoux  [:,   1] =  rhoux  [:,   2]
            rhoux  [:,ny-2] =  rhoux  [:,ny-3]
            rhoux  [:,ny-1] =  rhoux  [:,ny-3]
            rhouy  [:,   0] =  rhouy  [:,   2]
            rhouy  [:,   1] =  rhouy  [:,   2]
            rhouy  [:,ny-2] =  rhouy  [:,ny-3]
            rhouy  [:,ny-1] =  rhouy  [:,ny-3]
            rhoetot[:,   0] =  rhoetot[:,   2]
            rhoetot[:,   1] =  rhoetot[:,   2]
            rhoetot[:,ny-2] =  rhoetot[:,ny-3]
            rhoetot[:,ny-1] =  rhoetot[:,ny-3]
# === END update_boundary ===


# === START update_plot ===
def update_plot(cache=None, q=None, xr=None, yr=None, dr=None, sleep_sec=2e-2, dpi=100):
    # default quantity for plotting
    if q is None:
        q = rho
    # setup figure (initial call)
    if cache is None:
        # switch backend, as MacOSX does not support blitting
        matplotlib.use('TkAgg')
        # set up ticks and stuff
        rcParams['xtick.top'] = True
        rcParams['xtick.direction'] = 'in'
        rcParams['xtick.minor.visible'] = True
        rcParams['ytick.right'] = True
        rcParams['ytick.direction'] = 'in'
        rcParams['ytick.minor.visible'] = True
        # create figure
        if ndim == 1:
            fig = plt.figure(figsize=(10.0,6.0), dpi=dpi)
        if ndim == 2:
            fig = plt.figure(figsize=(10.0,8.0), dpi=dpi)
        # get current axes
        ax = plt.gca()
        # set axes ranges
        if xr is None: xr = [x[nguard]-dx/2, x[nx-nguard]-dx/2]
        if ndim == 2:
            if yr is None: yr = [y[nguard]-dy/2, y[ny-nguard]-dy/2]
        if dr is None: dr = [np.min(q), np.max(q)]
        ax.set_xlim(xr)
        if ndim == 1:
            ax.set_ylim(dr)
        # add axis labels
        plt.xlabel('x')
        if ndim == 1:
            plt.ylabel('q')
        if ndim == 2:
            plt.ylabel('y')
        # create initial plot and store animated plot in ln
        if ndim == 1:
            (ln,) = ax.plot(x, q, 'b.-', linewidth=0.5, markersize=1.0, animated=True)
        if ndim == 2:
            extent = [xr[0], xr[1], yr[0], yr[1]]
            ln = plt.imshow(q.T, cmap='plasma', origin='lower', interpolation='none', extent=extent)
            cb = plt.colorbar(label='q', pad=0.01, aspect=25)
            #if self.log is False: cb.ax.minorticks_on()
            cb.ax.yaxis.set_offset_position('left')
        
        # show and allow some time to draw the figure
        plt.show(block=False)
        #import ipdb; ipdb.set_trace()
        plt.pause(0.1)
        # get copy of entire figure (everything inside fig.bbox), but excluding animated artist
        bg = fig.canvas.copy_from_bbox(fig.bbox)
        # draw the animated artist, this uses a cached renderer
        ax.draw_artist(ln)
        # push updated RGBA buffer from the renderer to the GUI framework
        fig.canvas.blit(fig.bbox)
        # return those objects that we'll feed back in to reuse below (after init)
        cache = fig, ax, bg, ln
        return cache
    else: # we are redrawing with updated figure
        # restore from cache
        fig, ax, bg, ln = cache
        # restore background
        fig.canvas.restore_region(bg)
        # draw new y data
        if ndim == 1:
            ln.set_ydata(q)
        if ndim == 2:
            ln.set_data(q.T)
        # re-render the artist, updating the canvas state, but not the screen
        ax.draw_artist(ln)
        # copy the image to the GUI state
        fig.canvas.blit(fig.bbox)
        # flush any pending GUI events, re-painting the screen if needed
        fig.canvas.flush_events()
        # sleep a bit
        sleep(sleep_sec)
# === END update_plot ===


# ================== START advection_test ===================
def advection_test(scheme='upwind', fluxlim=None, boundary=None, cfl=0.8):
    print("Start advection_test...")
    if fluxlim is None: fluxlim = 'donor-cell'
    if boundary is None: boundary = 'periodic'
    # globals
    global ndim, gamma, rho, rhoux, rhouy, rhoetot, uxi, uyi
    # initialise grid and global variables
    init_grid(n=500, xr=[0.0,1.0], boundary=boundary)
    # set initial conditions
    gamma = 5.0/3.0
    if ndim == 1:
        rho[:] = 1.0 + np.exp(-(x-0.75)**2/0.1**2)
        index = (x >= 0.25-0.05) & (x <= 0.25+0.05)
        rho[index] = 2.0
    if ndim == 2:
        for j in range(0,ny):
            rho[:,j] = 1.0 + np.exp(-(x-0.75)**2/0.1**2)
            index = (x >= 0.25-0.05) & (x <= 0.25+0.05)
            rho[index,j] = 2.0
        for i in range(0,nx):
            index = (y >= 0.25-0.05) & (y <= 0.25+0.05)
            rho[i,index] = 2.0
    rhoux = rho # set for constant velocity
    if ndim == 2:
        rhouy = 0.5*rho # set for constant velocity
    rhoetot = 0.5*rhoux**2/rho
    # plot initial conditions
    d0 = 0.0
    d1 = 3.0
    cache = update_plot(dr=[d0,d1])
    # set number of time steps and start evolution loop
    input("Press enter to start evolution loop...")
    sleep(1.0)
    nt = 300
    time = 0.0
    for step in range(0, nt+1):
        # boundary condition
        update_boundary(boundary=boundary)
        # replot
        update_plot(cache=cache)
        # compute the velocity at the cell interfaces (by averaging)
        uxi = 0.5 * ( rhoux/rho + np.roll(rhoux,1,axis=0)/np.roll(rho,1,axis=0) )
        if ndim == 2:
            uyi = 0.5 * ( rhouy/rho + np.roll(rhouy,1,axis=1)/np.roll(rho,1,axis=1) )
        # compute timestep based on CFL
        dt = cfl * dx / np.max(abs(uxi))
        if ndim == 2:
            dt = np.min( [ dt, cfl * dy / np.max(abs(uyi)) ] )
        # check conservation
        #totmass = check_conservation(rho, xi, nguard)
        # do advection
        time += dt
        print('step={:6d}'.format(step), ' |  time={:10.6e}'.format(time), ' |  dt={:10.6e}'.format(dt))

        advect(rho, dt, sweep='x', scheme=scheme, fluxlim=fluxlim)
        if ndim == 2:
            advect(rho, dt, sweep='y', scheme=scheme, fluxlim=fluxlim)

        rhoux = rho # set for constant velocity!
        if ndim == 2:
            rhouy = 0.5*rho # set for constant velocity

# ================== END advection_test ===================

# ================== START hydro_test ===================
def hydro_test(scheme='upwind', fluxlim=None, boundary=None, cfl=0.8):
    print("Start hydro_test...")
    if fluxlim is None: fluxlim = 'vanleer'
    if boundary is None: boundary = 'periodic'
    # globals
    global ndim, gamma, rho, rhoux, rhouy, rhoetot
    # initialise grid and global variables
    init_grid(n=250, xr=[0.0,1.0], boundary=boundary)
    # set initial conditions
    gamma = 5.0/3.0
    if ndim == 1:
        rho[:] = 1.0 + np.exp(-(x-0.75)**2/0.1**2)
        index = (x >= 0.25-0.05) & (x <= 0.25+0.05)
        rho[index] = 2.0
    if ndim == 2:
        for j in range(0,ny):
            rho[:,j] = 1.0 + np.exp(-(x-0.75)**2/0.1**2)
            index = (x >= 0.25-0.05) & (x <= 0.25+0.05)
            rho[index,j] = 2.0
        for i in range(0,nx):
            index = (y >= 0.25-0.05) & (y <= 0.25+0.05)
            rho[i,index] = 2.0
    rhoux = 0.0*rho # set for constant velocity
    rhoetot = 0.5*rhoux**2/rho + 1.0*rho/(gamma-1) # eint = c_s^2 / (gamma-1), so this is c_s = 1
    if ndim == 2:
        rhouy = 0.0*rho # set for constant velocity
        rhoetot = 0.5*(rhoux**2+rhouy**2)/rho + 1.0*rho/(gamma-1) # eint = c_s^2 / (gamma-1), so this is c_s = 1
    # plot initial conditions
    d0 = 0.0
    d1 = 3.0
    cache = update_plot(dr=[d0,d1])
    # set number of time steps and start evolution loop
    input("Press enter to start evolution loop...")
    sleep(1.0)
    nt = 500
    time = 0.0
    for step in range(0, nt+1):
        # boundary condition
        update_boundary(boundary=boundary)
        # replot
        update_plot(cache=cache)
        # compute timestep based on CFL
        dt = compute_dt(cfl)
        # check conservation
        #totmass = check_conservation(rho, xi, nguard)
        # advance hydro
        time += dt
        print('step={:6d}'.format(step), ' |  time={:10.6e}'.format(time), ' |  dt={:10.6e}'.format(dt))
        if ndim == 1:
            hydro_step(dt, scheme=scheme, fluxlim=fluxlim)
        if ndim == 2:
            hydro_step_2d(dt, scheme=scheme, fluxlim=fluxlim)

# ================== END hydro_test ===================

# ================== START kh_test ===================
def kh_test(scheme='upwind', fluxlim=None, boundary=None, cfl=0.8):
    print("Start kh_test...")
    if fluxlim is None: fluxlim = 'vanleer'
    if boundary is None: boundary = 'periodic'
    # globals
    global ndim, gamma, rho, rhoux, rhouy, rhoetot
    # initialise grid and global variables
    ndim = 2 # Kelvin-Helmholtz is always 2D
    init_grid(n=250, xr=[0.0,1.0], boundary=boundary)
    # set initial conditions
    gamma = 5.0/3.0
    rhouy[:,:] = 0.0
    for j in range(0,ny):
        for i in range(0,nx):
            if y[j] < 0.5 + 0.05*np.sin(2*np.pi*x[i]): # lower half
                rho  [i,j] = 1.0
                rhoux[i,j] = +0.1 * rho[i,j]
            else: # upper half
                rho  [i,j] = 1.0
                rhoux[i,j] = -0.1 * rho[i,j]
    pres = 1.0/gamma # c_s = 1 at rho = 1
    rhoetot = 0.5*(rhoux**2+rhouy**2)/rho + pres/(gamma-1) # eint = pres / rho / (gamma-1)
    # plot initial conditions
    d0 = 0.0
    d1 = 3.0
    cache = update_plot(dr=[d0,d1], q=rhoux)
    # set number of time steps and start evolution loop
    input("Press enter to start evolution loop...")
    sleep(1.0)
    nt = 500
    time = 0.0
    for step in range(0, nt+1):
        # boundary condition
        update_boundary(boundary=boundary)
        # replot
        update_plot(cache=cache, q=rhoux)
        # compute timestep based on CFL
        dt = compute_dt(cfl)
        # check conservation
        #totmass = check_conservation(rho, xi, nguard)
        # advance hydro
        time += dt
        print('step={:6d}'.format(step), ' |  time={:10.6e}'.format(time), ' |  dt={:10.6e}'.format(dt))
        hydro_step_2d(dt, scheme=scheme, fluxlim=fluxlim)

# ================== END kh_test ===================


# ================== START rt_test ===================
def rt_test(scheme='upwind', fluxlim=None, boundary=None, cfl=0.8):
    print("Start rt_test...")
    if fluxlim is None: fluxlim = 'vanleer'
    if boundary is None: boundary = 'outflow'
    # globals
    global ndim, gamma, rho, rhoux, rhouy, rhoetot
    # initialise grid and global variables
    ndim = 2 # Rayleigh-Taylor is always 2D
    init_grid(n=250, xr=[0.0,1.0], boundary=boundary)
    # set initial conditions
    gamma = 5.0/3.0
    rhoux[:,:] = 0.0
    rhouy[:,:] = 0.0
    for j in range(0,ny):
        for i in range(0,nx):
            if y[j] < 0.8 - 0.1*np.sin(np.pi*x[i]): # lower half
                rho  [i,j] = 1.
            else: # upper half
                rho  [i,j] = 10.
    pres = 1.0/gamma # c_s = 1 at rho = 1
    rhoetot = 0.5*(rhoux**2+rhouy**2)/rho + pres/(gamma-1) # eint = pres / rho / (gamma-1)
    # plot initial conditions
    d0 = 0.0
    d1 = 3.0
    cache = update_plot(dr=[d0,d1], q=rho)
    # set number of time steps and start evolution loop
    #input("Press enter to start evolution loop...")
    sleep(1.0)
    nt = 500
    time = 0.0
    for step in range(0, nt+1):
        # boundary condition
        update_boundary(boundary=boundary)
        # replot
        update_plot(cache=cache, q=rho)
        # compute timestep based on CFL
        dt = compute_dt(cfl)
        # check conservation
        #totmass = check_conservation(rho, xi, nguard)
        # advance hydro
        time += dt
        print('step={:6d}'.format(step), ' |  time={:10.6e}'.format(time), ' |  dt={:10.6e}'.format(dt))
        hydro_step_2d(dt, scheme=scheme, fluxlim=fluxlim)
        # add gravity
        update_boundary(boundary=boundary)
        ei = eint()
        g = 100.0
        #rhouy = -rho * g * dt
        rhoetot = 0.5*(rhoux**2+rhouy**2)/rho + rho*ei

# ================== END rt_test ===================


# ===== the following applies in case we are running this in script mode =====
if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Hydrodynamics code.')
    parser.add_argument("-simulation", "--simulation", type=str, choices=['advection_test', 'hydro_test', 'kh_test', 'rt_test'], default='advection_test', help="The simulation setup/test.")
    parser.add_argument("-scheme", "--scheme", type=str, choices=['upwind', 'centred', 'downwind'], default='upwind', help="Finite difference scheme for advection.")
    parser.add_argument("-fluxlim", "--fluxlim", type=str, choices=['donor-cell', 'lax-wendroff', 'beam-warming', 'fromm', 'minmod', 'superbee', 'mc', 'vanLeer'], help="Flux limiter.")
    parser.add_argument("-boundary", "--boundary", type=str, choices=['periodic', 'outflow'], help="Boundary condition.")
    parser.add_argument("-cfl", "--cfl", type=float, default=0.8, help="CFL timestep safety factor.")
    parser.add_argument("-ndim", "--ndim", type=int, choices=[1,2], default=1, help="1D or 2D setup.")
    args = parser.parse_args()

    # set 1D or 2D globally (may be overwritten by specific setup requirement)
    ndim = args.ndim

    # advection_test
    if args.simulation == "advection_test":
        advection_test(scheme=args.scheme, fluxlim=args.fluxlim, boundary=args.boundary, cfl=args.cfl)
    # hydro_test
    if args.simulation == "hydro_test":
        hydro_test(scheme=args.scheme, fluxlim=args.fluxlim, boundary=args.boundary, cfl=args.cfl)
    # kh_test (Kelvin-Helmholtz instability)
    if args.simulation == "kh_test":
        kh_test(scheme=args.scheme, fluxlim=args.fluxlim, boundary=args.boundary, cfl=args.cfl)
    # rt_test (Rayleigh-Taylor instability)
    if args.simulation == "rt_test":
        rt_test(scheme=args.scheme, fluxlim=args.fluxlim, boundary=args.boundary, cfl=args.cfl)

    import ipdb; ipdb.set_trace()