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

import argparse
from time import sleep
import numpy as np
import matplotlib
from matplotlib import rcParams
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from cfpack import get_coords, eform, print, stop
class hydro_class:

    def __init__(self):
        self.ndim     = 1 # the dimensionality (1,2)
        self.ncells   = np.array([500, 500]) # number of cells (2D)
        self.nguard   = 2 # number of guard cells
        self.scheme   = "upwind" # upwind, downwind, centered
        self.fluxlim  = "minmod" # flux limiter
        self.boundary = "periodic" # boundary condition
        self.gamma    = 5/3 # adiabatic index
        self.cfl      = 0.8 # CFL safety factor
        self.nend     = 1000 # number of timesteps
        self.tmax     = 1e99 # end time
        self.dpi      = 90 # DPI for plotting
        self.verbose  = 1 # verbose level
        self.dx       = 1.0 # cell size x
        self.dy       = 1.0 # cell size y
        self.x        = None # cell-centred coordinates x
        self.y        = None # cell-centred coordinates y
        self.rho      = None # mass density
        self.rhoux    = None # momentum density x
        self.rhouy    = None # momentum density y
        self.rhoetot  = None # energy density
        self.uxi      = None # velocity (face-centred) x
        self.uyi      = None # velocity (face-centred) y

    # === START init_grid ===
    def init_grid(self, xr=[0.0,1.0], yr=[0.0,1.0]):
        # total number of cell centers (including guard cells)
        self.nx = self.ncells[0] + 2*self.nguard
        self.ny = 1
        if self.ndim == 2:
            self.ny = self.ncells[1] + 2*self.nguard
        # active cell indices
        self.ix = np.arange(self.nguard, self.nx-self.nguard)
        self.iy = 0
        if self.ndim == 2:
            self.iy = np.arange(self.nguard, self.ny-self.nguard)
        # cell width
        self.dx = (xr[1]-xr[0]) / (self.nx-2*self.nguard)
        self.dy = 1.0
        if self.ndim == 2:
            self.dy = (yr[1]-yr[0]) / (self.ny-2*self.nguard)
        # coordinates at cell centers
        self.x, self.y = get_coords([xr[0]-self.nguard*self.dx,yr[0]-self.nguard*self.dy],
                                    [xr[1]+self.nguard*self.dx,yr[1]+self.nguard*self.dy],
                                    [self.nx,self.ny], cell_centred=True)
        # define containers to store density, momentum, energy, etc.
        zeros = np.zeros([self.nx,self.ny])
        self.rho = np.copy(zeros) # mass density
        self.rhoux = np.copy(zeros) # momentum density x
        self.rhouy = np.copy(zeros) # momentum density y
        self.rhoetot = np.copy(zeros) # energy density
        self.uxi = np.copy(zeros) # used for velocity at interfaces x
        self.uyi = np.copy(zeros) # used for velocity at interfaces y
    # === END init_grid ===

    # === START parameter_info ===
    def parameter_info(self):
        print("number of dimensions    = ", self.ndim)
        print("number of cells         = ", self.ncells[0:self.ndim])
        print("scheme                  = ", self.scheme)
        print("flux limiter            = ", self.fluxlim)
        print("boundary condition      = ", self.boundary)
        print("CFL safety factor       = ", self.cfl)
        print("adiabtic gamma          = ", self.gamma)
        print("max number of timesteps = ", self.nend)
        print("max simulation time     = ", self.tmax)
    # === END parameter_info ===

    def get_integral_quantities(self, quiet=False):
        dV = self.dx * self.dy # cell volume element
        ix = [self.ix[0], self.ix[-1]+1] # active indices in x (excluding guard cells)
        if self.ndim == 2:
            iy = [self.iy[0], self.iy[-1]+1] # active indices in y (excluding guard cells)
        else:
            iy = [0,1]
        mass = np.sum(self.rho              [ix[0]:ix[1], iy[0]:iy[1]]) * dV
        momx = np.sum(self.rhoux            [ix[0]:ix[1], iy[0]:iy[1]]) * dV
        momy = np.sum(self.rhouy            [ix[0]:ix[1], iy[0]:iy[1]]) * dV
        etot = np.sum(self.rhoetot          [ix[0]:ix[1], iy[0]:iy[1]]) * dV
        eint = np.sum((self.rho*self.eint())[ix[0]:ix[1], iy[0]:iy[1]]) * dV
        ekin = etot - eint
        vals = [mass, momx, momy, etot, eint, ekin]
        if not quiet: print("mass, momx, momy, etot, eint, ekin = "+", ".join(["{:+.3e}".format(x) for x in vals]))
        return vals

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

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

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

    # === START advect ===
    def advect(self, q, dt, sweep='x'):

        if self.verbose > 1: print("before advect: q = ", q)

        # sweep direction
        if sweep == 'x':
            axis = 0
            ui = self.uxi
            delta = self.dx
        if sweep == 'y':
            axis = 1
            ui = self.uyi
            delta = self.dy
        # positive and negative velocity indices
        ind_up = ui >= 0
        ind_um = ui <  0

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

        # now construct the flux
        flux = np.zeros([self.nx,self.ny])
        if self.scheme == 'upwind':
            flux[ind_up] = ui[ind_up] * np.roll(q,1,axis=axis)[ind_up] # ui is at the left face of the cell
            flux[ind_um] = ui[ind_um] * np.roll(q,0,axis=axis)[ind_um] # ui is at the left face of the cell
        if self.scheme == 'centred':
            flux = ui * (np.roll(q,0,axis=axis)+np.roll(q,1,axis=axis))/2.0
        if self.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]

        # compute flux correction based on flux limiter
        fluxlimiter_correction = 0.5 * abs(ui) * (1.0 - abs(ui*dt/delta)) * phi * dq
        flux += fluxlimiter_correction
        if self.verbose > 1: print("fluxlimiter_correction = ", fluxlimiter_correction)

        # finally, update q
        q += - dt * (np.roll(flux,-1,axis=axis)-flux) / delta
        if self.verbose > 1: print("after advect: q = ", q)

    # === END advect ===

    # === START hydro_step_2d ===
    def hydro_step_2d(self, dt):
        # Strang splitting
        self.hydro_step(0.5*dt, sweep='x')
        self.hydro_step(0.5*dt, sweep='y')
        self.hydro_step(0.5*dt, sweep='y')
        self.hydro_step(0.5*dt, sweep='x')


    # === START hydro_step ===
    def hydro_step(self, dt, sweep='x'):

        if self.verbose > 1: print("START of hydro_step")

        # sweep direction dependent (compute velocity at cell interfaces, by averaging)
        if sweep == 'x':
            axis = 0
            uc = self.rhoux/self.rho # cell-centered velx
            ui = 0.5 * (uc + np.roll(uc,1,axis=axis)) # face-centered velx (index 0 is left face)
            self.uxi = ui # store in class (for access in advect)
            delta = self.dx
        if sweep == 'y':
            axis = 1
            uc = self.rhouy/self.rho # cell-centered vely
            ui = 0.5 * (uc + np.roll(uc,1,axis=axis)) # face-centered vely (index 0 is left face)
            self.uyi = ui # store in class (for access in advect)
            delta = self.dy

        # impose boundary conditions
        self.update_boundary()

        # advect rho
        self.advect(self.rho, dt, sweep=sweep)
        if self.verbose > 1: print("rho (after advect) = ", self.rho)
        # advect rho*ux and rho*uy
        self.advect(self.rhoux, dt, sweep=sweep)
        if self.verbose > 1: print("self.rhoux (after advect) = ", self.rhoux)
        self.advect(self.rhouy, dt, sweep=sweep)
        if self.verbose > 1: print("self.rhouy (after advect) = ", self.rhouy)
        # advect rho*etot
        self.advect(self.rhoetot, dt, sweep=sweep)
        if self.verbose > 1: print("self.rhoetot (after advect) = ", self.rhoetot)

        # re-impose boundary conditions after advection
        self.update_boundary()

        # compute pressure (calls equation of state)
        p = self.pres()
        if self.verbose > 1: print("p = ", p)

        # Now add the pressure force and energy, using operator-splitting
        if True: # somewhat more diffusive, but more stable
            pl = np.roll(p,+1,axis=axis) #  left cell-centered pres
            pr = np.roll(p,-1,axis=axis) # right cell-centered pres
            ul = 0.5 * (np.roll(ui, 0,axis=axis) + np.roll(ui,+1,axis=axis)) #  left cell-centered vel (ui is face-centred)
            ur = 0.5 * (np.roll(ui,-1,axis=axis) + np.roll(ui,-2,axis=axis)) # right cell-centered vel (ui is face-centred)
            if sweep == 'x': self.rhoux   += - dt/(2*delta) * (pr - pl) # (right-left) cell-centered pres
            if sweep == 'y': self.rhouy   += - dt/(2*delta) * (pr - pl) # (right-left) cell-centered pres
            self.rhoetot += - dt/(2*delta) * (pr*ur - pl*ul)
        else: # somewhat less diffusive, but less stable
            pl = 0.5*(p + np.roll(p,+1,axis=axis)) #  left face-centered pres
            pr = 0.5*(p + np.roll(p,-1,axis=axis)) # right face-centered pres
            ul = np.roll(ui, 0,axis=axis) #  left face-centered vel (ui is face-centred)
            ur = np.roll(ui,-1,axis=axis) # right face-centered vel (ui is face-centred)
            if sweep == 'x': self.rhoux   += - dt/(delta) * (pr - pl) # (right-left) face-centered pres
            if sweep == 'y': self.rhouy   += - dt/(delta) * (pr - pl) # (right-left) face-centered pres
            self.rhoetot += - dt/(delta) * (pr*ur - pl*ul)

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

        if self.verbose > 1:
            print("END of hydro_step")
            stop()

    # === END hydro_step ===

    # === START update_boundary ===
    def update_boundary(self):
        # loop over quantities that need guard-cell filling
        for q in [self.rho, self.rhoux, self.rhouy, self.rhoetot]:
            if self.boundary == 'periodic':
                # x boundaries
                for n in range(self.nguard):
                    q[n,:] = q[self.nx-2*self.nguard+n,:] # x lower boundary
                    q[self.nx-self.nguard+n,:] = q[self.nguard+n,:] # x upper boundary
                if self.ndim == 2:
                    # y boundaries
                    for n in range(self.nguard):
                        q[:,n] = q[:,self.ny-2*self.nguard+n] # y lower boundary
                        q[:,self.ny-self.nguard+n] = q[:,self.nguard+n] # y upper boundary
            if self.boundary == 'outflow':
                # x boundaries
                for n in range(self.nguard):
                    q[n,:] = q[self.nguard,:] # x lower boundary
                    q[self.nx-1-n,:] = q[self.nx-1-self.nguard,:] # x upper boundary
                if self.ndim == 2:
                    # y boundaries
                    for n in range(self.nguard):
                        q[:,n] = q[:,self.nguard] # y lower boundary
                        q[:,self.ny-1-n] = q[:,self.ny-1-self.nguard] # y upper boundary
    # === END update_boundary ===

    # === START update_plot ===
    def update_plot(self, cache=None, vars=None, xr=None, yr=None, dr=None, sleep_sec=2e-2):
        def prep_quantities(vars):
            q = [] # construct quantities to plot
            for var in vars:
                if var == "dens": q.append(self.rho)
                if var == "velx": q.append(self.rhoux/self.rho)
                if var == "vely": q.append(self.rhouy/self.rho)
                if var == "momx": q.append(self.rhoux)
                if var == "momy": q.append(self.rhouy)
                if var == "pres": q.append(self.pres())
                if var == "eint": q.append(self.eint())
                if var == "ener": q.append(self.rhoetot/self.rho)
            return q
        def redraw():
            # restore from cache
            fig, ax, bg, pdat, vars = cache
            q = prep_quantities(vars)
            # restore background
            fig.canvas.restore_region(bg)
            # for each q to plot
            for iq in np.arange(len(q)):
                # draw new y data
                if self.ndim == 1:
                    pdat[iq].set_ydata(q[iq])
                if self.ndim == 2:
                    pdat[iq].set_data(q[iq].T)
                # re-render the artist, updating the canvas state, but not the screen
                ax[iq].draw_artist(pdat[iq])
            # copy the image to the GUI state
            fig.canvas.blit(fig.bbox)
            # flush aself.ny pending GUI events, re-painting the screen if needed
            fig.canvas.flush_events()
            # sleep a bit
            sleep(sleep_sec)
        # setup figure (initial call)
        if cache is None:
            # print integral quantities
            hy.get_integral_quantities()
            # prep plot variables
            if vars is None: vars = "dens" # default variable for plotting
            if not isinstance(vars, list): vars = [vars] # turn into list
            q = prep_quantities(vars)
            nq = len(q) # number of variables to plot
            matplotlib.use('TkAgg') # switch backend, as MacOSX does not support blitting
            # 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 self.ndim == 1:
                figsize = (8.0,nq*2.5)
                if nq == 1: figsize = (8.0,5.0)
                fig, ax = plt.subplots(nrows=nq, ncols=1, dpi=self.dpi, figsize=figsize)
            if self.ndim == 2:
                fig, ax = plt.subplots(nrows=1, ncols=nq, dpi=self.dpi, figsize=(nq*5.0,6.0))
            if nq == 1: ax = [ax]
            # prep x (and y) range
            if xr is None: xr = [self.x[self.nguard,0]-self.dx/2, self.x[self.nx-self.nguard,0]-self.dx/2]
            if self.ndim == 2:
                if yr is None: yr = [self.y[0,self.nguard]-self.dy/2, self.y[0,self.ny-self.nguard]-self.dy/2]
            # prep data range(s)
            if dr is None:
                dr = []
                for iq in np.arange(nq):
                    dr.append([np.min(q[iq]), np.max(q[iq])])
            # for each q to plot
            pdat = []
            for iq in np.arange(nq):
                # set axes ranges
                ax[iq].set_xlim(xr)
                if self.ndim == 1:
                    ax[iq].set_ylim(dr[iq])
                # add axis labels
                ax[iq].set_xlabel(r'$x$')
                if self.ndim == 1:
                    ax[iq].set_ylabel(vars[iq])
                if self.ndim == 2:
                    ax[iq].set_ylabel(r'$y$')
                # create initial plot and store animated plot in pdat
                if self.ndim == 1:
                    (pd,) = ax[iq].plot(self.x, q[iq], 'b.-', linewidth=0.5, markersize=1.0, animated=True)
                if self.ndim == 2:
                    extent = [xr[0], xr[1], yr[0], yr[1]]
                    pd = ax[iq].imshow(q[iq].T, cmap='plasma', origin='lower', interpolation='none', extent=extent, vmin=dr[iq][0], vmax=dr[iq][1])
                    divider = make_axes_locatable(ax[iq])
                    cax = divider.append_axes('right', size='5%', pad=0.05)
                    cb = fig.colorbar(pd, cax=cax, label=vars[iq])
                    cb.ax.yaxis.set_offset_position('left')
                pdat.append(pd)
            # show and allow some time to draw the figure
            fig.tight_layout()
            plt.show(block=False)
            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)
            # for each q to plot
            for iq in np.arange(nq):
                # draw the animated artist, this uses a cached renderer
                ax[iq].draw_artist(pdat[iq])
            # 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, pdat, vars
            redraw()
            redraw()
            return cache
        else: # we are redrawing with updated figure
            redraw()
    # === END update_plot ===

    # === START evolve ===
    def evolve(self, cache=None, gravity=None):
        input("Press enter to start evolution loop...")
        sleep(1.0)
        time = 0.0
        for step in range(0, self.nend+1):
            # boundary condition
            self.update_boundary()
            # replot
            if cache: self.update_plot(cache=cache)
            # compute timestep based on CFL
            dt = self.compute_dt()
            # print integral quantities
            hy.get_integral_quantities()
            # advance hydro
            time += dt
            print('step={:6d}'.format(step), ' |  time={:10.6e}'.format(time), ' |  dt={:10.6e}'.format(dt))
            if self.ndim == 1:
                self.hydro_step(dt)
            if self.ndim == 2:
                self.hydro_step_2d(dt)
            # add gravity
            if gravity is not None:
                eint = self.eint()
                self.rhouy += -self.rho * gravity * dt
                self.rhoetot = 0.5*(self.rhoux**2+self.rhouy**2)/self.rho + self.rho*eint
                self.update_boundary()
            # return if reached max time
            if time > self.tmax:
                input("Reached max simulation time. Press enter to exit.")
                return
        input("Reached max number of time steps. Press enter to exit.")
    # === END evolve ===


