from numpy import array
from taskutil import get_global_namespace

def beam_minus_gaussian(beam, bmaj=None, bmin=None, bpa=None, x=None, y=None):
    """
    Returns an image of beam with a gaussian of FWHM bmaj and bmin and position
    angle bpa subtracted.  The gaussian will be fit to beam if its parameters
    are not specified.  Any specifcations should have units.
    x and y are assumed to be in J2000.  Do not specify them unless you are
    sure they are close to the correct values!
    """
    myf = get_global_namespace()
    ia  = myf['ia']
    qa  = myf['qa']
    
    ia.open(beam)
    b2dcsys = ia.coordsys([0, 1])

    # Set up fixed parameters.
    fixedstr = ""
    specifications = ia.maxfit(point=False)   # Does a parabolic fit to the maximum.

    # Don't bother fixing the flux, because fitsky is weird about that.
    [bmaj, bmin, bpa, x, y] = [quantity_or_nothing(f, qa) for f in [bmaj, bmin, bpa, x, y]]
    if bmaj:
        fixedstr += 'a'
        specifications['component0']['shape']['majoraxis'] = bmaj
    if bmin:
        fixedstr += 'b'
        specifications['component0']['shape']['minoraxis'] = bmin
    if bpa:
        fixedstr += 'p'
        specifications['component0']['shape']['positionangle'] = bpa
    if x:
        fixedstr += 'x'
        specifications['component0']['shape']['direction']['m0'] = x
    if y:
        fixedstr += 'y'
        specifications['component0']['shape']['direction']['m1'] = y

    # The result is in 'pixels'
    fitskyoutput = ia.fitsky(fit=True,         # Otherwise you won't get 'pixels'
                             includepix=[2],   # Allow negative pixels
                             deconvolve=False, # We don't want a delta function.
                             fixedparams=fixedstr,
                             estimate=specifications)
    ia.close()

    if fitskyoutput['converged'] != True:
        print "Warning!  ia.fitsky() did not converge!"

    result = beam + '_fitres'
    ia.fromarray(result, fitskyoutput['pixels'], csys=b2dcsys.torecord())
    ia.close()
    return result

def quantity_or_nothing(q, qa=None):
    """
    If q is a quantity, or a string which can be turned into a quantity,
    returns q as a quantity.

    Otherwise returns None.
    """
    if qa == None:
        myf = get_global_namespace()
        ia  = myf['qa']

    if(isinstance(q, dict) and q.has_key('unit') and q.has_key('value')):
        return q
    elif isinstance(q, str):
        try:
            return qa.quantity(q)
        except StandardError:
            print q, "could not be converted into a quantity."
    return None