# Script to image eMERLIN data with CASA 6 or later

import os, sys

if (sys.platform == 'darwin'):
	mac = True
	print('Mac detected')
else:
	mac = False

# 1252+5634.ms contains target 3C277.1 with all external calibration
# applied, including the phase-reference solutions 

target = '1252+5634' # 3C277.1
antref='Mk2'          # refant - usually - change if known to be bad.

gui=False              # whether to show plotms plots on screen

# *** 
# *** 

# Enter step(s) as runsteps

# Fill in the missing parameters and values

#############################
thesteps = []
step_title = {1: 'Check phase-reference calibrated target data',
			  2: 'Make first target image',
			  3: 'Get image information',
			  4: 'Plot model column',
			  5: 'Choose interval for first phase self-cal',
			  6: 'First phase self-cal',
			  7: 'Apply solutions and image',
			  8: 'Second phase self-cal',
			  9: 'Apply new phase calibration and image',
			  10: 'Chose interval for amp self-cal',
			  11: 'Amplitude self-calibration',
			  12: 'Apply amp self-cal and image',
			  13: 'Reweight according to sensitivity',
			  14: 'Multiscale clean',
			  15: 'Clean with uvtaper',
			  16: 'Clean with Briggs robust weighting'}

try:
	print('List of steps to be executed ...', runsteps)
	thesteps = runsteps
except:
	print('global variable runsteps not set')

if (thesteps==[]):
	thesteps = list(range(0,len(step_title)))
	print('Executing all steps: ', thesteps)

### 1) Check target MS
mystep = 1

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	os.system('rm target.ms.listobs')
	listobs(vis=target+'.ms',
			listfile='target.ms.listobs')

	# plot amp against uv distance
	plotms(vis=target+'.ms',
		   xaxis='uvwave',
		   yaxis='amp',
		   correlation='RR',
		   avgchannel='16',
		   showmajorgrid=True,
		   showminorgrid=True,
		   spw='3',
		   plotfile='Step_01_3C277.1_uvdist.png',
		   overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)


### 2) Image target
# Use known properties of 3C277.1 as guidance in initial masking.
# Expand mask as more emission appears above decreasing noise
# in successive cycles.

mystep = 2

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	os.system('rm -rf target0.*')
	if mac == False:
		tclean(vis=target+'.ms',
			   imagename='target0',
			   specmode='mfs',
			   niter=1500,
			   cycleniter=300,
			   threshold=0,
			   imsize=['**','**'],
			   cell='**',
			   weighting='**',
			   deconvolver='**',
			   savemodel='modelcolumn',
			   interactive=True)
	else:
		iclean(vis=target+'.ms',
				   imagename='target0',
				   specmode='mfs',
				   niter=1500,
				   cycleniter=300,
				   threshold=0,
				   imsize=['**','**'],
				   cell='**',
				   weighting='**',
				   deconvolver='**',
				   savemodel='modelcolumn')
		ft(vis=target+'.ms', 
		   model='target0.model',
		   usescratch=True) ##added as iclean doesn't include model after last iteration

mystep = 3

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	rms1=imstat(imagename='target0.image',
				  box='**,**,**,**')['rms'][0]
	peak1=imstat(imagename='target0.image',
				  box='**,**,**,**')['max'][0]

	print(('target0.image rms  = %7.3f mJy/bm' % (rms1*1000.)))
	print(('target0.image peak = %7.3f mJy/bm' % (peak1*1000.)))
	print(('target0.image S/N  = %7.3f' % (peak1/rms1)))

	if mac==False:
		os.system('rm Step_03_target0.image.png')
		imview(raster={'file': 'target0.image',
					'colormap':'Rainbow 2',
					'scaling': -2,
					'range': [-1.*rms1, 100.*rms1],
					'colorwedge': True},
			   zoom=1,
			   out='Step_03_target0.image.png')
	else:
		print('IMPORTANT: MAC OS detected .. Please view target0.image with CARTA')
		
mystep = 4

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	plotms(vis=target+'.ms',
		   xaxis='uvdist',
		   yaxis='amp',
		   ydatacolumn='model',
		   correlation='RR,LL',
		   avgchannel='16',	  
		   coloraxis='spw',
		   plotfile='',
		   overwrite=True, showgui=gui)


### 3) Chose interval for first phase self-cal
# If you like, inspect  phases against time to help decide solint.
# However, phase slopes are also due to source structure.
# Look at the shortest baseline (Mk2&Pi)
# But, inadvisable to use a very short solint while model imperfect.
# Half a scan length seems OK.

