import numpy as np
import os

#===== MAIN CODE ====#
#--- Initial parameters ---#
myMS = 'Z_CMa_TM2.split.cal' #--- This is the calibrated measurement set with the science target (generated during step 1 of scriptForImagingZCMa.py)

#--- Parameters for images shown Imaging Tutorial Extra ---#
contSPW = '0,1,2,3,4,5,6,7,8:300~959,9:0~85;115~959,10'
myCell = '0.152arcsec'   #--- cell/pixel size of your images. A string of format 'X.xxxarcsec'
myImsize = 300  #--- number of x by y pixels in your image. Format integer

myThreshold = '0.226mJy' #--- 3x your calculated theoretical threshold. A string of format 'Y.ymJy'

myMask = 'Z_CMa_ALL_SPW_CONTINUUM.mask' # generated during the interactive cleaning of the continuum in step 3 of scriptForImagingZCMa.py

#--- Logic switches ---#
""" Set the step you wish to run to = True and all others to = False """
#For symplicity, all the steps use non-interactive mode and the same mask for cleaning (defined with parameter myMask).
step0 = False       #--- Image the continuum with few iterations (undercleaning).
step1 = False       #--- Image the continuum with natural weighting (better sensitivity, worst angular resolution).
step2 = False       #--- Image the continuum with uniform weighting (better angular resolution, worst sensitivity).
step3 = False       #--- Image the continuum with uvtaper (emphysizes extended emission over a certain scale).
step4 = False       #--- Image the continuum with uvrange (removes emission over/under a certain scale).
step5 = False       #--- Image the continuum with hogbom deconvolver (to compare with multiscale deconvolver).
step6 = False       #--- Mask moments 0, 1 and 2 ussing moment 0 emission as threshold (uses the moms maps made during step 5 of scriptForImagingZCMa.py).


#-------------------------------#
#--- STEP 0: Image the continuum with few iterations (undercleaning) ---#
if step0:
	niter=50
	print '  >> Cleaning Z-CMa continuum with iterations: '+str(niter)
	tclean(vis = myMS,
          imagename = myMS+'_ALL_SPW_CONTINUUM_niter='+str(niter),
          spw=contSPW,
          threshold=myThreshold,
          cell=[myCell],
          imsize=[myImsize,myImsize], #should cover about the FoV
          outframe='LSRK',
          niter=niter,
          deconvolver='multiscale',
          scales=[0,5,15],# inpixels equivalent to point source and 1, 3 times the beam.
          weighting = 'briggs',
          robust=0.5,
          pbcor=True,
          specmode='mfs',
          restoringbeam='common',
          nterms=1,
          chanchunks=-1,
          gridder='standard',
	  usemask='user',
	  mask=myMask
          )

#-------------------------------#
#--- STEP 1: Image the continuum with natural weighting (better sensitivity, worst angular resolution) ---#
if step1:
	weighting='natural'
	print '  >> Cleaning Z-CMa continuum with weighting: '+ weighting
	tclean(vis = myMS,
          imagename = myMS+'_ALL_SPW_CONTINUUM_'+weighting,
          spw=contSPW,
          threshold=myThreshold,
          cell=[myCell],
          imsize=[myImsize,myImsize], #should cover about the FoV
          outframe='LSRK',
          niter=10000,
          deconvolver='multiscale',
          scales=[0,5,15],# inpixels equivalent to point source and 1, 3 times the beam.
          weighting = weighting,
          pbcor=True,
          specmode='mfs',
          restoringbeam='common',
          nterms=1,
          chanchunks=-1,
          gridder='standard',
	  usemask='user',
	  mask=myMask
          )

