import numpy as np
from astropy.io import fits
import tempfile
import warnings


def make_casa_mask(SpecCube, outname, append_to_image=True,
                   img=None):
    '''
    Takes a SpectralCube object as an input. Outputs the mask in a CASA
    friendly form.

    Parameters
    ----------
    SpecCube : SpectralCube
        SpectralCube object containing mask.
    outname : str
        Name of the outputted mask file.
    append_to_image : bool, optional
        Appends the mask to a given image.
    img : str, optional
        Image to be appended to. Must be specified if append_to_image is
        enabled.
    '''

    try:
        from taskinit import ia
        from tasks import immath
    except ImportError:
        print("Cannot import casac. Must be run in a CASA environment.")

    # Get the header info from the image
    # There's not wcs_astropy2casa (yet), so create a temporary file for
    # CASA to open.
    temp = tempfile.NamedTemporaryFile()
    # CASA is closing this file at some point so set it to manual delete.
    temp2 = tempfile.NamedTemporaryFile(delete=False)

    hdu = fits.PrimaryHDU(header=SpecCube.header,
                          data=np.empty(SpecCube.shape))

    hdu.writeto(temp.name)

    ia.fromfits(infile=temp.name, outfile=temp2.name, overwrite=True)

    temp.close()

    cs = ia.coordsys()

    ia.close()

    temp2.close()

    mask_arr = SpecCube.mask.include()

    # Transpose to match CASA axes
    mask_arr = mask_arr.T

    # CASA doesn't like bool? Using floats for now...
    ia.newimagefromarray(outfile=outname,
                         pixels=mask_arr.astype('float64'))

    ia.open(outname)
    ia.setcoordsys(cs.torecord())

    ia.close()

    if append_to_image:
        if img is None:
            raise TypeError("img argument must be specified to append the mask.")
        warnings.warn("Image appending not working yet.")
        # immath(imagename=img, mode='evalexpr', expr='IM0',
        #        outfile=img+"_test", mask='mask('+outname+')')

@@ -73,16 +73,64 @@
 rnings%0A%0A
+from ..wcs_utils import add_stokes_axis_to_wcs%0A%0A
 %0Adef mak
@@ -209,16 +209,33 @@
 img=None
+, add_stokes=True
 ):%0A    '
@@ -1174,16 +1174,376 @@
 e)%0A%0A    
+# Grab wcs%0A    # Optionally re-add on the Stokes axis%0A    if add_stokes:%0A        wcs = SpecCube.wcs%0A        new_wcs = add_stokes_axis_to_wcs(wcs, wcs.wcs.naxis)%0A        header = new_wcs.to_header()%0A        shape = (1,) + SpecCube.shape%0A    else:%0A        # Just grab the header from SpecCube%0A        header = SpecCube.header%0A        shape = SpecCube.shape%0A%0A    
 hdu = fi
@@ -1563,25 +1563,16 @@
 (header=
-SpecCube.
 header,%0A
@@ -1611,25 +1611,16 @@
 p.empty(
-SpecCube.
 shape))%0A
@@ -1792,24 +1792,24 @@
 p2.close()%0A%0A
-
     mask_arr
@@ -1832,24 +1832,109 @@
 .include()%0A%0A
+    # Reshape mask with possible Stokes axis%0A    mask_arr = mask_arr.reshape(shape)%0A%0A
     # Transp

Add option to re-add Stokes axis.