mystep = 5

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	plotms(vis=target+'.ms',
		   xaxis='**',
		   yaxis='**',
		   antenna='**',
		   correlation='**,**',
		   avgchannel='**',
		   coloraxis='spw',
		   plotfile='',
		   timerange='',
		   overwrite=True, showgui=gui)

### 4) First phase self-cal

mystep = 6

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	os.system('rm -rf target.p1')
	gaincal(vis=target+'.ms',
			calmode='**',
			caltable='target.p1',
			field=target,
			solint='**',
			refant=antref,
			minblperant=3, minsnr=5)

	# Plot solutions
	plotms(vis='target.p1',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='phase',
		   iteraxis='antenna',
		   coloraxis='spw',
		   plotrange=[-1,-1,-180,180],
		   correlation='L',
		   plotfile='Step_06_target.p1_L.png', overwrite=True, showgui=gui)
	os.system('mv Step_06_target.p1_L_*png Step_06_target.p1_L.png')

	plotms(vis='target.p1',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='phase',
		   iteraxis='antenna',
		   coloraxis='spw',
		   plotrange=[-1,-1,-180,180],
		   correlation='R',
		   plotfile='Step_06_target.p1_R.png', overwrite=True, showgui=gui)
	os.system('mv Step_06_target.p1_R_*png Step_06_target.p1_R.png')

### 5) Apply phase solutions and image again

mystep = 7

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	applycal(vis=target+'.ms',
			 gaintable=['**'])

	# clean phase self-calibrated image
	os.system('rm -rf '+target+'_p1.clean*')
	if mac == False:
		tclean(vis=target+'.ms',
			   imagename='targetp1',
			   imsize=['**','**'],
			   cell='**arcsec',
			   deconvolver='clark',
			   specmode='mfs',
			   niter=3000,  # stop before this if residuals reach noise
			   cycleniter=300,
			   threshold=0,
			   savemodel='modelcolumn',
			   interactive=True)
	else:
		iclean(vis=target+'.ms',
				   imagename='targetp1',
				   imsize=['**','**'],
				   cell='**arcsec',
				   deconvolver='clark',
				   specmode='mfs',
				   niter=3000,  # stop before this if residuals reach noise
				   cycleniter=300,
				   threshold=0,
				   savemodel='modelcolumn')
		ft(vis=target+'.ms',
	 	   model='targetp1.model',
		   usescratch=True)

	rmsp1=imstat(imagename='targetp1.image',
				  box='**,**,**,**')['rms'][0]
	peakp1=imstat(imagename='targetp1.image',
				  box='**,**,**,**')['max'][0]

	print(('targetp1.image rms  = %7.3f mJy/bm' % (rmsp1*1000.)))
	print(('targetp1.image peak = %7.3f mJy/bm' % (peakp1*1000.)))
	print(('targetp1.image SNR  = %7.3f' % (peakp1/rmsp1)))

	if mac==False:
		os.system('rm Step_07_targetp1.image.png')
		imview('targetp1.image')
		imview(raster={'file': 'targetp1.image',
					'range': [-1.*rmsp1, 100.*rmsp1],
					'colorwedge': True},
			zoom=1,
			out='Step_07_targetp1.image.png')
	else:
		print('IMPORTANT: MAC OS detected .. Please view targetp1.image with CARTA')

# Has SNR improved?

### 6) Self-cal target #2
# Better model, shorter solution interval
# This time, make spectral index image to prepare for amp self-cal

mystep = 8

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	os.system('rm -rf target.p2')
	gaincal(vis=target+'.ms',
			calmode='**',
			caltable='target.p2',
			field=target,
			solint='**',
			refant=antref,
			gaintable=['**'],
			interp=['**'],
			minblperant=3, minsnr=5)

	# Plot solutions
	plotms(vis='target.p2',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='phase',
		   iteraxis='antenna',
		   coloraxis='spw',
		   plotrange=[-1,-1,-180,180],
		   correlation='L',
		   plotfile='Step_08_target.p2_L.png', overwrite=True, showgui=gui)
	os.system('mv Step_08_target.p2_L*png Step_08_target.p2_L.png')
	
	plotms(vis='target.p2',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='phase',
		   iteraxis='antenna',
		   coloraxis='spw',
		   plotrange=[-1,-1,-180,180],
		   correlation='R',
		   plotfile='Step_08_target.p2_R.png', overwrite=True, showgui=gui)
	os.system('mv Step_08_target.p2_R*png Step_08_target.p2_R.png')

