"""
Make function to compute average of another function inside mesh triangles

Gareth Davies, Geoscience Australia 2014+
"""

import scipy
import anuga.utilities.spatialInputUtil as su
from anuga.geometry.polygon import inside_polygon

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


def make_spatially_averaged_function(q_function,
                                     domain,
                                     approx_grid_spacing=[1., 1.],
                                     chunk_size=1e+04,
                                     averaging='mean',
                                     polygons_for_averaging=None,
                                     verbose=True):
    """

    Given a function q_function (used to set quantities), return a function
    which estimates the spatial-average of q_function in each mesh triangle.

    It does this by generating a grid near the mesh triangle, with points
    spaced ~ approx_grid_spacing, then finding those points inside the mesh
    triangle, computing q_function at each, and averaging the result.

    If averaging='mean', then the mean is returned. Alternatively 'min' or
    'max' can be used.

    The returned function only applies to domain centroids, since mesh
    triangles are used to spatial average.

    All the domain centroids must be passed to the returned function at the
    same time, in the order that they appear in domain, as we use the ordering
    to associate centroids and vertices.

    INPUTS:

    q_function -- the function that you wish to spatially
                     average over the mesh triangles

    domain    -- the ANUGA domain

    approx_grid_spacing --
          averaging is computed from points in each triangle, generated by
          anuga.utilities.SpatialInputUtil.gridPointsInPolygon, with the
          value of approx_grid_spacing passed there

    chunk_size -- Number of mesh triangles to average in each call to
         q_function. A suitably large chunk_size can reduce function call
         overhead for some q_functions, but might have consume lots of
         memory if there are many grid-points in each triangle

    averaging -- character, what to do with the values inside the cell
                'mean' / 'min' / 'max'

    polygons_for_averaging -- list of polygons or None. If not None, then
         we only apply the averaging to points inside the polygons. Each
         polygon can be specified as either its filename (accepted by read_polygon)
         or directly in the list/array formats accepted by anuga.

    verbose -- print information

    OUTPUTS:

    function F(x,y) which can be passed as e.g.

    domain.set_quantity('elevation', F, location='centroids')

    """

    chunk_size = int(chunk_size)
    averaging = averaging

    def elevation_setter(xc, yc):

        # Return scipy array of values
        out = xc * 0.

        # Get multiple elevation values in each triangle.
        # Process triangles in chunks to reduce function call overhead
        lx = len(xc)
        lx_div_cs = scipy.ceil(lx * 1. / (1. * chunk_size)).astype(int)

        # Crude check that xc/yc are the centroid values
        #
        erMess = ' Result of make_meanFun can ONLY be applied to a vector' +\
            ' of ALL centroid coordinates\n' +\
            ' (since mesh triangles are used to spatially average)'
        assert scipy.all(xc == domain.centroid_coordinates[:, 0]), erMess
        assert scipy.all(yc == domain.centroid_coordinates[:, 1]), erMess

        # Find triangles in which we want to average
        if polygons_for_averaging is not None:

            averaging_flag = 0*xc

            # Need georeferenced centroid coordinates to find which
            # are in the polygon
            xll = domain.geo_reference.xllcorner
            yll = domain.geo_reference.yllcorner
            centroid_coordinates_georef = scipy.vstack([xc + xll, yc + yll]).transpose()

            for j in range(len(polygons_for_averaging)):
                poly_j = polygons_for_averaging[j]
                # poly_j can either be a polygon, or a filename
                if type(poly_j) is str:
                    poly_j = su.read_polygon(poly_j)
                
                points_in_poly_j = inside_polygon(centroid_coordinates_georef, 
                    poly_j)
                
                averaging_flag[points_in_poly_j] = 1
                
        else:
            averaging_flag = 1 + 0*xc
        

        for i in range(lx_div_cs):
            # Evaluate in triangles lb:ub
            lb = i * chunk_size
            ub = min((i + 1) * chunk_size, lx)

            if verbose:
                print 'Averaging in triangles ', lb, '-', ub - 1

            # Store x,y,triangleIndex
            px = scipy.array([])
            py = scipy.array([])
            p_indices = scipy.array([])

            for j in range(lb, ub):
                # If we average this cell, then get a grid
                # of points in it. Otherwise just get the centroid
                # coordinates. 
                if averaging_flag[j] == 1:
                    mesh_tri = \
                        domain.mesh.vertex_coordinates[
                            range(3 * j, 3 * j + 3), :].tolist()

                    pts = su.gridPointsInPolygon(
                        mesh_tri,
                        approx_grid_spacing=approx_grid_spacing)
                else:
                    # Careful to keep this a 2D array
                    pts = domain.centroid_coordinates[j,:, None].transpose()

                px = scipy.hstack([px, pts[:, 0]])

                py = scipy.hstack([py, pts[:, 1]])

                p_indices = scipy.hstack([p_indices,
                                          scipy.repeat(j, len(pts[:, 0]))])

            # Get function values at all px,py
            if verbose:
                print '  Evaluating function at ', len(px), ' points'

            allTopo = q_function(px, py)

            # Set output values in lb:ub
            for j in range(lb, ub):
                out_indices = (p_indices == j).nonzero()[0]
                assert len(out_indices) > 0
                if(averaging == 'mean'):
                    out[j] = allTopo[out_indices].mean()
                elif(averaging == 'min'):
                    out[j] = allTopo[out_indices].min()
                elif(averaging == 'max'):
                    out[j] = allTopo[out_indices].max()
                else:
                    raise Exception('Unknown value of averaging')
        return(out)

    return elevation_setter