# ================== START advect_test ===================
def advect_test(hy):
    print("Start advect_test...")
    # initialise grid
    hy.init_grid(xr=[0.0, 1.0])
    hy.parameter_info()
    # set initial conditions
    vx = 1.0
    vy = 0.5
    if hy.ndim == 1:
        hy.rho = 1.0 + np.exp(-(hy.x-0.75)**2/0.1**2)
        index = (hy.x >= 0.25-0.05) & (hy.x <= 0.25+0.05)
        hy.rho[index] = 2.0
    if hy.ndim == 2:
        hy.rho = 1.0 + np.exp(-(hy.x-0.75)**2/0.1**2)
        index = (hy.x >= 0.25-0.05) & (hy.x <= 0.25+0.05) & (hy.y >= 0.25-0.05) & (hy.y <= 0.25+0.05)
        hy.rho[index] = 2.0
    hy.rhoux = hy.rho * vx # set for constant velocity
    if hy.ndim == 2:
        hy.rhouy = hy.rho * vy # set for constant velocity
    hy.rhoetot = 0.5*(hy.rhoux**2+hy.rhouy**2)/hy.rho
    # plot initial conditions
    d0 = 0.0
    d1 = 3.0
    cache = hy.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
        hy.update_boundary()
        # replot
        hy.update_plot(cache=cache)
        # compute the velocity at the cell interfaces (by averaging)
        hy.uxi = 0.5 * ( hy.rhoux/hy.rho + np.roll(hy.rhoux,1,axis=0)/np.roll(hy.rho,1,axis=0) )
        if hy.ndim == 2:
            hy.uyi = 0.5 * ( hy.rhouy/hy.rho + np.roll(hy.rhouy,1,axis=1)/np.roll(hy.rho,1,axis=1) )
        # compute timestep based on CFL
        dt = hy.cfl * hy.dx / np.max(abs(hy.uxi))
        if hy.ndim == 2:
            dt = np.min( [ dt, hy.cfl * hy.dy / np.max(abs(hy.uyi)) ] )
        # print integral quantities
        hy.get_integral_quantities()
        # do advection
        time += dt
        print('step={:6d}'.format(step), ' |  time={:10.6e}'.format(time), ' |  dt={:10.6e}'.format(dt))
        hy.advect(hy.rho, dt, sweep='x')
        if hy.ndim == 2:
            hy.advect(hy.rho, dt, sweep='y')
        hy.rhoux = hy.rho * vx # set for constant velocity
        if hy.ndim == 2:
            hy.rhouy = hy.rho * vy # set for constant velocity
    input("Reached end of simulation. Press enter to exit.")
