"""Class Domain -
2D triangular domains for finite-volume computations of
the advection equation.

This module contains a specialisation of class Domain from module domain.py
consisting of methods specific to the advection equantion

The equation is

  u_t + (v_1 u)_x + (v_2 u)_y = 0

There is only one conserved quantity, the level u

The advection equation is a very simple specialisation of the generic
domain and may serve as an instructive example or a test of other
components such as visualisation.

Ole Nielsen, Stephen Roberts, Duncan Gray, Christopher Zoppou
Geoscience Australia, 2004   
"""

from domain import *
Generic_domain = Domain #Rename

class Domain(Generic_domain):

    def __init__(self, coordinates, vertices, boundary = None, velocity = None):

        
        Generic_domain.__init__(self, coordinates, vertices, boundary,
                                ['level'])

        if velocity is None:
            self.velocity = [1,0]
        else:
            self.velocity = velocity

        #Only first is implemented for advection    
        self.default_order = self.order = 1  

        #Realtime visualisation
        self.visualise = False
        self.smooth = True


    def check_integrity(self):
        Generic_domain.check_integrity(self)

        msg = 'Conserved quantity must be "level"'
        assert self.conserved_quantities[0] == 'level', msg
        

    def flux_function(self, normal, ql, qr, zl=None, zr=None):
        """Compute outward flux as inner product between velocity
        vector v=(v_1, v_2) and normal vector n.
            
        if <n,v> > 0 flux direction is outward bound and its magnitude is
        determined by the quantity inside volume: ql.
        Otherwise it is inbound and magnitude is determined by the
        quantity outside the volume: qr.
        """
    
        v1 = self.velocity[0]
        v2 = self.velocity[1]


        normal_velocity = v1*normal[0] + v2*normal[1]

        if normal_velocity < 0:
            flux = qr * normal_velocity
        else:
            flux = ql * normal_velocity
        
        max_speed = abs(normal_velocity)
        return flux, max_speed        

        
    def compute_fluxes(self):
        """Compute all fluxes and the timestep suitable for all volumes
        in domain.
        
        Compute total flux for each conserved quantity using "flux_function"
        
        Fluxes across each edge are scaled by edgelengths and summed up
        Resulting flux is then scaled by area and stored in
        domain.explicit_update

        The maximal allowable speed computed by the flux_function
        for each volume
        is converted to a timestep that must not be exceeded. The minimum of
        those is computed as the next overall timestep.

        Post conditions:
        domain.explicit_update is reset to computed flux values
        domain.timestep is set to the largest step satisfying all volumes. 
        """

        import sys
        from Numeric import zeros, Float
        from config import max_timestep

        N = self.number_of_elements
    
        neighbours = self.neighbours
        neighbour_edges = self.neighbour_edges
        normals = self.normals

        areas = self.areas
        radii = self.radii
        edgelengths = self.edgelengths
            
        timestep = max_timestep #FIXME: Get rid of this

        #Shortcuts
        Level = self.quantities['level']

        #Arrays
        level = Level.edge_values

        level_bdry = Level.boundary_values
    
        flux = zeros(1, Float) #Work array for summing up fluxes

        #Loop
        for k in range(N):
            optimal_timestep = float(sys.maxint)

            flux[:] = 0.  #Reset work array
            for i in range(3):
                #Quantities inside volume facing neighbour i
                ql = level[k, i]

                #Quantities at neighbour on nearest face
                n = neighbours[k,i] 
                if n < 0:
                    m = -n-1 #Convert neg flag to index
                    qr = level_bdry[m]
                else:    
                    m = neighbour_edges[k,i]
                    qr = level[n, m]

                    
                #Outward pointing normal vector   
                normal = normals[k, 2*i:2*i+2]

                #Flux computation using provided function
                edgeflux, max_speed = self.flux_function(normal, ql, qr)
                flux -= edgeflux * edgelengths[k,i]
                
                #Update optimal_timestep
                try:
                    optimal_timestep = min(optimal_timestep, radii[k]/max_speed)
                except ZeroDivisionError:
                    pass

            #Normalise by area and store for when all conserved
            #quantities get updated
            flux /= areas[k]
            Level.explicit_update[k] = flux[0]
            
            timestep = min(timestep, optimal_timestep)

        self.timestep = timestep    
	


    def evolve(self, yieldstep = None, finaltime = None):
        """Specialisation of basic evolve method from parent class
        """
        
        #Initialise real time viz if requested
        if self.visualise is True and self.time == 0.0:
            import realtime_visualisation as visualise
            visualise.create_surface(self)

        #Call basic machinery from parent class
        for t in Generic_domain.evolve(self, yieldstep, finaltime):
            #Real time viz
            if self.visualise is True:
                visualise.update(self)

            #Pass control on to outer loop for more specific actions    
            yield(t)