# Script to image eMERLIN data with CASA 6 or later

import os, sys

# 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

# *** I have left deconvolver 'clark' but 'hogbom' is normally used?
# *** plotms very high resolution png, would just displays be better?

# Enter step(s) as mysteps

# 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',
			  17: 'a',18: 'a'}

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 -rf '+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='64',
		   showmajorgrid=True,
		   showminorgrid=True,
		   spw='3',
		   plotfile='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 '+target+'_1.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_1.clean',
		   cell='0.012arcsec',
		   niter=1500,
		   cycleniter=300,
		   imsize=[512,512],
		   weighting='natural',
		   deconvolver='clark',
		   savemodel='modelcolumn',
		   interactive=True)

mystep = 3

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

	rms1=imstat(imagename=target+'_1.clean.image',
				  box='21,178,132,245')['rms'][0]
	peak1=imstat(imagename=target+'_1.clean.image',
				  box='49,46,203,195')['max'][0]

	print((target+'_1.clean.image rms %7.3f mJy' % (rms1*1000.)))
	print((target+'_1.clean.image peak %7.3f mJy/bm' % (peak1*1000.)))

	os.system('rm %s_1.clean.png'%target)
	imview(raster={'file': target+'_1.clean.image',
				   'colormap':'Rainbow 2',
				   'scaling': -2,
				   'range': [-1.*rms1, 100.*rms1],
				   'colorwedge': True},
		   zoom=1,
		   out=target+'_1.clean.png')

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='64',
		   coloraxis='spw',
		   plotfile='3C277.1_uvdist_model.png',
		   overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)


### 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='time',
		   yaxis='phase',
		   antenna='Mk2&Pi',
		   correlation='RR,LL',
		   avgchannel='64',
		   coloraxis='spw',
		   plotfile='3C277.1_phase.png',
		   timerange='02:02:40~02:18:40',
		   overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

### 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='p',
			caltable='target.p1',
			field=target,
			solint='180s',
			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='target.p1_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	plotms(vis='target.p1',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='phase',
		   iteraxis='antenna',
		   coloraxis='spw',
		   plotrange=[-1,-1,-180,180],
		   correlation='R',
		   plotfile='target.p1_R.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

### 5) Apply phase solutions and image again

mystep = 7

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

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

	# clean phase self-calibrated image
	os.system('rm -rf '+target+'_p1.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_p1.clean',
		   cell='0.012arcsec',
		   niter=2000,  # stop before this if residuals reach noise
		   cycleniter=300,
		   deconvolver='clark',
           savemodel='modelcolumn',
		   imsize=[512,512],
		   interactive=True)

	rmsp1=imstat(imagename=target+'_p1.clean.image',
				  box='21,178,132,245')['rms'][0]
	peakp1=imstat(imagename=target+'_p1.clean.image',
				  box='49,46,203,195')['max'][0]

	print((target+'_p1.clean.image rms %7.3f mJy' % (rmsp1*1000.)))
	print((target+'_p1.clean.image peak %7.3f mJy/bm' % (peakp1*1000.)))


	imview(target+'_p1.clean.image')

	imview(raster={'file': target+'_p1.clean.image',
				   'range': [-1.*rmsp1, 20.*rmsp1],
				   'colorwedge': True},
		   zoom=1,
		   out=target+'_p1.clean.png')