# ================== END advect_test ===================

# ================== START hydro_test ===================
def hydro_test(hy):
    print("Start hydro_test...")
    # initialise grid
    hy.init_grid(xr=[0.0, 1.0])
    hy.parameter_info()
    # set initial conditions
    if hy.ndim == 1:
        hy.rho = 1.0 + np.exp(-(hy.x-0.75)**2/0.1**2)
        index = (hy.x >= 0.25-0.05) & (hy.x <= 0.25+0.05)
        hy.rho[index] = 2.0
    if hy.ndim == 2:
        hy.rho = 1.0 + np.exp(-(hy.x-0.75)**2/0.1**2)
        index = (hy.x >= 0.25-0.05) & (hy.x <= 0.25+0.05) & (hy.y >= 0.25-0.05) & (hy.y <= 0.25+0.05)
        hy.rho[index] = 2.0
    hy.rhoetot = 0.5*hy.rhoux**2/hy.rho + 1.0*hy.rho/(hy.gamma-1) # eint = c_s^2 / (gamma-1), so this is c_s = 1
    if hy.ndim == 2:
        hy.rhoetot = 0.5*(hy.rhoux**2+hy.rhouy**2)/hy.rho + 1.0*hy.rho/(hy.gamma-1) # eint = c_s^2 / (gamma-1), so this is c_s = 1
    # plot initial conditions
    dr_rho = [0.0, 3.0]
    dr_u   = [-1.0, 1.0]
    dr_p   = [0.0, 10.0]
    cache = hy.update_plot(vars=["dens", "velx", "pres"], dr=[dr_rho, dr_u, dr_p])
    # evolve
    hy.evolve(cache=cache)
