from numarray import *


"""
  ---------------------------------------------------------------------
  Copyright (C) 2003, 2004  Michael Zingale

  This software is distributed under the terms of the GNU General
  Public License, version 2.  See COPYING for details.  It is covered
  by no warantee whatsoever, explicit or implied.
  ---------------------------------------------------------------------
"""


def _restrict(nx_f, ng_f, U_f, nx_c, ng_c, U_c):

    imin_f = ng_f
    imax_f = ng_f + nx_f - 1

    imin_c = ng_c
    imax_c = ng_c + nx_c - 1

    # fill the coarse patch with the restricted data -- just average
    # the fine zones into their corresponding coarse zone.  This is
    # done by shifting our view into the U_f array and using a stride
    # of 2 in the indexing.
    U_c[imin_c:imax_c+1] = 0.5*(U_f[imin_f:imax_f+1:2] +
                                U_f[imin_f+1:imax_f+1:2])



def _prolong(nx_c, ng_c, U_c, nx_f, ng_f, U_f):

    """
    We will reconstruct the data in the zone from the zone-averaged
    variables using the same limited slopes as in the advection
    routine.
    
              (x)         
    f(x,y) = m    x/dx + <f>

    where the m's are the limited differences in each direction.  When
    averaged over the parent cell, this reproduces <f>.

    Each zone's reconstrution will be averaged over 2 children.

    |           |     |     |     |
    |    <f>    | --> |  1  |  2  |
    +-----------+     +-----+-----+                           

    We will fill each of the finer resolution zones by filling
    all the 1's together, using a stride 2 into the fine array.
    Then the 2's, this allows us to operate in a vector
    fashion.  All operations will use the same slopes for their
    respective parents.
    
    """

    imin_c = ng_c
    imax_c = ng_c + nx_c - 1

    imin_f = ng_f
    imax_f = ng_f + nx_f - 1

    
    # get the limited differences in each direction for all interior
    # zones in our current mesh
    m_x = zeros((nx_c + 2*ng_c), Float64)
    
    m_x[imin_c:imax_c+1] = 0.5*(U_c[imin_c+1:imax_c+2] - U_c[imin_c-1:imax_c])
            
    # fill the '1' children
    U_f[imin_f:imax_f+1:2] = U_c[imin_c:imax_c+1] - 0.25*m_x[imin_c:imax_c+1]
    U_f[imin_f+1:imax_f+1:2] = U_c[imin_c:imax_c+1] + 0.25*m_x[imin_c:imax_c+1]


