# This script describes the imaging for the ALMA SV Band 4 data on IRAS16293
#
# Tested in CASA 4.5.0. Meant to be run interactively.

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


########################################
# Check CASA version

import re


if re.search('^4.5.0', casadef.casa_version) == None:
 sys.exit('ERROR: PLEASE USE THE SAME VERSION OF CASA THAT YOU USED FOR GENERATING THE SCRIPT: 4.5.0')


########################################
# Getting a list of ms files to image

import glob

vislist=glob.glob('uid___A002_X867766_Xa7.ms.split.cal')

########################################
# Removing pointing table

# This step removes the pointing table from the data to avoid
# a bug with mosaics in CASA 4.2.2

for vis in vislist:
    tb.open( vis + '/POINTING',
            nomodify = False)
    a = tb.rownumbers()
    tb.removerows(a)
    tb.close()


###################################
# Splitting off science target data

for vis in vislist:
    listobs(vis=vis)

for vis in vislist:
    sourcevis=vis+'.source'
    rmtables(sourcevis)
    os.system('rm -rf ' + sourcevis + '.flagversions')
    split(vis=vis,
          intent='*TARGET*', # split off the target sources
          outputvis=sourcevis,
          datacolumn='data')
    # Check that split worked as desired.
    listobs(vis=sourcevis) 

############################################
# Rename and backup data set

os.system('mv -i ' + sourcevis + ' ' + 'calibrated_final.ms')

# At this point you should create a backup of your final data set in
# case the ms you are working with gets corrupted by clean. 

os.system('cp -ir calibrated_final.ms calibrated_final.ms.backup')


##################################################
# Create an Averaged Continuum MS

finalvis='calibrated_final.ms' # This is your output ms from the data
                               # preparation script.

# Use plotms to identify line and continuum spectral windows
plotms(vis=finalvis, xaxis='channel', yaxis='amplitude',
       ydatacolumn='data',
       avgtime='1e8', avgscan=True, avgchannel='4', # you should only lightly average over frequency
#       avgbaseline=True, # try if you don't see anything with the time and frequency averaging
       iteraxis='spw' )


# Set spws to be used to form continuum
contspws = '0'

# If you have complex line emission and no dedicated continuum
# windows, you will need to flag the line channels prior to averaging.
flagmanager(vis=finalvis,mode='save',
            versionname='before_cont_flags')

# Flag the "line channels"
flagchannels='0:840~1360,0:1600~1840,0:2200~2480,0:2720~2860,0:3240~3400,0:3440~3760' 

flagdata(vis=finalvis,mode='manual',
          spw=flagchannels,flagbackup=False)

# check that flags are as expected, NOTE must check reload on plotms
# gui if its still open.
plotms(vis=finalvis,yaxis='amp',xaxis='channel',
       avgchannel='4',avgtime='1e8',avgscan=True,iteraxis='spw') 

# Average the channels within spws
contvis='calibrated_final_cont.ms'
rmtables(contvis)
os.system('rm -rf ' + contvis + '.flagversions')
split(vis=finalvis,
      spw=contspws,      
      outputvis=contvis,
      width=[3840], # number of channels to average together. change to appropriate value for each spectral window in contspws (use listobs to find) and make sure to use the native number of channels per SPW (that is, not the number of channels left after flagging any lines)
      datacolumn='data')

# Note: There is a bug in split that does not average the data
# properly if the width is set to a value larger than the number of
# channels in an SPW. Specifying the width of each spw (as done above)
# is necessary for producing properly weighted data.

# If you flagged any line channels, restore the previous flags
flagmanager(vis=finalvis,mode='restore',
            versionname='before_cont_flags')

# Inspect continuum for any problems
plotms(vis=contvis,xaxis='uvdist',yaxis='amp',coloraxis='spw')

# #############################################
# Image Parameters

# source parameters
# ------------------

field='3' 
imagermode='csclean'
cell='0.15arcsec' # cell size for imaging.
imsize = [256,256] # size of image in pixels.

# velocity parameters
# -------------------

start=0 # start velocity. See science goals for appropriate value.
width=4 # velocity width. See science goals.
nchan = 960 # number of channels. See science goals for appopriate value.
outframe='bary' # velocity reference frame. See science goals.
veltype='radio' # velocity type. See note below.

# imaging control
# ----------------

# The cleaning below is done interactively, so niter and threshold can
# be controlled within clean. 

weighting = 'briggs'
robust=0.5
niter=1000
threshold = '0.0mJy'

#############################################
# Imaging the Continuuum

# Set the ms and continuum image name.
contvis = 'calibrated_final_cont.ms'         
contimagename = 'calibrated_final_cont_image'

# If necessary, run the following commands to get rid of older clean
# data.

#clearcal(vis=contvis)
#delmod(vis=contvis)

for ext in ['.flux','.image','.mask','.model','.pbcor','.psf','.residual','.flux.pbcoverage']:
    rmtables(contimagename+ext)