# ================== END hydro_test ===================

# ================== START shocktube ===================
def shocktube(hy):
    print("Start shocktube...")
    # initialise grid
    hy.init_grid(xr=[-50, 50], yr=[-50, 50])
    hy.parameter_info()
    # set initial conditions
    index = hy.x < 0 # left state
    hy.rho[index] = 1e5
    hy.rhoetot[index] = 1.0/(hy.gamma-1) # eint = c_s^2 / (gamma-1), so this is c_s = 1
    index = hy.x >= 0 # right state
    hy.rho[index] = 1.25e4
    hy.rhoetot[index] = 0.1/(hy.gamma-1) # eint = c_s^2 / (gamma-1), so this is c_s = 0.1
    # plot initial conditions
    dr_rho = [0.0, 1.1e5]
    dr_u   = [-1e-2, 1e-2]
    dr_p   = [0.0, 1.1]
    cache = hy.update_plot(vars=["dens", "velx", "pres"], dr=[dr_rho, dr_u, dr_p])
    # evolve
    hy.evolve(cache=cache)
# ================== END shocktube ===================

# ================== START shockpiston ===================
def shockpiston(hy):
    print("Start shockpiston...")
    # initialise grid
    hy.init_grid(xr=[-0.5, 0.5], yr=[-0.5, 0.5])
    hy.parameter_info()
    # set initial conditions
    prs = hy.rho*0 # make local pressure grid
    vel = hy.rho*0 # make local velocity grid
    # right state
    index = hy.x >= 0
    hy.rho[index] = 1.0
    vel[index] = 0.0
    prs[index] = 1.0/hy.gamma # set pressure to 1/gamma, such that cs = 1 for rho = 1
    # left state
    index = hy.x < 0
    hy.rho[index] = (hy.gamma+1.0) * 100 / ( (hy.gamma-1.0)*100 + 2 )
    vel[index] = 10.0 # piston speed; Mach 10 => results in a shock speed of about 4/3 * 10 ~ 13.33
    prs[index] = ( 1.0 - hy.gamma + 2.0*hy.gamma*100 ) / (hy.gamma+1.0)
    # set remaining quantities
    hy.rhoux = hy.rho * vel
    hy.rhoetot = 0.5*hy.rhoux**2/hy.rho + prs/(hy.gamma-1.0)
    # plot initial conditions
    dr_rho = [0.0, 6.0]
    dr_u   = [-1, 11]
    dr_p   = [-10.0, 150]
    cache = hy.update_plot(vars=["dens", "velx", "pres"], dr=[dr_rho, dr_u, dr_p])
    # evolve
    hy.evolve(cache=cache)
