#
# This script explores CASA imaging of spectral line data sets using
# ALMA Science Verification data. So far we have looked at imaging
# continuum data. The main difference with line data is that we image
# each frequency separately. CLEAN includes several parameters that
# let you do this. Also, the continuum that we imaged before is still
# present in each data set and will need to be removed before we can
# study the line.
#

# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%
# SET INPUTS AND OUTPUTS
# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%

#
# Pick the input data set and set the root name of the output
# data. These are read into variables that are then fed to parameters
# in the calls to CLEAN and continuum subtraction below.
#

# Path to the calibrated data (pick one of these)

data = "../../calib/NGC3256_Band3_CalibratedData/ngc3256_line_target.ms"
local = "ngc3256_line_target.ms"

#data = "../../calib/TWHYA_BAND7_CalibratedData/TWHydra_corrected.ms"
#local = "TWHydra_corrected.ms"

#data = "../../calib/TWHYA_BAND6_CalibratedData/TWHydra_corrected.ms"
#local = "TWHydra_corrected.ms"

#data = "../../calib/TWHYA_BAND3_CalibratedData/TWHydra_corrected.ms"
#local = "TWHydra_corrected.ms"

#
# List the basic parameters of the input data set. 
#

listobs(vis=data)

# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%
# COPY THE DATA
# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%

#
# We're going to be manipulating the data. To keep things clean, use
# split to make a copy that contains ONLY the line that we care
# about. Note that you need to know something about your data before
# you would know the mapping of SPW to line. Included below are
# "SPLIT" calls to separate out lines for each data set.
#

# NGC 3256 CO
local = 'ngc3256_co.ms'
os.system("rm -rf "+local)
split(vis=data,
      outputvis=local,
      spw='0',
      datacolumn='data')

# NGC 3256 CN
#local = 'ngc3256_cn.ms'
#os.system("rm -rf "+local)
#split(vis=data,
#      outputvis=local,
#      spw='1',
#      datacolumn='data')

# TW HYDRA BAND 3
#local = 'tw_hydra_band3_hcoplus.ms'
#os.system("rm -rf "+local)
#split(vis=data,
#      outputvis=local,
#      spw='1',
#      datacolumn='data')

# TW HYDRA BAND 6 CO AND DCN3-2
#local = 'tw_hydra_band6_co.ms'
#os.system("rm -rf "+local)
#split(vis=data,
#      outputvis=local,
#      spw='0',
#      datacolumn='data')

#local = 'tw_hydra_band6_dcn32.ms'
#os.system("rm -rf "+local)
#split(vis=data,
#      outputvis=local,
#      spw='1',
#      datacolumn='data')

# TW HYDRA BAND 7 CO AND HCO+
#local = 'tw_hydra_band7_co.ms'
#os.system("rm -rf "+local)
#split(vis=data,
#      outputvis=local,
#      spw='2',
#      datacolumn='data')

#local = 'tw_hydra_band7_hcoplus.ms'
#os.system("rm -rf "+local)
#split(vis=data,
#      outputvis=local,
#      spw='0',
#      datacolumn='data')

# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%
# UV CONTINUUM SUBTRACTION
# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%

#
# Find the line free channels by making an average spectrum that
# combines all times and baselines (so maximizes
# signal-to-noise). This won't always work but it will work for our
# test data sets. Note the channel ranges where it looks like there is
# no line. You will feed these in to the uv continuum subtraction in
# the next step.
#

plotms(vis = local,
       spw='0',
       xaxis='channel',
       yaxis='amp',
       avgtime='1e9',
       avgscan=True,
       avgbaseline=True,
       ydatacolumn='data')

# SUGGESTED CALLS TO UV CONTINUUM SUBTRACTION

# Try to figure the line-free channels out on your own from the plots
# above but these are the suggested calls to uvcontsub2 from the CASA
# guides and associated scripts. Notice what happens: these calls
# produce a NEW measurement set with the extension ".contsub". We will
# image that measurement set.

# ... for NGC 3256 CO
uvcontsub(vis = local,
          fitspw='0:20~53;71~120',
          solint ='int',
          fitorder = 0,
          combine='')

# ... for NGC 3256 CN
#uvcontsub(vis = local,
#           fitspw='0:70~120',
#           solint ='int',
#           fitorder = 0,
#           combine='')

# ... for TW Hydra Band 3
#uvcontsub(vis=local,
#           fitspw='0:20~200,0:220~380',
#           solint='int',
#           fitorder=0,
#           combine='')

# ... for TW Hydra Band 6 CO
#uvcontsub(vis=local,
#           fitspw='0:20~160,0:190~380',
#           solint='int',
#           fitorder=0,
#           combine='')