# Quick test
if __name__ == '__main__':
    import anuga
    
    domain = anuga.rectangular_cross_domain(10, 5, len1=10.0, len2=5.0)
   
    # Define a topography function where the spatial scale of variation matches
    # the scale of a mesh triangle
    def topography(x, y):
        return x%0.5

    # Do 'averaging' where 2 <= y <= 3
    polygon_for_averaging = [ [[0.0, 2.0], [0.0, 3.0], [10.0, 3.0], [10.0, 2.0]] ]

    topography_smooth = make_spatially_averaged_function(topography, domain, 
        approx_grid_spacing = [0.1, 0.1], averaging = 'min', 
        polygons_for_averaging = polygon_for_averaging,
        verbose=False)
    
    domain.set_quantity('elevation', topography_smooth, location='centroids') # Use function for elevation

    # Check that it worked

    inpol = ((domain.centroid_coordinates[:,1] >= 2.0) * 
         (domain.centroid_coordinates[:,1] <= 3.0)).nonzero()

    outpol = ((domain.centroid_coordinates[:,1] <= 2.0) + 
         (domain.centroid_coordinates[:,1] >= 3.0)).nonzero()

    elv = domain.quantities['elevation'].centroid_values

    # Check that the elevation in the 'averaging' band is very small
    # (since we used 'min' averaging)
    if elv[inpol].mean() < 1.0e-06:
        print 'PASS'
    else:
        print 'FAIL'

    # Check that no 'averaging' occurred outside the polygon
    x = domain.centroid_coordinates[:,0]
    if all(elv[outpol] - x[outpol]%0.5 == 0.0):
        print 'PASS'
    else:
        print 'FAIL'

    # Another test which can catch index errors   
 
    topography_smooth2 = make_spatially_averaged_function(topography, domain, 
        approx_grid_spacing = [0.1, 0.1], averaging = 'min', 
        verbose=False)
    
    domain.set_quantity('elevation', topography_smooth2, location='centroids')

    # If we get to here, then the above function did not hit an index error.
    print 'PASS' 
    

    # Another test which can catch index errors   
    # Make the polygon entirely outside of the domain!
    polygon_for_averaging3 = [ [[0.0, -2.0], [0.0, -3.0], [10.0, -3.0], [10.0, -2.0]] ]
    topography_smooth3 = make_spatially_averaged_function(topography, domain, 
        approx_grid_spacing = [0.1, 0.1], averaging = 'min', 
        verbose=False)
    
    domain.set_quantity('elevation', topography_smooth3, location='centroids')
    
    # If we get to here, then the above function did not hit an index error.
    print 'PASS' 
    
