########################################################################
#
#
#
########################################################################

import numpy as np
import math

#-----------------------------------------------------------------------
#
#
#
#-----------------------------------------------------------------------
class wave_system:
    '''This class defines methods and updates for the scalar advection
       equation.
    '''
    #-------------------------------------------------------------------
    #   public methods
    #-------------------------------------------------------------------
    # The values xmin and xmax need to be modified if one is picky about
    # periodic boundary conditions.
    def __init__(self,xmin,xmax,Nx,add_ghost_zones=False):
        self._N     = Nx
        if add_ghost_zones == True:
            ng = 3
            self._dx    = float(xmax - xmin)/(Nx-1-ng)
            self._xmin  = xmin - 2.0*self._dx
            self._xmax  = xmax + self._dx
        else:
            self._dx    = float(xmax - xmin)/(Nx-1)
            self._xmin  = xmin
            self._xmax  = xmax
 
        self._u_n   = np.zeros(Nx)
        self._u1    = np.zeros(Nx)
        self._u_np1 = np.zeros(Nx)
        self._f_n   = np.zeros(Nx)
        self._dt    = 0.1 * self._dx 
        self._flux_type  = 1
        self._time   = 0.0
        self._bcond  = 1
        self._ode_type = 'Huen'
        self._slope_type = 'Constant'
        print "dx = " + str(self._dx)
        print "dt = " + str(self._dt)

    #-------------------------------------------------------------------
    def set_u_n(self, ivp, args):
        self._u_n = ivp(self.x(), args)

    #-------------------------------------------------------------------
    def cal_flux(self):
        if (self._flux_type == 1):
          flux_godunov(self._f_n, self._u_n, self._slope_type, self._N)
        elif (self._flux_type == 2):
          flux_ftbs(self._f_n, self._u_n, self._N)
        elif (self._flux_type == 3):
          cfl = self._dt/self._dx
          LaxF(self._f_n, self._u_n, cfl, self._N)
        else:
          print "Unknown flux_type ",self._flux_type

    #-------------------------------------------------------------------
    def set_boundary_condition(self, bctype):
        self._bcond = bctype

    #-------------------------------------------------------------------
    def set_cfl(self, cfl):
        self._dt = cfl * self._dx

    #-------------------------------------------------------------------
    def set_flux_type(self, ft):
        self._flux_type = ft

    #-------------------------------------------------------------------
    def x(self):
        return np.linspace(self._xmin, self._xmax, self._N)

    #-------------------------------------------------------------------
    def u_n(self):
        return self._u_n[:]

    #-------------------------------------------------------------------
    def u_np1(self):
        return self._u_np1[:]

    #-------------------------------------------------------------------
    def f_n(self):
        return self._f_n[:]
 
    #-------------------------------------------------------------------
    def dx(self):
        return self._dx

    #-------------------------------------------------------------------
    def dt(self):
        return self._dt

    #-------------------------------------------------------------------
    def time(self):
        return self._time

    #-------------------------------------------------------------------
    def flux_type(self):
        return self._flux_type

    #-------------------------------------------------------------------
    def set_slope_type(self, slp):
        self._slope_type = slp

    #-------------------------------------------------------------------
    def set_ode_type(self, ode):
        self._ode_type = ode

    #-------------------------------------------------------------------
    def slope_type(self):
        return self._slope_type

    #-------------------------------------------------------------------
    def ode_type(self):
        return self._ode_type

    #-------------------------------------------------------------------
    def step(self, nsteps):
        N = self._N
        dx = self._dx
        dt = self._dt
        cfl = dt/dx

        for ii in range (0,nsteps):
            if self._ode_type == "Huen":
                Huen(self._u_np1, self._u_n, self._u1, self._f_n, 
                    self._slope_type, self._bcond, cfl, N)
            elif self._ode_type == "RK2":
                RK2(self._u_np1, self._u_n, self._u1, self._f_n,
                    self._slope_type, self._bcond, cfl, N)
            elif self._ode_type == "Euler":
                Euler(self._u_np1, self._u_n, self._f_n, 
                      self._slope_type, self._bcond, cfl, N)

            #rotate pointers
            temp = self._u_n; self._u_n = self._u_np1; self._u_np1 = temp
            self._time += self._dt
            print "...time now " + str(self._time)

           
#-------------------------------------------------------------------
# Flux routines
#-------------------------------------------------------------------
def flux_ftbs(f, u, n):
    for i in range(0,n-1):
        f[i] = 0.5*u[i]**2
    f[n-1] = f[n-2]

def flux_godunov(f, u, slope, n):
    ul, ur = Reconstruct(u, slope, n)
    ustar = Riemann(ul, ur, n)
    for i in range(n):
        f[i] = 0.5*ustar[i]*ustar[i]


def LaxF(f, u, cfl, n):
    for i in range(0,n-1):
        f[i] = 0.5*(0.5*u[i+1]*u[i+1] + u[i]*u[i]) - 0.5*(u[i+1] - u[i])/cfl
    f[n-1] = f[n-2]