clean(vis=contvis,
      imagename=contimagename,
      field=field,
      mode='mfs',
      psfmode='clark',
      imsize = imsize, 
      cell= cell, 
      weighting = weighting, 
      robust = robust,
      niter = niter, 
      threshold = threshold, 
      interactive = True,
      mask='',
      imagermode = imagermode)

# Ran 15 iterations of clean using the continuum mask.  Final RMS in the spw=0 map is 0.76 mJy

# If you'd like to redo your clean, but don't want to make a new mask
# use the following commands to save your original mask. This is an optional step.
#contmaskname = 'cont.mask'
##rmtables(contmaskname) # if you want to delete the old mask
#os.system('cp -ir ' + contimagename + '.mask ' + contmaskname)

########################################
# Continuum Subtraction for Line Imaging (CH3OH)

# CH3OH is shown as an example. For CH3CN you can follow a similar procedure to what follows
# but for spectral window 2.

fitspw = '0:0~839;1361~1599;1841~2199;2481~2719;2861~3239;3401~3439;3761~3839' # line-free channel for fitting continuum
linespw = '0' # line spectral windows. You can subtract the continuum from multiple spectral line windows at once.

finalvis='calibrated_final.ms'

uvcontsub(vis=finalvis,
          spw=linespw, # spw to do continuum subtraction on
          fitspw=fitspw, # select spws to fit continuum. exclude regions with strong lines.
          combine='spw', 
          solint='int',
          fitorder=1,
          want_cont=False) # This value should not be changed.

# NOTE: Imaging the continuum produced by uvcontsub with
# want_cont=True will lead to extremely poor continuum images because
# of bandwidth smearing effects. For imaging the continuum, you should
# always create a line-free continuum data set using the process
# outlined above.

linevis = finalvis+'.contsub'

##############################################
# Image CH3OH line emission 

finalvis = 'calibrated_final.ms'
linevis = finalvis + '.contsub' # uncomment if continuum subtracted
# linevis = linvis + '.selfcal' # uncommment if self-calibrated

sourcename ='IRAS16293' # name of source
linename = 'Methanol' # name of transition (see science goals in OT for name) 
lineimagename = sourcename+'_'+linename+'_image' # name of line image

restfreq='157.17902GHz' # Typically the rest frequency of the line of
                        # interest. If the source has a significant
                        # redshift (z>0.2), use the observed sky
                        # frequency (nu_rest/(1+z)) instead of the
                        # rest frequency of the
                        # line.
# spw='1' # uncomment and replace with appropriate spw if necessary.

# If necessary, run the following commands to get rid of older clean
# data.

#clearcal(vis=linevis)
#delmod(vis=linevis)

for ext in ['.flux','.image','.mask','.model','.pbcor','.psf','.residual','.flux.pbcoverage']:
    rmtables(lineimagename + ext)

clean(vis=linevis,
      imagename=lineimagename, 
      field=field,
      mode='channel',  #change to mode = 'channel' to accomodate the above channel selection
      start=start,
      width=width,
      nchan=nchan, 
      outframe=outframe, 
      veltype=veltype, 
      restfreq=restfreq, 
      niter=niter,  
      threshold=threshold, 
      interactive=True,
      mask='',
      cell=cell,
      imsize=imsize, 
      weighting=weighting, 
      robust=robust,
      imagermode=imagermode)
      

# Ran 9 iterations of clean down to a final map RMS of 2.5 mJy/beam

# If you'd like to redo your clean, but don't want to make a new mask
# use the following commands to save your original mask. This is an
# optional step.
# linemaskname = 'line.mask'
## rmtables(linemaskname) # uncomment if you want to overwrite the mask.
# os.system('cp -ir ' + lineimagename + '.mask ' + linemaskname)
#
#
##############################################
# Apply a primary beam correction

import glob

myimages = glob.glob("*.image")

rmtables('*.pbcor')
for image in myimages:
    impbcor(imagename=image, pbimage=image.replace('.image','.flux'), outfile = image.replace('.image','.pbcor'))

##############################################
# Export the images

import glob

myimages = glob.glob("*.image")
for image in myimages:
    exportfits(imagename=image, fitsimage=image+'.fits',overwrite=True)

myimages = glob.glob("*.pbcor")
for image in myimages:
    exportfits(imagename=image, fitsimage=image+'.fits',overwrite=True)

myimages = glob.glob("*.flux")
for image in myimages:
    exportfits(imagename=image, fitsimage=image+'.fits',overwrite=True)


#Moment 0 map of CH3OH - 6(0,6) - 6(-1,6) transition (157.04862 GHz):
os.system('rm -rf IRAS16293_Methanol_image.image.mom0.*')
immoments(imagename='IRAS16293_Methanol_image.image',
          outfile='IRAS16293_Methanol_image.image.mom0',
          chans='412~450',
          moments=0)

exportfits(imagename='IRAS16293_Methanol_image.image.mom0', fitsimage='IRAS16293_Methanol_image.image.mom0.fits')

##############################################
# Analysis

# For examples of how to get started analyzing your data, see 
#     http://casaguides.nrao.edu/index.php?title=TWHydraBand7_Imaging_4.2