# Template script for combining interferometric (multi-configuration) data with Single Dish Total power data.
# Bits that are to be changed (or at least checked) are marked as 'EDIT THIS'
# This template only deals with ONE science target and produces only ONE spectral cube.
# It has to be correspondingly expanded if several targets are to be imaged and/or several lines/spectral ranges


import re
import os

thesteps = []
step_title = {0: 'Split out and concatenate all science interferometry data',
              1: 'Diagnostic plots: weight vs uv-dist, weight density vs. uv-dist, amp vs uv-dist',
              2: '(Interferometry) continuum imaging',
              3: '(Interferometry) continuum subtraction',
              4: 'Interferometry line imaging',
              5: 'Create single dish cube',
              6: 'Prepare interferometry and single dish cubes for combination',
              7: 'Combine data'}

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


# EDIT THIS ... NOTE: steps 4...7 have to be repeated for every cube that has to be created.

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

# input datasets

interf_12m_TE = 'calibrated/calibrated.12m.ms'     # EDIT THIS tell where the data for the most extended 12m configuration is
interf_12m_TC = ''                                 # EDIT tell where data for the compact configuration is (if there is any)
interf_7m = ['calibrated/calibrated.7m.ms']        # EDIT THIS should be only one (always?)

# EDIT THIS
SD_msNames = ['calibrated/uid___A002_X998e51_X10e9.ms.cal.jy', \
  'calibrated/uid___A002_X998e51_Xda2.ms.cal.jy', \
  'calibrated/uid___A002_X9998b8_X21d9.ms.cal.jy', \
  'calibrated/uid___A002_X9998b8_X251d.ms.cal.jy', \
  'calibrated/uid___A002_X9998b8_X27b3.ms.cal.jy', \
  'calibrated/uid___A002_X9998b8_X29bb.ms.cal.jy', \
  'calibrated/uid___A002_X9998b8_X2c13.ms.cal.jy', \
  'calibrated/uid___A002_X99c183_X24f5.ms.cal.jy', \
  'calibrated/uid___A002_X99c183_X27d9.ms.cal.jy', \
  'calibrated/uid___A002_X99c183_X3d00.ms.cal.jy', \
  'calibrated/uid___A002_X9a8f80_Xb21.ms.cal.jy', \
  'calibrated/uid___A002_X9aca45_X1e73.ms.cal.jy']


scienceTargets = ['Circinus']    # EDIT THIS This may be a list of several objects observed in one SB
phasecenters = ['J2000 14h13m09.906 -65d20m20.4684']    # EDIT THIS This may be a list for several objects observed in one SB

spws = ['0','1','2','3']      # EDIT THIS 
spws_obs = {'0':'0,4,8,12', '1':'1,5,9,13', '2':'2,6,10,14', '3':'3,7,11,15'}      # EDIT THIS
spw_cont = '0:31~1320;1520~1999,1:31~1350;1600~1999,2,3,4:31~1320;1520~1999,5:31~1350;1600~1999,6,7,8:31~1320;1520~1999,9:31~1350;1600~1999,10,11,12:10~1260;1460~1919,13:10~1000;1150~1300;1500~1919,14,15'      # EDIT THIS

# set some general imaging related parameters
sp_outframe = 'bary'      # EDIT THIS take from imaging script of most extended interferometry combination; should be 'lsrk' for galactic, 'bary' for extragalactic
intf_imsize = [320,320]   # EDIT THIS take from imaging script of most extended interferometry combination
intf_cellsize = '0.5arcsec' # EDIT THIS take from imaging script of most extended interferometry combination

interf_12m = []
if (interf_12m_TE != ''):
    interf_12m.append(interf_12m_TE)

if (interf_12m_TC != ''):
    interf_12m.append(interf_12m_TC)

interf_msNames = []
interf_msNames = interf_12m
for myname in interf_7m:
    interf_msNames.append(myname)

interf_name = ''
if (interf_12m_TE != ''):
    interf_name = interf_name+'TE'

if (interf_12m_TC != ''):
    if (interf_name == ''):
        interf_name = 'TC'
    else:
        interf_name = interf_name+'.TC'

if (interf_7m != ''):
    if (interf_name == ''):
        interf_name = '7m'
    else:
        interf_name = interf_name+'.7m'

scienceTargetsString = ''
for i in range(scienceTargets.__len__()):
    if (i == 0):
        scienceTargetsString = scienceTargets[i]
    else:
        scienceTargetsString = scienceTargetsString+','+scienceTargets[i]

###########################################################################
# Provide a number of parameters for the single dish imaging (taken from
# the scriptForSDimaging.py)    EDIT THIS
###########################################################################

SD_spwIds = ['17', '19', '21', '23']

# Define SD imaging parameters - taken from scriptForSDimaging.py

maxsize = 203.0

theorybeam = {}
theorybeam['17'] = 50.7559791418 # mean freq = 114.714470664
theorybeam['19'] = 51.5688455596 # mean freq = 112.906256037
theorybeam['21'] = 57.7638034162 # mean freq = 100.797470664
theorybeam['23'] = 56.7458349288 # mean freq = 102.605685291
  