#-------------------------------#
#--- STEP 2: Image the continuum with uniform weighting (better angular resolution, worst sensitivity) ---#
if step2:
	weighting='uniform'
	print '  >> Cleaning Z-CMa continuum with weighting: '+ weighting
	tclean(vis = myMS,
          imagename = myMS+'_ALL_SPW_CONTINUUM_'+weighting,
          spw=contSPW,
          threshold=myThreshold,
          cell=[myCell],
          imsize=[myImsize,myImsize], #should cover about the FoV
          outframe='LSRK',
          niter=10000,
          deconvolver='multiscale',
          scales=[0,5,15],# inpixels equivalent to point source and 1, 3 times the beam.
          weighting = weighting,
          pbcor=True,
          specmode='mfs',
          restoringbeam='common',
          nterms=1,
          chanchunks=-1,
          gridder='standard',
	  usemask='user',
	  mask=myMask
          )

#-------------------------------#
#--- STEP 3: Image the continuum with uvtaper (emphysizes extended emission over a certain scale). ---#
if step3:
	uvtaper='200klambda'
	print '  >> Cleaning Z-CMa continuum with uvtaper: '+uvtaper
        tclean(vis = myMS,
          imagename = myMS+'_ALL_SPW_CONTINUUM_uvtaper='+uvtaper,
          spw=contSPW,
          threshold=myThreshold,
          cell=[myCell],
          imsize=[myImsize,myImsize], #should cover about the FoV
          outframe='LSRK',
          niter=10000,
	  uvtaper=uvtaper,
          deconvolver='multiscale',
          scales=[0,5,15],# inpixels equivalent to point source and 1, 3 times the beam.
          weighting = 'briggs',
          robust=0.5,
          pbcor=True,
          specmode='mfs',
          restoringbeam='common',
          nterms=1,
          chanchunks=-1,
          gridder='standard',
          usemask='user',
	  mask=myMask
          )

#-------------------------------#
#--- STEP 4: Image the continuum with uvrange (removes extended emission under/over a certain scale). ---#
if step4:
	uvrange='>200klambda'
	print '  >> Cleaning Z-CMa continuum with uvrange: '+uvrange
        tclean(vis = myMS,
          imagename = myMS+'_ALL_SPW_CONTINUUM_uvrange='+uvrange,
          spw=contSPW,
          threshold=myThreshold,
          cell=[myCell],
          imsize=[myImsize,myImsize], #should cover about the FoV
          outframe='LSRK',
          niter=10000,
	  selectdata=True,
	  uvrange=uvrange,
          deconvolver='multiscale',
          scales=[0,5,15],# inpixels equivalent to point source and 1, 3 times the beam.
          weighting = 'briggs',
          robust=0.5,
          pbcor=True,
          specmode='mfs',
          restoringbeam='common',
          nterms=1,
          chanchunks=-1,
          gridder='standard',
          usemask='user',
	  mask=myMask
          )

#-------------------------------#
#--- STEP 5: Image the continuum with hogbom deconvolver (to compare with multiscale deconvolver) ---#
if step5:
	deconvolver='hogbom'
	print '  >> Cleaning Z-CMa continuum with deconvolver: '+deconvolver
	tclean(vis = myMS,
          imagename = myMS+'_ALL_SPW_CONTINUUM_'+deconvolver,
          spw=contSPW,
          threshold=myThreshold,
          cell=[myCell],
          imsize=[myImsize,myImsize], #should cover about the FoV
          outframe='LSRK',
          niter=10000,
          deconvolver=deconvolver,
          weighting = 'briggs',
          robust=0.5,
          pbcor=True,
          specmode='mfs',
          restoringbeam='common',
          nterms=1,
          chanchunks=-1,
          gridder='standard',
	  usemask='user',
	  mask=myMask
          )

#-------------------------------#
#--- STEP 6: Mask moments 0, 1 and 2 ussing moment 0 emission as threshold (uses the moms maps made during step 5 of scriptForImagingZCMa.py). ---#
if step6:
	imagename='Z_CMa_CO_cube.image_red' #rootname
	threshold=0.07*5
	mask = "'"+imagename+".integrated'>="+str(threshold) #masking using integrated moment
	print '  >> Blanking moments using: '+imagename+'.integrated > '+str(round(threshold,4))
	for end_name in ['.integrated', '.weighted_coord','.weighted_dispersion_coord']:
		immath(imagename=imagename+end_name, expr='IM0', mask=mask, outfile=imagename+end_name+'_B')