class patch1d:
    """
    the 2-d patch class.  The patch object will contain the coordinate
    information (at various centerings) and the unknown data itself.

    a patch is built in a multi-step process before it can be used:

    create the patch object:

        myPatch = patch1d(nx, ny)

    register any variables that w eexpect to live on this patch

        myPatch.registerVar('density')
        myPatch.registerVar('x-momentum')
        myPatch.registerVar('y-momentum')
        ...

    finally, initialize the patch

        myPatch.init()

    This last step actually allocates the storage for the unknowns.  Once
    this is done, the patch is considered to be locked.  New variables cannot
    be added.
    """

    # the constructor function.  The only data that we require is the number
    # of points that make up the mesh.
    #
    # We optionally take the extrema of the domain (assume it is [0,1]x[0,1])
    # and the number of guardcells (assume 1)

    def __init__ (self, nx, ng=1, \
                  xmin=0.0, xmax=1.0, \
                  xlb = "reflect", xrb = "reflect"):

        self.vars = []
        self.nvar = 0

        self.nx = nx
        self.ng = ng

        # domain extrema
        self.xmin = xmin
        self.xmax = xmax

        # domain boundary conditions
        self.xlb = xlb
        self.xrb = xrb

        # compute the indices of the block interior (excluding guardcells)
        self.imin = ng
        self.imax = ng+nx-1

        # define the coordinate information at the left, center, and right
        # zone positions
        self.dx = (xmax - xmin)/nx

        self.xl = (arange(nx+2*ng) - ng)*self.dx + xmin
        self.xr = (arange(nx+2*ng) + 1.0 - ng)*self.dx + xmin
        self.x = 0.5*(self.xl + self.xr)

        # time/timestep information
        self.time = 0.0
        self.dt = 0.0
        self.nstep = 0
        

    def registerVar(self, name):
        """ register a variable with patch object """

        self.vars.append(name)
        self.nvar += 1


    def init(self):
        """
        called after all the variables are registered and allocates
        the storage for the unknowns
        """
        
        self.data = reshape(array(zeros( (2*self.ng+self.nx)*(self.nvar), Float64 )), \
                            (self.nvar, 2*self.ng+self.nx))

        
    def __str__(self):
        """ print out some basic information about the patch object """
        string = "1-d patch: nx = " + `self.nx` + ', ng = ' + `self.ng` + "\n" + \
                 "           nvars = " + `self.nvar` + "\n" + \
                 "           variables: " + `self.vars`
        return string
        

    def getVar(self, name):

        n = self.vars.index(name)
        return self.data[n,:]


    def getVarByIndex(self, index):
        return self.data[index,:]


    def zero(self, name):
        n = self.vars.index(name)
        self.data[n,:,:] = 0.0
        
        
    def prolong(self, varList=["all"]):
        """ create a new mesh object from the existing one by prolonging to
            a grid that is twice as large.
            """

        # create the new mesh object
        fP = patch1d(self.nx*2, ng=self.ng,
                     xmin=self.xmin, xmax=self.xmax,
                     xlb=self.xlb, xrb=self.xrb)

        try: varList.index("all")
        except ValueError:
            rvars = varList
        else:
            rvars = self.vars
            
        for var in rvars:    
            fP.registerVar(var)

        fP.init()
        
        for var in rvars:    

            nself = self.vars.index(var)
            nfp = rvars.index(var)

            _prolong(self.nx, self.ng, self.data[nself,:],
                     fP.nx, fP.ng, fP.data[nfp,:])
            
        return fP
    

        
    def restrict(self, varList=["all"]):
        """ restrict the current mesh to a coarser mesh.  Assume a factor of
            two in their resolution """

        # we can only restrict if our mesh is an even size
        if (self.nx%2 != 0):
            print "ERROR: restriction only operates on meshes with even sizes"
            return NULL
        
        # create the new mesh object
        coarsePatch = patch1d(self.nx/2, ng=self.ng,
                              xmin=self.xmin, xmax=self.xmax,
                              xlb=self.xlb, xrb=self.xrb)

        try: varList.index("all")
        except ValueError:
            rvars = varList
        else:
            rvars = self.vars

        for var in rvars:    
            coarsePatch.registerVar(var)


        coarsePatch.init()

        for var in rvars:    

            nself = self.vars.index(var)
            ncp = rvars.index(var)
            
            _restrict(self.nx, self.ng, self.data[nself,:],
                      coarsePatch.nx, coarsePatch.ng, coarsePatch.data[ncp,:])

        return coarsePatch
    

        
    def fillBC(self):
        """ fill the boundary conditions """
    
        # do the boundaries one at a time.  The single grid case
        # is easy -- every boundary is a physical boundary of the
        # domain, so there is not much trickery to do

        # we do periodic, outflow, reflect, and in the y-direction,
        # HSE boundaries.
        
        # lower X boundary
        if (self.xlb == "reflect" or self.xlb == "outflow"):
            n = 0
            while n < self.nvar:

                i = 0
                while i < self.imin:
                    self.data[n,i] = self.data[n,self.imin]
                    
                    i += 1
                
                n+= 1

            # if we are reflecting, actually reflect the x-momentum now
            if (self.xlb == "reflect"):
                n = self.vars.index("x-momentum")

                i = 0
                while i < self.imin:
                    self.data[n,i] = -self.data[n,2*self.ng-i-1]
                    
                    i += 1


        elif (self.xlb == "periodic"):
            n = 0
            while n < self.nvar:

                i = 0
                while i < self.imin:
                    self.data[n,i] = self.data[n,self.imax-self.ng+i+1]

                    i += 1

                n += 1
            

        # upper X boundary
        if (self.xrb == "reflect" or self.xrb == "outflow"):
            n = 0
            while n < self.nvar:

                i = self.imax+1
                while i < self.nx+2*self.ng:
                    self.data[n,i] = self.data[n,self.imax]

                    i += 1
                
                n+= 1

            # if we are reflecting, actually reflect the x-momentum now
            if (self.xrb == "reflect"):
                n = self.vars.index("x-momentum")

                i = self.imax+1
                while i < self.nx+2*self.ng:
                    self.data[n,i] = -self.data[n,2*self.imax-i+1+self.ng]
                    
                    i += 1


        elif (self.xrb == "periodic"):
            n = 0
            while n < self.nvar:

                i = self.imax+1
                while i < 2*self.ng + self.nx:
                    self.data[n,i] = self.data[n,i-self.imax-1+self.ng]

                    i += 1

                n += 1