# the values below were calculated assuming cell = theorybeam[spw]/9.0

sfbeam = {}
sfbeam['17'] = 57.9547742481 # mean freq = 114.714470664
sfbeam['19'] = 58.8206965208 # mean freq = 112.906256037
sfbeam['21'] = 65.4480586513 # mean freq = 100.797470664
sfbeam['23'] = 64.3564966286 # mean freq = 102.605685291


################################################################################
# Prepare interferometric data for combination:
# Split out science interferometry data, concatenate all science interferometry into one measurement set
mystep = 0
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  print "# Concatenating interferometry data"

  os.system('rm -rf calibrated.'+interf_name+'.ms')
  concat(vis = interf_msNames, concatvis = 'calibrated.'+interf_name+'.ms')

  listobs('calibrated.'+interf_name+'.ms', listfile = 'calibrated.'+interf_name+'.ms.listobs')

  print "# Split out science target data"

  interf_science = 'calibrated.'+interf_name+'.science.ms'
  os.system('rm -rf '+interf_science+'*')
  split(vis = 'calibrated.'+interf_name+'.ms', outputvis = interf_science,
      datacolumn = 'data',
      field = scienceTargetsString,
      timebin = '33s', # EDIT THIS be careful here particularly for very long baselines, timebin should not be too long, to avoid uv smearing
      keepflags = True)

  listobs(interf_science, listfile = interf_science+'.listobs')
  

################################################################################
# Prepare interferometric data for combination:
# plot weights vs. uv-distance to check if weights are consistent between
# different interferometry configurations:
# - different 12m array configurations should have comparable weights
# - ACA 7m data should have weights which are lower by a factor of a few (canonical value:5.2) than the corresponding 12m array weights

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

  print "# Plot: weight vs. uv-distance"

  for my_spw in spws:
      plotfilename = 'calibrated.'+interf_name+'.ms.wt_vs_uvdist.'+my_spw+'.png'
      plotms(vis = 'calibrated.'+interf_name+'.ms',
          xaxis = 'uvdist', yaxis = 'wt',
          field = '',
          spw = spws_obs[my_spw],
          coloraxis = 'observation',
          clearplots = True,
          plotfile = plotfilename)

  print "# Plot: weight density vs. uv-distance"

  print "# Plot: amplitude vs. uv-distance (for central? representative? Field)"

  print "# Calculating representative weight ratios:"
  print "#                 TE/7m (should be close to canonical ratio of 5.2);"
  print "#                 TE/TC (if there is more than one 12m configurations)"

##########################################################################
# Start the cleaning
# - continuum (interferometry only!)
# - continuum subtraction (definitely needed if interferometry and SD line cubes are created)
# - line cubes (interferometry and single dish)
##########################################################################

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

  # Running clean: Continuum
  # - needs to be repeated for each science source
  #   (could be done as a loop over science fields, e.g.)
  #   (may need a separate definition of spw_cont for each target)

  target_to_image = scienceTargets[0]     # EDIT THIS
  cont_image_name = target_to_image+'interf.cont'
  os.system('rm -rf '+cont_image_name+'.*')
  default(clean)
  clean(vis = interf_science,
  imagename = cont_image_name,
      field = target_to_image,
      spw = spw_cont,
      mode = 'mfs', outframe = sp_outframe,
      imagermode = 'mosaic',    # 'mosaic' if it's a mosaic OR combines 7m with 12m data
      interactive = F,
      imsize = intf_imsize,
      cell = intf_cellsize,
      phasecenter = phasecenters[0],                                              # EDIT THIS
      mask = ['circle[[160pix,160pix],20pix]','circle[[159pix,112pix],10pix]'],   # EDIT THIS
      weighting = 'briggs', robust = 0.5,
      niter = 1000, threshold = '0.5mJy')                                         # EDIT THIS


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

  # Continuum subtraction
  # - needs to be repeated for each science source
  #   (could be done as a loop over science fields, e.g.)
  #   (may need a separate definition of spw_cont for each target)

  target_to_image = scienceTargets[0]     # EDIT THIS
  cont_image_name = 'interf.cont.'+target_to_image
  os.system('rm -rf '+interf_science+'.cont*')
  uvcontsub(vis = interf_science,
          field = target_to_image,
          fitspw = spw_cont,
          combine = 'spw',
          solint= 'int',
          fitorder = 1,
          want_cont = False)

  os.system('rm -rf calibrated.'+interf_name+'.'+target_to_image+'.ms.contsub')
  os.system('mv calibrated.'+interf_name+'.science.ms.contsub calibrated.'+interf_name+'.'+target_to_image+'.ms.contsub')