# ================== END shockpiston ===================

# ================== START sedov ===================
def sedov(hy):
    print("Start sedov...")
    # initialise grid
    hy.init_grid(xr=[-0.5, 0.5], yr=[-0.5, 0.5])
    hy.parameter_info()
    # set initial conditions
    hy.rho[:] = 1.0
    hy.rhoetot = 1.0*hy.rho/(hy.gamma-1) # eint = c_s^2 / (gamma-1), so this is c_s = 1
    ind = np.sqrt(hy.x**2+hy.y**2) <= 0.05
    hy.rhoetot[ind] = 1e1*hy.rho[ind]/(hy.gamma-1) # eint = c_s^2 / (hy.gamma-1), so this is c_s^2 = 10
    # plot initial conditions
    dr_rho = [0.0, 3.5]
    dr_u   = [-2, 2]
    dr_p   = [0.0, 11.0]
    cache = hy.update_plot(vars=["dens", "velx", "pres"], dr=[dr_rho, dr_u, dr_p])
    # evolve
    hy.evolve(cache=cache)
# ================== END sedov ===================

# ================== START kelvin_helmholtz ===================
def kelvin_helmholtz(hy):
    print("Start kelvin_helmholtz...")
    # initialise grid
    hy.init_grid(xr=[0.0, 1.0], yr=[0.0, 1.0])
    hy.parameter_info()
    # set initial conditions
    layer_thickness = 0.01
    vy_perturb      = 0.01
    smoothing_size  = 0.2
    pres            = 2.5
    dens            = 1.5
    dens_perturb    = 0.5
    vx              = 0.5
    # geometry
    xmid = 0.5*(hy.x.max()-hy.x.min())
    ymid = 0.5*(hy.y.max()-hy.y.min())
    y_tilde = abs(hy.y - ymid) - 0.25
    # rho and rhoux, rhouy, rhoetot
    hy.rho = dens - dens_perturb * np.tanh(y_tilde/layer_thickness)
    hy.rhoux = hy.rho*(vx * np.tanh(y_tilde/layer_thickness))
    hy.rhouy = hy.rho*(vy_perturb * np.cos(4*np.pi*(hy.x-xmid)) * np.exp(-y_tilde**2/smoothing_size**2))
    hy.rhoetot = 0.5*(hy.rhoux**2+hy.rhouy**2)/hy.rho + pres/(hy.gamma-1)
    # plot initial conditions
    dr_rho = [0.8, 2.2]
    dr_ux  = [-0.2, 0.2]
    dr_uy  = [-0.2, 0.2]
    cache = hy.update_plot(vars=["dens", "velx", "vely"], dr=[dr_rho, dr_ux, dr_uy])
    # evolve
    hy.evolve(cache=cache)