def Reconstruct(u, slope, n):
    ''' This routine calculates the reconstructions for the left and right
        States.  
           Note:  What I call slope here is really the divided difference,
                  because I do not divide by dx.
    '''
    ul = np.zeros_like(u)
    ur = np.zeros_like(u)
    slp = np.zeros_like(u)

    if slope == "Constant":
        slp[:] = 0.0  # slp is already zero, but for good practice...

    elif slope == "Centered":
        for i in range(1,n-1):
            slp[i] = 0.5*(u[i+1] + u[i-1]) 

        slp[0] = 0.0   # should already be zero, but good practice...
        slp[n-1] = 0.0

    elif slope == "Minmod":
        for i in range(1,n-1):
            slp[i] = minmod(u[i+1]-u[i], u[i]-u[i-1])

        slp[0] = 0.0   # should already be zero, but good practice...
        slp[n-1] = 0.0

    elif slope == "MC":
        for i in range(1,n-1):
            p = 2.0*minmod( u[i]-u[i-1], u[i+1]-u[i])
            q = 0.5*(u[i+1]-u[i-1])
            slp[i] = minmod(p, q)

    elif slope == "Superbee":
        for i in range(1,n-1):
            p = minmod(u[i+1]-u[i], 2.0*(u[i]-u[i-1]))
            q = minmod(u[i]-u[i-1], 2.0*(u[i+1]-u[i]))
            slp[i] = maxmod(p, q)

    else:
        print "unknown slope " + slope
    
    for i in range(1,n):
        ul[i]   = u[i] + 0.5*slp[i]
        ur[i-1] = u[i] - 0.5*slp[i]

    ul[0] = ul[1]
    return ul, ur 

def Riemann(ul, ur, n):
    ''' This routine solves the Riemann Problem for Burger's Equation'''
    ustar = np.zeros_like(ur)
    for i in range(n):
        if ul[i] >= ur[i]:
            # Shock
            s = 0.5*(ul[i] + ur[i])
            if s > 0.0:
                ustar[i] = ul[i]
            else:
                ustar[i] = ur[i]
        else:
            # Rarefaction
            if ul[i] > 0.0:
                ustar[i] = ul[i]
            elif ur[i] < 0.0:
                ustar[i] = ur[i]
            else:
                ustar[i] = 0.0

    return ustar

    
def minmod(a,b):
    if a*b > 0.0:
        return np.sign(a)*min(abs(a),abs(b))
    else:
        return 0.0

def maxmod(a,b):
    if a*b > 0.0:
        return np.sign(a)*max(abs(a),abs(b))
    else:
        return 0.0
    
#-------------------------------------------------------------------
# Euler 
#-------------------------------------------------------------------
def Euler(u2, u, nf, slope_type, bcond, cfl, N):
    ''' Take one step with Euler's method'''
    flux_godunov(nf, u, slope_type, N)
    u2[1:N-1] = u[1:N-1] - cfl*(nf[1:N-1] - nf[0:N-2])
    set_bc(u2, bcond)


#-------------------------------------------------------------------
#  RK2
#-------------------------------------------------------------------
def RK2(u2, u, u1, nf, slope_type, bcond, cfl, N):
    ''' Take one step with RK2 '''

    flux_godunov(nf, u, slope_type, N)
    u1[1:N-1] = u[1:N-1] - 0.5*cfl*(nf[1:N-1] - nf[0:N-2])
    set_bc(u1, bcond)

    flux_godunov(nf, u1, slope_type, N)
    u2[1:N-1] = u[1:N-1] - cfl*(nf[1:N-1] - nf[0:N-2])
    set_bc(u2, bcond)

#-------------------------------------------------------------------
#  Huen's method
#-------------------------------------------------------------------
def Huen(u2, u, u1, nf, slope_type, bcond, cfl, N):
    ''' Take one step with Huen's method'''

    flux_godunov(nf, u, slope_type, N)
    u1[1:N-1] = u[1:N-1] - cfl*(nf[1:N-1] - nf[0:N-2])
    set_bc(u1, bcond)

    flux_godunov(nf, u1, slope_type, N)
    u2[1:N-1] = 0.5*(u[1:N-1] + u1[1:N-1]) - 0.5*cfl*(nf[1:N-1] - nf[0:N-2])
    set_bc(u2, bcond)



#-------------------------------------------------------------------
# Initial data routines
#-------------------------------------------------------------------
def id_step(x, args):
    ul = args[0]
    ur = args[1]
    return np.where (x < 0, ul, ur)

def exact_step(x, x0, args):
    ul = args[0]
    ur = args[1]
    return np.where(x < x0, ul, ur)

#-------------------------------------------------------------------
#
#  id_gaussian: Gaussian initial data
#      Here args = [A, x0, sigma]
#
#-------------------------------------------------------------------
def id_gaussian(x, args):
    A = args[0]
    x0 = args[1]
    sigma = args[2]
    return 0.1 + A*np.exp(-(x-x0)**2/sigma**2)

def id_box(x, args):
    ul = args[0]
    ur = args[1]
    return np.where (abs(x) < 0.5, ul, ur)

def id_sine(x, args):
    amp = args[0]
    u0 = args[1]
    return np.where( abs(x) < 0.5, u0 + amp*np.sin(2.0*math.pi*(x-0.5)), u0)

def set_bc(f, bctype):
    N = np.size(f)
    if (bctype == 1):
        f[0] = f[2]
        f[1] = f[2]
        f[N-2] = f[N-3]
        f[N-1] = f[N-3]
    elif (bctype ==2):
        f[0] = f[N-4]
        f[1] = f[N-3]
        f[N-2] = f[2]
        f[N-1] = f[3]
    else:
        print "unknown bctype = " + str(bctype)
        exit(2)