################################################################################
# Create line cube(s)
################################################################################

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

  target_to_image = scienceTargets[0]     # EDIT THIS
  sp_phasecenter = phasecenters[0]        # EDIT THIS
  sp_win = '0'                            # EDIT THIS
  sp_line = '12CO'                        # EDIT THIS give some string to identify the cube (e.g., line name, spw, freq, ...)
  sp_mode = 'velocity'                    # EDIT THIS
  sp_restfr = '115271.202MHz'             # EDIT THIS
  sp_start = '160km/s'                    # EDIT THIS
  sp_width = '3km/s'                      # EDIT THIS
  sp_nchan = 180                          # EDIT THIS

  invis = 'calibrated.'+interf_name+'.'+target_to_image+'.ms.contsub'

  # Create interferometric cube

  imbase = target_to_image+'.'+sp_line

  os.system('rm -rf '+imbase+'.'+interf_name+'.*')
  default(clean)
  clean(vis = invis,
    imagename = imbase+'.'+interf_name,
    field = target_to_image,
    spw = spws_obs[sp_win],
    mode = sp_mode, restfreq = sp_restfr, start = sp_start, width = sp_width, nchan = sp_nchan, outframe = sp_outframe,
    imagermode = 'mosaic',
    interactive = F,
    imsize = intf_imsize,
    cell = intf_cellsize,
    phasecenter = sp_phasecenter,
    mask = 'circle[[160pix,160pix],112pix]',                   # EDIT THIS (take from imaging script of most extended config)
    weighting = 'briggs', robust = 0.5,
    niter = 20000, threshold = '24mJy', multiscale = [0,5],    # EDIT THIS (take from imaging script of most extended config), introducing multiscale cleaning might be an idea here...
    usescratch = True)


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

  # Create single dish cube(s)

  # EDIT THIS make sure to use identical spectral specifications as for interferometry!!!

  spw = '17'
  cell = theorybeam[spw]/9.0
  imsize = int(round(maxsize/cell)*2)

  target_to_image = scienceTargets[0]    # EDIT THIS
  sp_phasecenter = phasecenters[0]       # EDIT THIS

  os.system('rm -rf '+imbase+'.TP')
  sdimaging(infiles = SD_msNames,
      field = target_to_image,
      spw = spw,
      mode = sp_mode, restfreq = sp_restfr, start = sp_start, width = sp_width, nchan = sp_nchan, outframe = sp_outframe,
      gridfunction = 'SF',
      convsupport = 6,
      phasecenter = sp_phasecenter,
      imsize = imsize,
      cell = str(cell)+'arcsec',
      overwrite = True,
      outfile = imbase+'.TP')

  # Correct the brightness unit in the image header
  imhead(imagename = imbase+'.TP',
      mode = 'put',
      hdkey = 'bunit',
      hdvalue = 'Jy/beam')

  # Add Restoring Beam Header Information to the Science Image
  ia.open(imbase+'.TP')
  ia.setrestoringbeam(major = str(sfbeam[spw])+'arcsec', minor = str(sfbeam[spw])+'arcsec', pa = '0deg')
  ia.done()



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

  # Prepare TP cube and interferometric cube for combination

  # Regrid TP cube to same spatial grid as interferometric cube

  os.system('rm -rf '+imbase+'.TP.regrid*')
  imregrid(imagename = imbase+'.TP',
         template = imbase+'.'+interf_name+'.image',
         axes = [0,1],
         output = imbase+'.TP.regrid')

  # Trim interferometric cube, TP cube, and primary beam pattern to more sensible size

  trim_box = '42,46,278,274'     # EDIT THIS visually inspect the interferometry image to identify 'empty' edge regions and trim them

  os.system('rm -rf '+imbase+'.interf.subimage*')
  imsubimage(imagename = imbase+'.'+interf_name+'.image',
           outfile = imbase+'.'+interf_name+'.subimage',
             box = trim_box)

  os.system('rm -rf '+imbase+'.TP.regrid.subimage*')
  imsubimage(imagename = imbase+'.TP.regrid',
           outfile = imbase+'.TP.regrid..subimage',
             box = trim_box)

  os.system('rm -rf '+imbase+'.'+interf_name+'.flux.subimage*')
  imsubimage(imagename = imbase+'.'+interf_name+'.flux',
           outfile = imbase+'.'+interf_name+'.flux.subimage',
             box = trim_box)

  # Correct TP cube for interferometric Primary beam coverage

  os.system('rm -rf '+imbase+'.TP.regrid.subimage.dePB*')
  immath(imagename = [imbase+'.TP.regrid.subimage',imbase+'.'+interf_name+'.flux.subimage'],
       expr = 'IM0*IM1',
       outfile = imbase+'.TP.regrid.subimage.dePB')

# Feather TP cube with interferometric cube

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

  outimage = imbase+'.'+interf_name+'.TP.image'
  os.system('rm -rf '+outimage+'*')
  feather(imagename = outimage,
        highres = imbase+'.'+interf_name+'.subimage',
        lowres = imbase+'.TP.regrid.subimage.dePB')

  #print "# Apply primary beam correction, export to fits"

  impbcor(imagename=outimage, pbimage=imbase+'.'+interf_name+'.flux.subimage', outfile=outimage+'.pbcor', overwrite=True) # perform PBcorr
  exportfits(imagename=outimage+'.pbcor', fitsimage=outimage+'.pbcor.fits') # export the corrected image
  exportfits(imagename=imbase+'.'+interf_name+'.flux.subimage', fitsimage=imbase+'.'+interf_name+'.flux.subimage.fits') # export the PB image