from taskutil import get_global_namespace

def annular_mask(template, outer, inner=None, center=None,
                 output=None, overwrite=False, name=None):
    """
    NOTE: setting center doesn't work yet!  It defaults to the center of
    template.  Also, it only masks the first xy plane.
    
    Returns the name of a float "mask" based on template with pixel values of 1
    for positions with distances from center between outer and inner and 0
    elsewhere.  center will default to the center of template if not specified.
    The output name will be template.mask if not specified.

    template will get a new mask corresponding to masktool.  If name is not
    supplied, it will default to something like 'mask0'.

    No antialiasing attempt is made, partly because I don't know if it would do
    any good.  The annulus is circular in template's projection, not
    necessarily on the sky.

    Example:
       masktool = annular_mask('fascinating.im', '5arcsec',
                               center='J2000 18h00m00.03s -23d0m0.0s',
                               output='fascinating_5arcsec_disk.mask')

       sets masktool to an open ia tool containing a mask which is a disk
       (inner defaults to 0) centered around J2000 18h00m00.03s -23d0m0.0s with
       radius 5 arcseconds.  The mask will be written to
       fascinating_5arcsec_disk.mask on disk when masktool.close() or .done()
       is called.
    """
    myf = get_global_namespace()
    ia = myf['ia']
    pl = myf['pl']
    qa = myf['qa']

    ia.open(template)
    summary = ia.summary()['header']
    bb = ia.boundingbox()
    shape = ia.shape()
    cs = ia.coordsys()
    ia.close()

    xpixqa = qa.quantity(str(summary['incr'][0]) + summary['axisunits'][0])
    ypixqa = qa.quantity(str(summary['incr'][1]) + summary['axisunits'][1])

    # Convert everything to arcseconds for convenience.
    onearcsec = qa.quantity('1arcsec')
    xpixas = qa.canonical(qa.div(xpixqa, onearcsec))['value']
    ypixas = qa.canonical(qa.div(ypixqa, onearcsec))['value']
    xoy = abs(xpixas / ypixas)
    
    if isinstance(outer, str):
        outerqa = qa.quantity(outer)
    outeras = qa.canonical(qa.div(outerqa, onearcsec))['value']

    if inner:
        if isinstance(inner, str):
            innerqa = qa.quantity(inner)
        inneras = qa.canonical(qa.div(innerqa, onearcsec))['value']
    else:
        inneras = 0.0

    if inneras > outeras:
        print "Warning!  outer =", outer, "< inner =", inner
        print "Switching them around in annular_mask()"
        temp = outeras
        outeras = inneras
        inneras = temp
    
    # Use newimagefromshape instead of newimagefromimage to allow making an OTF
    # mask, not that seems to work anyway.
    if not output:
        output = template + '.mask'
    masktool = ia.newimagefromshape(output, shape, csys=cs.torecord(),
                                    overwrite=overwrite)
    maskarr = pl.zeros(shape[0:2], dtype=int)

    if center:
        # For some unknown reason topixel returns a dictionary containing an array
        # of floats.
        centpix = masktool.topixel(center)['numeric'].tolist()
    else:
        centpix = 0.5 * (bb['blc'] + bb['trc'])

    cx, cy = centpix[0:2]
    #firstplane = tuple(pl.zeros([len(shape[2:]) + 1]).tolist())

    def setdisk(radas, value):
        """
        Sets the elements within radas of (cx, cy) to value.
        """
        outxrad = abs(radas / xpixas)
        outyrad = abs(radas / ypixas)
        minx = int(pl.ceil(max(0, cx - outxrad)))
        maxx = int(pl.floor(min(shape[0], cx + outxrad)))
        Miny = int(pl.ceil(max(0, cy - outyrad)))
        Maxy = int(pl.floor(min(shape[0], cy + outyrad)))
        maxd2 = outxrad**2
        for i in xrange(minx, maxx):
            yspreado2 = xoy * pl.sqrt(maxd2 - (i - cx)**2)
            miny = max(Miny, int(pl.ceil(cy - yspreado2)))
            maxy = min(Maxy, int(pl.floor(cy + yspreado2)))
            #maskarr[i, miny:maxy][firstplane] = 1
            maskarr[i, miny:maxy] = value

    setdisk(outeras, 1) # I figure it's easier to set some pixels twice
    setdisk(inneras, 0) # than to make the loops more complicated.

    # replicate=True will copy the 2D array to all planes in a cube.
    masktool.putchunk(maskarr, replicate=True)
    masktool.close()    # Flush to disk.

    # AAAARRGGHH! ia.statistics complains (and doesn't work) if mask is a
    # float, and putchunk converts things to floats!
    # This might not work if masktool isn't on disk!
    ia.open(template)
    if name:     # Guard against + or - in filenames
        ia.calcmask('"%s" > 0' % output, name=name)
    else:
        ia.calcmask('"%s" > 0' % output) # Probably 'mask0'
    ia.close()
    return output