# ... for TW Hydra Band 6 DCN3-2
#uvcontsub(vis=local,
#           fitspw='0:20~80;120~380',
#           solint='int',
#           fitorder=0,
#           combine='')

# ... for TW Hydra Band 7 CO
#uvcontsub(vis=local,
#           fitspw='0:20~200,0:240~380',
#           solint='int',
#           fitorder=0,
#           combine='')

# ... for TW Hydra Band 7 HCOplus
#uvcontsub(vis=local,
#           fitspw='0:20~150,0:190~380',
#           solint='int',
#           fitorder=0,
#           combine='')

# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%
# IMAGE THE CONTINUUM SUBTRACTED DATA
# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%

#
# Set the parameters that CLEAN will use to grid and deconvolve the
# u-v data into your final image. The main difference between this and
# the continuum case is that we will specify the third (spectral) axis
# of the data. We will see several ways to do this.

default('clean')

# image the continuum-subtracted data

vis=local+'.contsub'

# specify the output
# (NB we cheat on the name here using a minor python trick to strip the trailing '.ms')

imagename=str.rstrip(local, '.ms')
print imagename

# we want to do this interactively
interactive=True

# ... STOPPING CONDITIONS
thresh = '1mJy'
niter = 10000

# ... CELL SIZE

# we want ~3-5 cells per synthesized beam. These are recommended for
# the different examples below, as in previous scripts

# cell = '0.3arcsec'
# cell = '0.5arcsec'
cell = '1.0arcsec'

# ... IMAGE SIZE

# we want to image out to the primary beam by default, as before we
# pick the size appropriate to your image.

# imsize = 100
imsize = 300

imagermode = 'csclean'
psfmode='hogbom'

# ... U-V WEIGHTING
# weighting = 'natural'
weighting = 'briggs'
robust = 0.5

# ----------------------
# THE LINE-SPECIFIC PART
# ----------------------

# You can specify your spectral range and gridding in a few ways. The
# CASA guides adopt different approaches, see them for more
# details. 

# For each of these cases you will want to start by specifying the
# rest frequency:

# NGC 3256 CO 1-0
restfreq = '115.271201800GHz'

# NGC 3256 CN N=1-0 J=1/2-1/2
#restfreq = '113.17049GHz'

# NGC 3256 CN N=1-0 J=3/2-1/2
#restfreq = '113.48812GHz'

# TW Hydra HCO+ 1-0
#restfreq = '89.188526GHz'

# TW Hydra CO 2-1
#restfreq = '230.538GHz'

# TW Hydra DCN
#restfreq = '217.2385GHz'

# TW Hydra CO 3-2
#restfreq='345.79599GHz'

# TW Hydra HCO+ 4-3
#restfreq='356.7342GHz'

#
# APPROACH 1: mode 'channel'. Pick a channel range via spw
# selection. You can refine this with a specified start channel,
# width, and number of channels to image. See examples in the NGC 3256
# CASA guide.
#

# ... Here for NGC 3256 CO
mode='channel'
start=''
spw='0:38~87'
width=''
nchan=50

# ... Here for NGC 3256 CN "low" line
#mode='channel'
#start=''
#spw = '0:29~54'
#width=''
#nchan=26

# ... Here for NGC 3256 CN "high" line
#mode='channel'
#start=''
#spw = '0:50~76'
#width=''
#nchan=27

#
# APPROACH 2: mode 'velocity'. Pick a starting velocity, step size,
# and number of channel maps to image. See examples in the TW Hydra
# CASA guide. CASA handles the appropriate regridding.
#

# ... for TW Hydra Band 7 (either HCO+ or CO 3-2)
#mode='velocity'
#start='-6km/s'
#width='1.2km/s'
#nchan=15
#outframe='LSRK'

# ... for TW Hydra Band 6
#mode='velocity'
#start='-6km/s'
#width='1.7km/s'
#nchan=15
#outframe='LSRK'

# ... for TW Hydra Band 3
#mode='velocity'
#start='-6km/s'
#width='2.1km/s'
#nchan=10
#outframe='LSRK'

# ----------------------
# FIRE OFF CLEAN
# ----------------------

# erase previous images (skip this step if you want to continue cleaning)
os.system('rm -rf '+imagename+'.image')
os.system('rm -rf '+imagename+'.model')
os.system('rm -rf '+imagename+'.psf')
os.system('rm -rf '+imagename+'.mask')
os.system('rm -rf '+imagename+'.residual')
os.system('rm -rf '+imagename+'.flux')

# review inputs
inp

# execute CLEAN
go

# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%
# VIEW
# =%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%=%

#
# Display the resulting image. Note that if you are using nterms > 1
# this image is instead called .image.tt0
#

viewer(image_out+'.image')

#
# You may also want to inspect the residual, model, flux (primary beam
# response), and (if using nterms > 1) alpha images.
#