# ================== END kelvin_helmholtz ===================

'''
# ================== START rayleigh_taylor ===================
def rayleigh_taylor(hy):
    print("Start rayleigh_taylor...")
    # initialise grid
    hy.init_grid(xr=[0.0, 1.0], yr=[0.0, 2.0])
    hy.parameter_info()
    # set initial conditions
    ypos = 1.5
    yperturb = 0.1
    pres = 2.5
    # rho and rhoux, rhouy, rhoetot
    index = hy.y > ypos - yperturb*np.sin(np.pi*hy.x)
    hy.rho[index] = 10.0 # top
    hy.rho[~index] = 1.0 # bottom
    hy.rhoux = hy.rho*0.0
    hy.rhouy = hy.rho*0.0
    hy.rhoetot = 0.5*(hy.rhoux**2+hy.rhouy**2)/hy.rho + pres/(hy.gamma-1)
    # plot initial conditions
    dr_rho = [0.5, 20]
    dr_ux  = [-1, 1]
    dr_uy  = [-1, 1]
    cache = hy.update_plot(vars=["dens", "velx", "vely"], dr=[dr_rho, dr_ux, dr_uy])
    # evolve
    hy.evolve(cache=cache, gravity=10.)
# ================== END rayleigh_taylor ===================
'''

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

    parser = argparse.ArgumentParser(description='Hydroself.dynamics code.')
    parser.add_argument("-simulation", "--simulation",
                        choices=['advect_test', 'hydro_test', 'shocktube', 'shockpiston', 'sedov', 'kh', 'rt'],
                        default='advect_test', help="The simulation setup/test (default: %(default)s).")
    parser.add_argument("-scheme", "--scheme", choices=['upwind', 'centred', 'downwind'],
                        help="Finite difference scheme for advection (default: %(default)s).")
    parser.add_argument("-fluxlim", "--fluxlim", choices=['donor-cell', 'lax-wendroff', 'beam-warming', 'fromm', 'minmod', 'superbee', 'mc', 'vanLeer'],
                        help="Flux limiter.")
    parser.add_argument("-boundary", "--boundary", choices=['periodic', 'outflow'],
                        help="Boundary condition.")
    parser.add_argument("-cfl", "--cfl", type=float,
                        help="CFL timestep safety factor (default: %(default)s).")
    parser.add_argument("-gamma", "--gamma", type=float,
                        help="Adiabatic index (default: %(default)s).")
    parser.add_argument("-ndim", "--ndim", type=int, choices=[1,2],
                        help="1D or 2D setup (default: %(default)s).")
    parser.add_argument("-ncells", "--ncells", nargs='*', type=int,
                        help="Number of cells (default: %(default)s).")
    parser.add_argument("-nend", "--nend", type=int,
                        help="max number of timesteps (default: %(default)s).")
    parser.add_argument("-tmax", "--tmax", type=float,
                        help="end time (default: %(default)s).")
    parser.add_argument("-dpi", "--dpi", type=int,
                        help="DPI for plotting (default: %(default)s).")
    parser.add_argument("-verbose", "--verbose", type=int, choices=[0,1,2],
                        default=1, help="Verbose level (default: %(default)s).")
    args = parser.parse_args()

    # new hydro class object
    hy = hydro_class()

    # set hydro class parameters
    if args.ndim     is not None: hy.ndim     = args.ndim
    if args.ncells   is not None: hy.ncells   = args.ncells*hy.ndim
    if args.scheme   is not None: hy.scheme   = args.scheme
    if args.fluxlim  is not None: hy.fluxlim  = args.fluxlim
    if args.boundary is not None: hy.boundary = args.boundary
    if args.cfl      is not None: hy.cfl      = args.cfl
    if args.gamma    is not None: hy.gamma    = args.gamma
    if args.nend     is not None: hy.nend     = args.nend
    if args.tmax     is not None: hy.tmax     = args.tmax
    if args.dpi      is not None: hy.dpi      = args.dpi
    if args.verbose  is not None: hy.verbose  = args.verbose

    # advection test
    if args.simulation == "advect_test":
        if args.fluxlim is None: hy.fluxlim = 'superbee'
        advect_test(hy)

    # hydro test
    if args.simulation == "hydro_test":
        if args.fluxlim is None: hy.fluxlim = 'vanLeer'
        if args.tmax is None: hy.nend = 400
        hydro_test(hy)

    # shocktube
    if args.simulation == "shocktube":
        if args.fluxlim is None: hy.fluxlim = 'minmod'
        if args.boundary is None: hy.boundary = 'outflow'
        if args.cfl is None: hy.cfl = 0.5
        if args.tmax is None: hy.tmax = 7.5e3
        shocktube(hy)

    # shockpiston
    if args.simulation == "shockpiston":
        if args.fluxlim is None: hy.fluxlim = 'donor-cell'
        if args.boundary is None: hy.boundary = 'outflow'
        if args.cfl is None: hy.cfl = 0.5
        if args.tmax is None: hy.tmax = 0.03
        shockpiston(hy)

    # sedov (Sedov-Taylor explosion)
    if args.simulation == "sedov":
        if args.fluxlim is None: hy.fluxlim = 'donor-cell'
        if args.boundary is None: hy.boundary = 'outflow'
        if args.cfl is None: hy.cfl = 0.5
        if args.tmax is None: hy.tmax = 0.1
        sedov(hy)

    # kh (Kelvin-Helmholtz instability)
    if args.simulation == "kh":
        if args.ndim is None: hy.ndim = 2
        if args.ncells is None: hy.ncells = [128, 128]
        if args.fluxlim is None: hy.fluxlim = 'donor-cell'
        if args.gamma is None: hy.gamma = 1.4
        if args.tmax is None: hy.tmax = 3.0
        if args.nend is None: hy.nend = 100000
        kelvin_helmholtz(hy)

    # rt (Rayleigh-Taylor instability)
    if args.simulation == "rt":
        if args.ndim is None: hy.ndim = 2
        if args.ncells is None: hy.ncells = [128, 256]
        if args.fluxlim is None: hy.fluxlim = 'donor-cell'
        if args.gamma is None: hy.gamma = 1.4
        if args.tmax is None: hy.tmax = 10.0
        if args.nend is None: hy.nend = 100000
        rayleigh_taylor(hy)