# ALMA Data Reduction Script
# Created by A.Avison for the Durham ALMA Data Reduction Workshop
# Dublin 14-16 Nov 2018
# 
# This script is based upon and behaves like the typical output of
# the ALMA imaging script generator. But contains steps not normally
# seen in such a script.
#
# This script is designed to accompany the Tutorial which can be found 
# here: 
#
# As part of this tutorial, tutees are intended to fill in task parameters
# throughout this script to make it run correctly.

# Imaging


""" Imaging ALESS49                                                       """
""" Following Wardlow et al. 18 we will image +/- 8arcsec around ALESS49  """
""" SPW 1 which covers the detected CO line .                             """


#=== STEP SET-UP FOR USING THIS SCRIPT ===#
thesteps = []
step_title = {0: 'Full spectral window imaging',
              1: 'Continuum subtraction',
              2: 'Image a spectral line from within spw 1: CO',
              3: 'Immoments ',
              4: 'Export images to FITS file',
              }

if 'applyonly' not in globals(): applyonly = False
try:
  print 'List of steps to be executed ...', mysteps
  thesteps = mysteps
except:
  print 'global variable mysteps not set.'
if (thesteps==[]):
  thesteps = range(0,len(step_title))
  print 'Executing all steps: ', thesteps

# The Python variable 'mysteps' will control which steps
# are executed when you start the script using
#   execfile('scriptForCalibration.py')
# e.g. setting
#   mysteps = [2,3,4]# before starting the script will make the script execute
# only steps 2, 3, and 4
# Setting mysteps = [] will make it execute all steps.


#=============================================================================#
#========================== THE SCRIPT PROPER ================================#
#=============================================================================#

#=== VARIBLES FOR USE THROUGHOUT ===#
thevis = 'ALESS49_ALL_EBS.ms'                           #- The measurement set described in the tutorials
phCen = 'J2000 03:31:24.595000 -27.50.42.60000 '        #- The phase centre of the observation above
target = 'ALESS49'                                      #- The target name
restnus = ['86798.625MHz', '87576.9377MHz', '96626.750MHz', '98397.250MHz'] #- The central frequencies of each SPW

#==========================================================#
#               SELECT LINE FREE CHANNELS                  #
#==========================================================#

""" To make a continuum subtracted images we need to select all 
line free channels from each spectral window (SPW). This is fairly
 easy for ALESS49 we only have one line detectd. 

This step will great full cubes of the each of the 4 SPWs so you
can isolate the line free channels and fill in uvcontsubs fitspw
parameter in Step 1.
"""

mystep = 0
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
  
  for ix in range(4):       #- This `for' loop cycles through each of the 4 spectral windows to image each in turn
    if ix == 0 or ix ==1:   #- This `if' statement set which channel to start the full cube on as there are bad edge...
        startCh=9           #- ...channels we want to avoid.
    else:
        startCh=10       
     
    os.system('rm -rf '+target+'_line_SPW'+str(ix))

    tclean(vis=thevis,
        imagename=target+'_line_SPW'+str(ix),
        field=target,
        phasecenter=phCen,
        restfreq=restnus[ix],
        spw=str(ix),
        threshold='',                     #=========== FILL THIS IN
        imsize=[,],                       #=========== FILL THIS IN
        cell='0.2arcsec',
        niter=10000,
        width=1,
        start=startCh,
        outframe='BARY', 
        nchan=109,
        deconvolver='hogbom', 
        weighting='briggs',
        robust = 0.5,
        pbcor=True, 
        specmode='cube', 
        gridder='standard', 
        interactive=True
        )

#==========================================================#
#               SELECT LINE FREE CHANNELS                  #
#==========================================================#
""" The fitspw parameters allows the CASA task uvcontsub to fit a 
continuum level and subtract it directly from the visibilities. 
"""
mystep = 1
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -rf '+thevis+'.contsub')

  uvcontsub(vis = thevis,
      field = target,
      fitspw = '',                     #=========== FILL THIS IN
      want_cont=False
      ) 

#=======================================================================#
# Image data cube CO line
#=======================================================================#
""" Now we switch to using the continuum subtracted MS to 
generated a spectral cube of the CO emission seen in 
ALESS49.

Again you will need to fill in the cell, imsize and
threshold values. In addition to this you will need to set 
the start, nchan and width parameters to deinfe the cubes start 
and end points.
"""
mystep = 2
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]
   
  os.system('rm -rf '+target+'_line_SPW1_contsubbed')

  tclean(vis=thevis+'.contsub',
        imagename=target+'_line_SPW1_contsubbed',
        field=target,
        phasecenter=phCen,
        restfreq=restnus[1],
        spw='1',
        threshold='',                     #=========== FILL THIS IN
        imsize=[,],                       #=========== FILL THIS IN
        cell='',                          #=========== FILL THIS IN
        niter=10000,
        width=,                           #=========== FILL THIS IN
        start=,                           #=========== FILL THIS IN
        outframe='BARY', 
        nchan=,                           #=========== FILL THIS IN
        deconvolver='hogbom', 
        weighting='briggs',
        robust = 0.5,
        pbcor=True, 
        specmode='cube', 
        gridder='standard', 
        interactive=True
        )

#=======================================================================#
# Make momentum map of the SO2 line
#=======================================================================#
""" When you have a cube from steps 2 we proceed to generate
0th to 2nd moment images of the data. 

You will need to fill in values for the excludepix and chans 
parameters. (This task runs fairly quickly, so you can play around 
with these inputs a bit).

"""
mystep = 3
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]


    immoments(imagename=target+'_line_SPW1_contsubbed.image',#- The name of the file you created in Step 2
        moments=[0,1,2],
        outfile='ALESS49_CO.I.manual.image.mom',#- Or whatever you would like to call the output image
        excludepix=[,],       #=========== FILL THIS IN
        chans=''             #=========== FILL THIS IN
        )

#=======================================================================#
# Export images to FITS format
#=======================================================================#
""" This example looks for any CASA image files ending in .image or 
.image.pbcor in the current directory and converts them to FITS file, 
dropping all degenerate axis (axis length=1)

Note it misses the immoments outputs. Add some code (or do it manually
in CASA) to export these files to FITS too.
"""

mystep = 4
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  for item in os.listdir('.'):
        if item.endswith('.image') or item.endswith('.image.pbcor') or item.startswith('ALESS49_CO.I.'):
               exportfits(imagename = item,
                fitsimage = item+'.fits',
                overwrite = True,
                dropdeg = True
                )