### 7) Apply new phase calibration and image
# This time, make spectral index image to prepare for amp self-cal

mystep = 9

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	applycal(vis=target+'.ms',
			 gaintable=['**', '**'])

	# clean phase self-calibrated image
	os.system('rm -rf targetp2.*')
	if mac == False:
		tclean(vis=target+'.ms',
			   imagename='targetp2',
			   imsize=['**','**'],
			   cell='**arcsec',
			   specmode='mfs',
			   niter=4000, # stop before this if residuals reach noise
			   cycleniter=500,
			   gain=0.2,
			   threshold=0,
			   deconvolver='mtmfs',
			   savemodel='modelcolumn',
			   nterms=2, # Make spectral index image
			   interactive=True)
	else:
		iclean(vis=target+'.ms',
				   imagename='targetp2',
				   imsize=['**','**'],
				   cell='**arcsec',
				   specmode='mfs',
				   niter=4000, # stop before this if residuals reach noise
				   cycleniter=500,
				   gain=0.2,
				   threshold=0,
				   deconvolver='mtmfs',
				   savemodel='modelcolumn',
				   nterms=2) # Make spectral index image
		ft(vis=target+'.ms',
	 	   model=['targetp2.model.tt0','targetp2.model.tt1'],
		   nterms=2,
		   usescratch=True)

	rmsp2=imstat(imagename='targetp2.image.tt0',
				  box='**,**,**,**')['rms'][0]
	peakp2=imstat(imagename='targetp2.image.tt0',
				  box='**,**,**,**')['max'][0]

	print(('targetp2.image.tt0 rms  = %7.3f mJy/bm' % (rmsp2*1000.)))
	print(('targetp2.image.tt0 peak = %7.3f mJy/bm' % (peakp2*1000.)))
	print(('targetp2.image.tt0 S/N  = %7.3f' % (peakp2/rmsp2)))
	
	if mac==False:
		os.system('rm Step_09_targetp2.image.png')
		imview('targetp2.image.tt0')
		imview(raster={'file': 'targetp2.image.tt0',
					   'range': [-1.*rmsp2, 100.*rmsp2],
					   'colorwedge': True},
			zoom=1,
			out='Step_09_targetp2.image.png')
	else:
		print('IMPORTANT: MAC OS detected .. Please view targetp2.image.tt0 with CARTA')

### 8) Chose interval for amplitude self-cal

mystep = 10

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	plotms(vis=target+'.ms',
		   xaxis='**',
		   yaxis='**',
		   ydatacolumn='**',
		   antenna='**&**',
		   correlation='RR,LL',
		   avgchannel='16',	  
		   coloraxis='spw',
		   timerange='',
		   plotfile='',
		   overwrite=True, showgui=gui)


### 9) amplitude self-cal

mystep = 11

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	os.system('rm -rf target.ap3')
	gaincal(vis=target+'.ms',
			calmode='**',
			caltable='target.ap3',
			field=target,
			solint='**',
			refant=antref,
			solnorm=True,
			gaintable=['**','**'],
			minblperant=3, minsnr=5)

	# Plot solutions
	# If these are all more than ~20% from unity, check for a poor model
	# If just a few are  more than ~20% from unity, check for bad data
	plotms(vis='target.ap3',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='Step_11_target.ap3_L.png', overwrite=True, showgui=gui)
	os.system('mv Step_11_target.ap3_L*png Step_11_target.ap3_L.png')

	plotms(vis='target.ap3',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='R',
		   plotfile='Step_11_target.ap3_R.png', overwrite=True, showgui=gui)
	os.system('mv Step_11_target.ap3_R*png Step_11_target.ap3_R.png')

### 10) Apply amplitude calibration and image