# 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='p',
			caltable='target.p2',
			field=target,
			solint='30s',
			refant=antref,
			gaintable=['target.p1'],
			interp=['linear'],
			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='target.p2_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)
	
	plotms(vis='target.p2',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='phase',
		   iteraxis='antenna',
		   coloraxis='spw',
		   plotrange=[-1,-1,-180,180],
		   correlation='R',
		   plotfile='target.p2_R.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)


### 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=['target.p1', 'target.p2'])

	# clean phase self-calibrated image
	os.system('rm -rf '+target+'_p2.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_p2.clean',
		   cell='0.012arcsec',
		   niter=3500, # stop before this if residuals reach noise
		   cycleniter=300,
		   gain=0.2,
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)

	rmsp2=imstat(imagename=target+'_p2.clean.image.tt0',
				  box='21,178,132,245')['rms'][0]
	peakp2=imstat(imagename=target+'_p2.clean.image.tt0',
				  box='49,46,203,195')['max'][0]

	print((target+'_p2.clean.image.tt0 rms %7.3f mJy' % (rmsp2*1000.)))
	print((target+'_p2.clean.image.tt0 peak %7.3f mJy/bm' % (peakp2*1000.)))

	imview(target+'_p2.clean.image.tt0')

	imview(raster={'file': target+'_p2.clean.image.tt0',
				   'range': [-1.*rmsp2, 100.*rmsp2],
				   'colorwedge': True},
		   zoom=1,
		   out=target+'_p2.clean.png')

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

mystep = 10

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


	plotms(vis=target+'.ms',
		   xaxis='time',
		   yaxis='amp',
		   ydatacolumn='corrected',
		   antenna='Mk2&Pi',
		   correlation='RR,LL',
		   avgchannel='64',
		   coloraxis='spw',
		   timerange='',
		   plotfile='3C277.1_amp.png',
		   overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)



### 9) amplitude self-cal

mystep = 11

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

	os.system('rm %s_pre_ampsc.tar.gz'%(target))
	os.system('tar -cvzf %s_pre_ampsc.tar.gz %s.ms'%(target,target))
	os.system('rm -rf target.ap1')
	gaincal(vis=target+'.ms',
			calmode='ap',
			caltable='target.ap1',
			field=target,
			solnorm=True,
			solint='inf',
			refant=antref,
			gaintable=['target.p1','target.p2'],
			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.ap1',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='target.ap1_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	plotms(vis='target.ap1',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='R',
		   plotfile='target.ap1_R.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)


### 10) Apply amplitude calibration and image

mystep = 12

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

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

	# clean amp&phase self-calibrated image
	os.system('rm -rf '+target+'_ap1.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_ap1.clean',
		   cell='0.012arcsec',
		   cycleniter=500,
		   gain=0.2,
		   niter=4000, # stop before this if residuals reach noise
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)

	rmsap1=imstat(imagename=target+'_ap1.clean.image.tt0',
				  box='21,178,132,245')['rms'][0]
	peakap1=imstat(imagename=target+'_ap1.clean.image.tt0',
				  box='49,46,203,195')['max'][0]

	print((target+'_ap1.clean.image.tt0 rms %7.3f mJy/bm' % (rmsap1*1000.)))
	print((target+'_ap1.clean.image.tt0 peak %7.3f mJy/bm' % (peakap1*1000.)))

	imview(target+'_ap1.clean.image.tt0')

	imview(raster={'file': target+'_ap1.clean.image.tt0',
				   'range': [-1.*rmsap1, 100.*rmsap1],
				   'colorwedge': True},
		   zoom=1,
		   out=target+'_ap1.clean.png')

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


### 13) Phase sc again 15s

mystep = 13

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

	gaincal(vis=target+'.ms',
			calmode='p',
			caltable='target.p3',
			field=target,
			solnorm=True,
			solint='15s',
			refant=antref,
			gaintable=['target.p1','target.p2','target.ap1'],
			minblperant=3, minsnr=5)

	plotms(vis='target.p3',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='target.p3_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	applycal(vis=target+'.ms',
			 gaintable=['target.p1','target.p2','target.ap1','target.p3'], calwt=False)

	# clean amp&phase self-calibrated image
	os.system('rm -rf '+target+'_p3.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_p3.clean',
		   cell='0.012arcsec',
		   cycleniter=300,
		   gain=0.2,
		   niter=4000, # stop before this if residuals reach noise
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)

#Amp self cal
mystep = 14

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

	gaincal(vis=target+'.ms',
			calmode='ap',
			caltable='target.ap2',
			field=target,
			solnorm=True,
			solint='90s',
			refant=antref,
			gaintable=['target.p1','target.p2','target.ap1','target.p3'],
			minblperant=3, minsnr=5)

	plotms(vis='target.ap2',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='target.ap2_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	applycal(vis=target+'.ms',
			 gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2'], calwt=False)

	# clean amp&phase self-calibrated image
	os.system('rm -rf '+target+'_ap3.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_ap2.clean',
		   cell='0.012arcsec',
		   cycleniter=300,
		   gain=0.2,
		   niter=4000, # stop before this if residuals reach noise
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)

mystep = 15

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

	gaincal(vis=target+'.ms',
			calmode='ap',
			caltable='target.ap3',
			field=target,
			solnorm=True,
			solint='60s',
			refant=antref,
			gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2'],
			minblperant=3, minsnr=5)

	plotms(vis='target.ap3',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='target.ap3_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	applycal(vis=target+'.ms',
			 gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2','target.ap3'], calwt=False)

	# clean amp&phase self-calibrated image
	os.system('rm -rf '+target+'_ap3.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_ap3.clean',
		   cell='0.012arcsec',
		   cycleniter=300,
		   gain=0.2,
		   niter=4000, # stop before this if residuals reach noise
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)

mystep = 16

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

	gaincal(vis=target+'.ms',
			calmode='ap',
			caltable='target.ap4',
			field=target,
			solnorm=True,
			solint='30s',
			refant=antref,
			gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2','target.ap3'],
			minblperant=3, minsnr=5)

	plotms(vis='target.ap4',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='target.ap4_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	applycal(vis=target+'.ms',
			 gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2','target.ap3','target.ap4'], calwt=False)

	# clean amp&phase self-calibrated image
	os.system('rm -rf '+target+'_ap4.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_ap4.clean',
		   cell='0.012arcsec',
		   cycleniter=300,
		   gain=0.2,
		   niter=4000, # stop before this if residuals reach noise
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)

mystep = 17

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

	gaincal(vis=target+'.ms',
			calmode='ap',
			caltable='target.ap5',
			field=target,
			solnorm=True,
			solint='15s',
			refant=antref,
			gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2','target.ap3','target.ap4'],
			minblperant=3, minsnr=5)

	plotms(vis='target.ap5',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='target.ap5_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	applycal(vis=target+'.ms',
			 gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2','target.ap3','target.ap4','target.ap5'], calwt=False)

	# clean amp&phase self-calibrated image
	os.system('rm -rf '+target+'_ap5.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_ap5.clean',
		   cell='0.012arcsec',
		   cycleniter=1000,
		   gain=0.2,
		   niter=10000, # stop before this if residuals reach noise
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)


mystep = 18

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

	gaincal(vis=target+'.ms',
			calmode='ap',
			caltable='target.ap6',
			field=target,
			solnorm=True,
			solint='15s',
			refant=antref,
			gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2','target.ap3','target.ap4','target.ap5'],
			minblperant=3, minsnr=5)

	plotms(vis='target.ap6',
		   gridcols=2, gridrows=3,
		   xaxis='time', yaxis='amp',
		   iteraxis='antenna',
		   coloraxis='spw',
		   correlation='L',
		   plotfile='target.ap6_L.png', overwrite=True, showgui=gui, highres=True, dpi = 800, width=1500, height=750)

	applycal(vis=target+'.ms',
			 gaintable=['target.p1','target.p2','target.ap1','target.p3','target.ap2','target.ap3','target.ap4','target.ap5','target.ap6'], calwt=False)

	# clean amp&phase self-calibrated image
	os.system('rm -rf '+target+'_ap6.clean*')
	tclean(vis=target+'.ms',
		   imagename=target+'_ap6.clean',
		   cell='0.012arcsec',
		   cycleniter=1000,
		   gain=0.2,
		   niter=10000, # stop before this if residuals reach noise
		   imsize=[512,512],
		   deconvolver='mtmfs',
		   savemodel='modelcolumn',
		   nterms=2, # Make spectral index image
		   interactive=True)