import numpy
import itertools

def fit_polynomial(image, mask=None, degree=2, poly_args=None):
    """Return an array containing the values of a polynomial fit to the input.

    Parameters:
        image: 2D array.
        mask: if specified, use only the values in the image at points where
            the mask is true
        degree: degree of the fit polynomial in x and y
        poly_args: if not None, must have been generated by get_poly_args(). Useful
            when multiple images of the same shape are to be fit to a polynomial
            so the poly_args array does not need to be recalculated.
    """
    image = numpy.asarray(image)
    shape = image.shape
    if poly_args is None:
        poly_args = get_poly_args(shape, degree)
    if mask is not None:
        mask = numpy.asarray(mask, dtype=bool)
        if not numpy.any(mask):
            return numpy.zeros_like(image)
        a = poly_args[mask, :]
        b = image[mask]
    else:
        a = poly_args.reshape((shape[0]*shape[1], poly_args.shape[-1]))
        b = image.ravel()
    # solve ax = b, where a contains the polynomial expansion of the spatial positions
    # of the image samples, and b contains the image samples
    coefficients, residues, rank, singval = numpy.linalg.lstsq(a, b)
    output = (poly_args * coefficients).sum(axis=-1)
    return output

def subtract_background(image, mask=None, degree=2):
    """Fit the input to a low-order polynomial to approximate the background
    and then subtract that from the image. See fit_polynomial() for further
    details.

    Note: effectively a high-pass filter.
    """
    background = fit_polynomial(image, mask, degree)
    sub = (image - background).clip(0, numpy.inf)
    return sub.astype(image.dtype)

def get_poly_args(shape, degree):
    """Return an array containing the arguments for a polynomial in x and y defined
    over an array of the given shape.

    Let 'binomial(degree)' refer to the number of parameters of a polynomial function
    of the given degree. Then the output shape = (shape[0], shape[1], binomial(degree))

    Example: let degree = 2, and let x, y = numpy.indices(shape)
       result = get_poly_args(shape, degree)
       result[:,:,0] == 1
       result[:,:,1] == x
       result[:,:,2] == y
       result[:,:,3] == x*x
       result[:,:,4] == x*y
       result[:,:,5] == y*y
    """
    x, y = numpy.indices(shape)
    poly_args = [numpy.ones(shape)]
    for i in range(1, degree+1):
        for v in itertools.combinations_with_replacement([x, y], i):
            poly_args.append(numpy.prod(v, axis=0))
    poly_args = numpy.array(poly_args) # shape = binomial(degree), shape[0], shape[1]
    poly_args = poly_args.transpose((1,2,0)) # shape = shape[0], shape[1], binomial(degree)
    return poly_args