mystep = 12

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	applycal(vis=target+'.ms',
			 gaintable=['**','**','**'], calwt=False)

	# clean amp&phase self-calibrated image
	os.system('rm -rf targetap3.*')
	if mac == False:
		tclean(vis=target+'.ms',
			   imagename='targetap3',
			   imsize=['**','**'],
			   cell='**arcsec',
			   specmode='mfs',
			   niter=5000, # stop before this if residuals reach noise
			   cycleniter=500,
			   gain=0.2,
			   threshold=0,
			   deconvolver='mtmfs',
			   savemodel='modelcolumn',
			   nterms=2, # Make spectral index image
			   interactive=True)
	else:
		iclean(vis=target+'.ms',
				   imagename='targetap3',
				   imsize=['**','**'],
				   cell='**arcsec',
				   specmode='mfs',
				   niter=5000, # stop before this if residuals reach noise
				   cycleniter=500,
				   gain=0.2,
				   threshold=0,
				   deconvolver='mtmfs',
				   savemodel='modelcolumn',
				   nterms=2) # Make spectral index image

	rmsap1=imstat(imagename='targetap3.image.tt0',
				  box='**,**,**,**')['rms'][0]
	peakap1=imstat(imagename='targetap3.image.tt0',
				  box='**,**,**,**')['max'][0]

	print(('targetap3.image.tt0 rms  = %7.3f mJy/bm' % (rmsap1*1000.)))
	print(('targetap3.image.tt0 peak = %7.3f mJy/bm' % (peakap1*1000.)))
	print(('targetap3.image.tt0 S/N  = %7.3f' % (peakap1/rmsap1)))

	if mac==False:
		os.system('rm Step_12_targetap3.image.png')
		imview('targetap3.image.tt0')
		imview(raster={'file': 'targetap3.image.tt0',
					   'range': [-1.*rmsap1, 100.*rmsap1],
					   'colorwedge': True},
			   zoom=1,
			   out='Step_12_targetap3.image.png')
	else:
		print('IMPORTANT: MAC OS detected .. Please view targetap3.image.tt0 with CARTA')

# 1252+5634_ap1.clean.image.tt0 rms   0.107 mJy
# 1252+5634_ap1.clean.image.tt0 peak 172.318 mJy/bm


### 11) Re-weight according to sensitivity
# Re-weight the data according to sensitivity (rms of residual in short intervals)

mystep = 13

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	# View weights
	plotms(vis=target+'_adv_im.ms', field=target, xaxis='uvdist',
		   yaxis='weight',antenna=antref+'&*',
		   correlation='LL,RR',
		   coloraxis='antenna2',plotfile='Step_13_target_weight-uv.png',
		   overwrite=True, showgui=gui)

	# The weights may look odd because e-MERLIN weights are not handled
	# correctly by all versions of CASA

	# For statwt, one can now use residuals which are calculated on the fly from the 
	# corrected and model columns.  This means that the variance of the data is not
	# biased by source structure, if the imaging was OK.
	# It can only be used properly if there is a reasonable model in the MS.

	# re-weight
	statwt(vis='**',
		   datacolumn='residual_data')

	# View weights again
	plotms(vis=target+'_adv_im.ms', field=target, xaxis='uvdist',
		   yaxis='weight',antenna=antref+'&*',
		   correlation='LL,RR',
		   coloraxis='antenna2',
		   plotfile='Step_13_target_re-weight-uv.png',
		   overwrite=True, showgui=gui)


### 12) Multiscale clean
# Try searching for multiple scales in clean

mystep = 14

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	# clean amp&phase self-calibrated image
	os.system('rm -rf advIMmulti.*')
	if mac == False:
		tclean(vis=target+'_adv_im.ms',
			   imagename='advIMmulti',
			   imsize=['**','**'],
			   cell='**arcsec',
			   specmode='mfs',
			   niter=3000, # stop before this if residuals reach noise
			   cycleniter=300,
			   gain=0.2,
			   threshold=0,
			   deconvolver='mtmfs',
			   scales=['**','**','**'],
			   nterms=2,
			   interactive=True)
	else:
		iclean(vis=target+'_adv_im.ms',
				   imagename='advIMmulti',
				   imsize=['**','**'],
				   cell='**arcsec',
				   specmode='mfs',
				   niter=3000, # stop before this if residuals reach noise
				   cycleniter=300,
				   gain=0.2,
				   threshold=0,
				   deconvolver='mtmfs',
				   scales=['**','**','**'],
				   nterms=2)	

	# More components are subtracted per cycle so clean runs faster
	rmsmulti=imstat(imagename='advIMmulti.image.tt0',
				  box='**,**,**,**')['rms'][0]
	peakmulti=imstat(imagename='advIMmulti.image.tt0',
				  box='**,**,**,**')['max'][0]

	print('advIMmulti.image.tt0 rms  = %7.3f mJy/bm' % (rmsmulti*1000.))
	print('advIMmulti.image.tt0 peak = %7.3f mJy/bm' % (peakmulti*1000.))
	print('advIMmulti.image.tt0 S/N  = %7.3f' % (peakmulti/rmsmulti))


