"""
Spatially average a function inside mesh triangles


Gareth Davies, Geoscience Australia 2014+
"""

import scipy
import anuga.utilities.spatialInputUtil as su

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


def make_spatially_averaged_function(q_function,
                                     domain,
                                     approx_grid_spacing=[1., 1.],
                                     chunk_size=1e+04,
                                     averaging='mean',
                                     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' or 'harmonic_mean' 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' / 'harmonic_mean'

    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

        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):

                mesh_tri = \
                    domain.mesh.vertex_coordinates[
                        range(3 * j, 3 * j + 3), :].tolist()

                pts = su.gridPointsInPolygon(
                    mesh_tri,
                    approx_grid_spacing=approx_grid_spacing)

                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()
                elif(averaging == 'harmonic_mean'):
                    out[j] = 1.0 / (1.0 / allTopo[out_indices]).mean()
                else:
                    raise Exception, 'Unknown value of averaging'
        return(out)

    return elevation_setter