### 13) Different imaging weights: tapering

mystep = 15

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	# Experiment with changing the relative weights given to longer
	# baselines.  Applying a taper reduces this, giving a larger
	# synthesised beam more sensitive to extended emission. This just
	# changes the weights during gridding for imaging, no permanent change
	# to the uv data. However, strong tapering increases the noise.

	# Look at the uv distance plot and choose an outertaper (half-power of
	# Gaussian taper applied to data) which is about 2/3 of the longest
	# baseline. Units are wavelengths, so calculate that or plot amp
	# v. uvwave

	os.system('rm -rf advIMtaper.*')
	if mac == False:
		tclean(vis=target+'_adv_im.ms',
			   imagename='advIMtaper',
			   imsize=['**','**'],
			   cell='**arcsec',
			   specmode='mfs',
			   niter=4000, # stop before this if residuals reach noise
			   cycleniter=300,
			   gain=0.2,
			   threshold=0,
			   deconvolver='mtmfs',
			   uvtaper=['**Mlambda'],
			   scales=[0],
			   nterms=2,
			   interactive=True)
	else:
		iclean(vis=target+'_adv_im.ms',
				   imagename='advIMtaper',
				   imsize=['**','**'],
				   cell='**arcsec',
				   uvtaper=['**Mlambda'],
				   specmode='mfs',
				   niter=4000, # stop before this if residuals reach noise
				   cycleniter=300,
				   gain=0.2,
				   threshold=0,
				   deconvolver='mtmfs',
				   scales=[0],
				   nterms=2)

	rmstaper=imstat(imagename='advIMtaper.image.tt0',
				  box='**,**,**,**')['rms'][0]
	peaktaper=imstat(imagename='advIMtaper.image.tt0',
				  box='**,**,**,**')['max'][0]

	print('advIMtaper.image.tt0 rms  = %7.3f mJy/bm' % (rmstaper*1000.))
	print('advIMtaper.image.tt0 peak = %7.3f mJy/bm' % (peaktaper*1000.))
	print('advIMtaper.image.tt0 S/N  = %7.3f' % (peaktaper/rmstaper))


### 14) Different imaging weights: Briggs robust

mystep = 16

if(mystep in thesteps):
	print('Step ', mystep, step_title[mystep])

	# Experiment with changing the relative weights given to longer
	# baselines.  Default natural weighting gives zero weight to parts of
	# the uv plane with no samples.  Briggs weighting interpolates between
	# samples, giving more weight to the interpolated areas for lower
	# values of robust.  For an array like e-MERLIN with sparser sampling
	# on long baselines, this has the effect of increasing their
	# contribution, i.e. making the synthesised beam smaller. This may
	# bring up the jet/core, at the expense of the extended lobes.  Low
	# values of robust increase the noise. For now, remove tapering and
	# multiscale.

	os.system('rm -rf advIMrobust*')
	if mac == False:
		tclean(vis=target+'_adv_im.ms',
			   imagename='advIMrobust',
			   imsize=['**','**'],
			   cell='**arcsec',
			   specmode='mfs',
			   niter=4000, # stop before this if residuals reach noise
			   cycleniter=300,
			   gain=0.2,
			   threshold=0,
			   deconvolver='mtmfs',
			   nterms=2,
			   weighting='**',
			   robust='**',
			   interactive=True)
	else:		
		iclean(vis=target+'_adv_im.ms',
			 	   imagename='advIMrobust',
			 	   imsize=['**','**'],
			 	   cell='**arcsec',
			 	   specmode='mfs',
			 	   niter=4000, # stop before this if residuals reach noise
			 	   cycleniter=300,
			 	   gain=0.2,
			 	   threshold=0,
			 	   deconvolver='mtmfs',
			 	   nterms=2,
			 	   weighting='**',
			 	   robust='**')

	rmsrobust=imstat(imagename='advIMrobust.image.tt0',
				  box='**,**,**,**')['rms'][0]
	peakrobust=imstat(imagename='advIMrobust.image.tt0',
				  box='**,**,**,**')['max'][0]

	print(('advIMrobust.image.tt0 rms  = %7.3f mJy/bm' % (rmsrobust*1000.)))
	print(('advIMrobust.image.tt0 peak = %7.3f mJy/bm' % (peakrobust*1000.)))
	print(('advIMrobust.image.tt0 S/N  = %7.3f' % (peakrobust/rmsrobust)))