########################################################################
# Data reduction script for SgrA* Band 3:
# - Calibration script: "SgrA_Band3_Calibration.py" -
# Tested in CASA Version 3.4.0 (r19988)
########################################################################
"""
See accompanying README file for details of the necessary input files
and comments on the data

"""
#----------------------------------------------------------------------------------
#----- Optional Steps -------------------------------------------------------------
#----------------------------------------------------------------------------------

#----- Set this option to true if you want to make diagnostic plots along the way.
#----- Note that this will slow down the reduction significantly.
#----- Default is no plots (you can make them later anyway).
calplots=F

#----- Set this option to true if you want to run this script interactively to make inspection plots.
#----- You'll need to hit <enter> at various stages to continue the script.
#----- Default is true (user input is required).
#----- Set to false if you want no user interaction except clean.
interact=T 

#----------------------------------------------------------------------------------
#----- Some setup steps -----------------------------------------------------------
#----------------------------------------------------------------------------------

#----- Input MS names 
#-----  (the data provided have already been coverted to CASA measurement set
#-----  format (.ms) using the task importasdm)
basename=["uid___A002_X230515_X107","uid___A002_X230515_X264","uid___A002_X230515_X409"]

#----- Re-name the MS to something shorter
basename_all=["X107","X264","X409"]

for i in range(len(basename)):
	os.system('mv '+ basename[i]+'.ms '+ basename_all[i]+'.ms')

#----- List of field names to be used
field_names = ['3c279','nrao530','SgrA']

#----------------------------------------------------------------------------------
#----- Initial summary information ------------------------------------------------
#----------------------------------------------------------------------------------

#----- Create a summary text file for each dataset
for asdm in basename_all:
	os.system('rm -rf '+asdm+'.listobs.txt')
        listobs(vis=asdm+'.ms', listfile=asdm+'.listobs.txt', verbose=True)
        print "<< Printed listobs output to "+asdm+'.listobs.txt'

#----- Plot the antenna configuration for each dataset to a file
for asdm in basename_all:
        print "<< Plotting antenna configuration for dataset "+asdm+' to '+asdm+'.plotants.png'
	os.system('rm -rf '+asdm+'.plotants.png')
        plotants(vis=asdm+'.ms', figfile=asdm+'.plotants.png')
        if(interact):
		dummy_string = raw_input("Hit <Enter> to see the antenna configuration for the next data set")

#----------------------------------------------------------------------------------
#----- Inspect the data -----------------------------------------------------------
#----------------------------------------------------------------------------------

#----- Inspect amplitude versus time for each dataset
if(interact):
	for asdm in basename_all:
		plotms(vis=asdm+'.ms',
		       interactive=T,
		       spw='1',
		       xaxis='time', yaxis='amp',
		       correlation='XX',
		       antenna='*&*',
		       avgchannel='10000',coloraxis='field')
		dummy_string = raw_input("Once the plot has loaded, hit <Enter> to see the next data set")

#----------------------------------------------------------------------------------
#----- Initial Flagging -----------------------------------------------------------
#----------------------------------------------------------------------------------

for asdm in basename_all:
#----- Remove all the previous flags
	flagdata(vis=asdm+'.ms',mode='manualflag', unflag=T, flagbackup=F)
#----- Flagging shadowed data
	flagdata(vis=asdm+'.ms',mode = 'shadow', diameter=12.0, flagbackup = F)
#----- Flagging pointing and atmosphere calibration scans
	flagdata(vis=asdm+'.ms', mode='manualflag', intent='*POINTING*',flagbackup = F)
	flagdata(vis=asdm+'.ms', mode='manualflag', intent='*ATMOSPHERE*',flagbackup = F)
#----- Flagging autocorrelation data
	flagdata(vis=asdm+'.ms',autocorr=True,flagbackup=F)
#----- Saving 'a priori' flags
	flagmanager(vis = asdm+'.ms', mode = 'save', versionname = 'Apriori')
#----- Restoring up 'a priori' flags"+asdm
#        flagmanager(vis = asdm+'.ms', mode = 'restore', versionname = 'Apriori')

#----------------------------------------------------------------------------------
#----- Create, Examine and Apply Tsys and WVR calibration tables --------------------------
#----------------------------------------------------------------------------------
#----- Creating Tsys tables in TDM
os.system('rm -rf *tdm.tsys')
for asdm in basename_all:
    print "Creating TDM Tsys Table for "+asdm
    gencal(vis=asdm+'.ms',caltable=asdm+'.tdm.tsys',spw='9,11,13,15',caltype='tsys')
#----- Mapping the TDM Tsys to FDM spws
from recipes.almahelpers import tsysspwmap
for asdm in basename_all:
    tsysspwmap(vis=asdm+'.ms',tsystable=asdm+'.tdm.tsys')

#----- Plotting Tsys tables
if(calplots):
	for asdm in basename_all:
#----- Plotting Tsys vs. time for antenna DV01
		plotcal(caltable=asdm+'.tdm.tsys', 
			xaxis="time",yaxis="amp",
			plotsymbol=".", subplot=221,
			antenna='DV01', plotrange = [0,0,50,100],
			iteration='spw', figfile=asdm+'.tsys_vs_time.png',
			fontsize=6.0)    
#----- Plotting Tsys vs. frequency for antenna DV01
		plotcal(caltable=asdm+'.tdm.tsys',
			xaxis="freq",yaxis="amp",
			plotsymbol=".", subplot=221,
			antenna='DV01', plotrange = [0,0,50,100],
			iteration='spw', figfile=asdm+'.tsys_vs_freq.png',
			fontsize=6.0)    

if(calplots):
	for asdm in basename_all:
#----- Plotting Tsys vs. time (all antennas) for spw 9
		plotcal(caltable=asdm+'.tdm.tsys', 
			xaxis="time",yaxis="amp",
			spw='9:10~118',plotsymbol=".", subplot=231,
			antenna='0~5',
			iteration='antenna', figfile=asdm+'.tsys_vs_time.page1.png',
			fontsize=6.0)    
		plotcal(caltable=asdm+'.tdm.tsys', 
			xaxis="time",yaxis="amp",
			antenna='6~10',
			spw='9:10~118',plotsymbol=".", subplot=231,
			iteration='antenna', figfile=asdm+'.tsys_vs_time.page2.png',
			fontsize=6.0)    
#----- Plotting Tsys vs. frequency (all antennas) for spw 9
		plotcal(caltable=asdm+'.tdm.tsys',
			xaxis="freq",yaxis="amp",
			spw='', plotsymbol=".", subplot=231,
			iteration='antenna', figfile=asdm+'.tsys_vs_freq.page1.png',
			antenna='0~5', fontsize=6.0)    
		plotcal(caltable=asdm+'.tdm.tsys',
			xaxis="freq",yaxis="amp",
			spw='', plotsymbol=".", subplot=231,
			iteration='antenna', figfile=asdm+'.tsys_vs_freq.page2.png',
			antenna='6~10', fontsize=6.0)

#----- Correct for delay problem in DV13 in dataset X107

#----- Creating WVR cal tables
os.system('rm -rf *.wvrgcal')
for asdm in basename_all:
    wvrgcal(vis=asdm+'.ms',caltable=asdm+'.wvrgcal',toffset=-1)

#----- Apply Tsys, WVR calibrations
for asdm in basename_all:
	for field in field_names:
		applycal(vis=asdm+".ms", spw='1,3,5,7', 
			field=field, gainfield=[field,field],
			interp=['linear,spline','nearest'],
			gaintable=[asdm+".tdm.tsys",asdm+'.wvrgcal'],
			spwmap=[[0, 9, 9, 11, 11, 13, 13, 15, 15, 9, 9, 11, 11, 13, 13, 15, 15],[]],
			flagbackup=F, calwt=T)

#----------------------------------------------------------------------------------
#----- Inspect the data again -----------------------------------------------------
#----------------------------------------------------------------------------------

#----- Plot spectral plot in amplitude and phase
if(interact):
	for asdm in basename_all:
		plotms(vis=asdm+'.ms', 
		      field='3c279',
		      xaxis='frequency', yaxis='amp',
		      selectdata=T, spw='1,3,5,7', 
		      coloraxis='corr',
		      iteraxis='baseline',
		      antenna='*&*',
		      ydatacolumn='data')
		dummy_string = raw_input("Next plot for "+asdm+" . Hit <Enter> to continue.")
		plotms(vis=asdm+'.ms', 
		      field='3c279',
		      xaxis='frequency', yaxis='phase',
		      selectdata=T, spw='1,3,5,7', 
		      coloraxis='corr',
		      iteraxis='baseline',
		      antenna='*&*',
		      ydatacolumn='data')
		dummy_string = raw_input("Next plot for "+asdm+" . Hit <Enter> to continue.")

#----------------------------------------------------------------------------------
#----- Split out corrected data and save the flags --------------------------------
#----------------------------------------------------------------------------------

#----- First save the flags
for asdm in basename_all:
	flagmanager(vis = asdm+'.ms', mode = 'save', versionname = 'Before_flagging')

#----- Now do the split and output another summary file
for asdm in basename_all:
	os.system('rm -rf '+ asdm+'.wvrtsys.ms')
	split(vis=asdm+'.ms', outputvis=asdm+'.wvrtsys.ms', 
	      datacolumn='corrected', spw='1,3,5,7', keepflags=F)
	os.system('rm '+asdm+'.wvrtsys.listobs.txt')
	listobs(vis=asdm+'.wvrtsys.ms', listfile=asdm+'.wvrtsys.listobs.txt', verbose=True)

#----------------------------------------------------------------------------------
#----- Further Data Flagging ------------------------------------------------------
#----------------------------------------------------------------------------------

#----- Flag shadowed data
for asdm in basename_all:
	flagdata(vis=asdm+'.wvrtsys.ms', flagbackup = F, mode = 'shadow')

#----- Flag DV02 SPW 2 (was 5 before split) BOTH POLARIZATIONS
for asdm in basename_all:
	flagdata(vis=asdm+'.wvrtsys.ms',antenna='DV02',spw='2',correlation='', selectdata=True, flagbackup = F)
#	flagdata2(flagbackup = F, vis=asdm+'.wvrtsys.ms',manualflag=T,mf_antenna='DV02',spw='2',correlation='', selectdata=True)

#----- Flag the less sensitive edge channels of every data set
for asdm in basename_all:
	flagdata(vis = asdm+'.wvrtsys.ms',spw ='*:0~150;3700~3840', selectdata=True, flagbackup = F)
#	flagdata2(vis = asdm+'.wvrtsys.ms',manualflag=T,mf_spw = ['*:0~150;3700~3840'],  selectdata=True, flagbackup = F)

#----------------------------------------------------------------------------------
#----- Concatenate datasets ---------------------------------------------
#----------------------------------------------------------------------------------

#----- Remove the pointing tables
for asdm in basename_all:
	tb.open(asdm+'.wvrtsys.ms/POINTING',nomodify=False)
	a = tb.rownumbers()
	tb.removerows(a)
	tb.close()

cal_vis = [vis+'.wvrtsys.ms' for vis in basename_all]

os.system('rm -rf SgrA.ms')
concat(vis=cal_vis, concatvis='SgrA.ms', timesort=T)

#----------------------------------------------------------------------------------
#----- Calibration Steps on the concatenated dataset ------------------------------
#----------------------------------------------------------------------------------
name='SgrA'
#----- Output another summary file
os.system('rm -rf '+name+'.listobs.txt')
listobs(vis=name+'.ms', listfile=name+'.listobs.txt', verbose=True)
print "<< Printed listobs output to "+name+'.listobs.txt'

#----- Running a short solution interval phase calibration for bandpass
os.system('rm -rf '+name+'.bpphase.gcal')
gaincal(vis = name+'.ms',
	selectdata=T,field = '3c279',spw = '*:1100~1300',
	caltable = name+'.bpphase.gcal',
	solint = 'int',refant = 'DV11',calmode='p',
	minblperant=4, minsnr=2)

# ----- To avoid a problem with passband calibration on the third dataset
flagdata(vis = name+'.ms',spw ='1', scan='39', selectdata=True, flagbackup = F)

#----- Running a bandpass calibration
os.system('rm -rf '+name+'.bandpass.bcal')
bandpass(vis = name+'.ms',
	field = '3c279', spw='*:151~3699',
	gaintable = name+'.bpphase.gcal',
	caltable = name+'.bandpass.bcal',
	bandtype='B',
	solint = 'inf',combine = 'scan,field', solnorm=T,refant = 'DV11',
	minblperant=3,minsnr=2,fillgaps=62)

#----- Plotting bandpass solutions
if(calplots):
	plotcal(caltable = name+'.bpphase.gcal',
			xaxis = 'time', yaxis = 'phase',
			iteration = 'antenna', plotrange=[0,0,-180,180],
			showgui=False, subplot=421, figfile=name+'.bpphase.page1.png',
			antenna='0~7')
	plotcal(caltable = name+'.bpphase.gcal',
			xaxis = 'time', yaxis = 'phase',
			iteration = 'antenna', plotrange=[0,0,-180,180],
			showgui=False, subplot=421, figfile=name+'.bpphase.page2.png',
			antenna='8~15')
	plotcal(caltable = name+'.bandpass.bcal', 
			xaxis = 'freq',yaxis = 'amp',
			plotrange = [0,0,0,1],
			antenna='0~5', iteration='antenna',
			showgui=False, subplot=321, figfile=name+'.bcal_amp.page1.png')
	plotcal(caltable = name+'.bandpass.bcal', 
			xaxis = 'freq',yaxis = 'amp',
			plotrange = [0,0,0,1],
			antenna='6~10', iteration='antenna',
			showgui=False, subplot=321, figfile=name+'.bcal_amp.page2.png')
	plotcal(caltable = name+'.bandpass.bcal', 
			xaxis = 'freq',yaxis = 'phase', iteration='antenna',
			antenna='0~5', subplot=321, figfile=name+'.bcal_phase.page1.png',
			plotrange = [0,0,-180,180], showgui=False)
	plotcal(caltable = name+'.bandpass.bcal', 
			xaxis = 'freq',yaxis = 'phase', iteration='antenna',
			antenna='6~10', subplot=321, figfile=name+'.bcal_phase.page2.png',
			plotrange = [0,0,-180,180], showgui=False)

#----- Carrying out short timescale (integration, ~ 8 sec) phase solution
os.system('rm -rf '+name+'.intphase.gcal')
gaincal(vis=name+'.ms',
		gaintable=name+'.bandpass.bcal', 
                caltable=name+'.intphase.gcal',
                calmode='p',
                field='3c279,nrao530',
                spw='*:40~3800',
                refant='DV11', solint='int',minsnr=2.0,minblperant=4)

#----- Flag Low  or rising fluxes of the las two observations
flagdata(flagbackup = F, vis = name+'.ms',
	 timerange='05:11:0~08:40:0',
	 spw='2,3')

#----- Flag No nrao530 observations after scan 47 on SgrA
flagdata(flagbackup = F, vis = name+'.ms',
	 scan='47')

#------ TITAN contaminated by Saturn
#----- Use nrao530 instead
#----- nrao530 has a flux of 4.7 Jy in Band 3 (from CARMA calibrator tool)
setjy(vis = name+'.ms',
      field = 'nrao530',
      spw = '0~3:10~3830',
      fluxdensity=[4.7,0,0,0])

#----- Carrying out longer timescale (by scan) phase solution
os.system('rm -rf '+name+'.scanphase.gcal')
gaincal(vis=name+'.ms',
                gaintable=name+'.bandpass.bcal', 
                caltable=name+'.scanphase.gcal',
                calmode='p',
                field='3c279,nrao530',
                spw='*:40~3800',
                refant='DV11', solint='Inf',minsnr=2.0,minblperant=4)

#----- Solving for longer (scan) interval amplitude solution
os.system('rm -rf '+name+'.amp.cal')
gaincal(vis = name+'.ms',
               gaintable =[name+'.bandpass.bcal',name+'.intphase.gcal'],
               caltable = name+'.amp.cal',
               calmode='a',
               field = '3c279,nrao530',
               spw='*:40~3800',
               refant = 'DV11',solint = 'inf')

#----- Flux transfer to the Bandpass Calibrator. This step is not needed but provides a flux for 3c279 of 18 Jy
fluxscale(vis=name+'.ms',
	  caltable=name+'.amp.cal', fluxtable=name+'.flux.cal',
	  reference ='nrao530', transfer =['3c279'])

#----- Plotting solutions
if(calplots):
	plotcal(caltable = name+'.scanphase.gcal',
			xaxis = 'time', yaxis = 'phase',
			iteration = 'antenna', plotrange=[0,0,-180,180],
			showgui=False, subplot=421, figfile=name+'.scanphase.page1.png',
			antenna='0~5', fontsize=6.0)
	plotcal(caltable = name+'.scanphase.gcal',
			xaxis = 'time', yaxis = 'phase',
			iteration = 'antenna', plotrange=[0,0,-180,180],
			showgui=False, subplot=421, figfile=name+'.scanphase.page2.png',
			antenna='6~10', fontsize=6.0)

#----- Applying calibrations to the different sources:

applycal(vis=name+'.ms',field='3c279',
	gaintable=[name+'.bandpass.bcal',name+'.intphase.gcal',name+'.flux.cal'],
	interp=['nearest','nearest','nearest'], spw='*:10~3830',
	gainfield=['3c279','3c279','3c279'],flagbackup=F, calwt=F)
applycal(vis=name+'.ms',field='nrao530',
	gaintable=[name+'.bandpass.bcal',name+'.intphase.gcal',name+'.flux.cal'],
	interp=['nearest','nearest','nearest'], spw='*:10~3830',
	gainfield=['3c279','nrao530','nrao530'],flagbackup=F, calwt=F)
applycal(vis=name+'.ms',field='SgrA',
	interp=['nearest','linear','linear'], spw='*:10~3830',
	gaintable=[name+'.bandpass.bcal',name+'.scanphase.gcal',name+'.flux.cal'],
	gainfield=['3c279','nrao530','nrao530'],flagbackup=F, calwt=F)

#----- Inspecting calibrated data
if(interact):
	plotms(vis = name+'.ms', xaxis='uvdist', yaxis='amp',
		        ydatacolumn='corrected', field='3c279',
		        averagedata=True, avgchannel='3840', avgtime='',
		        avgscan=F, avgbaseline=F, coloraxis='corr')
	dummy_string = raw_input("Next plot for "+name+" . Hit <Enter> to continue.")
	plotms(vis = name+'.ms', xaxis='uvdist', yaxis='phase',
		        ydatacolumn='corrected', field='3c279',
			avgchannel='3840', avgscan=F, avgbaseline=F, coloraxis='corr')
	dummy_string = raw_input("Next plot for "+name+" . Hit <Enter> to continue.")
	plotms(vis = name+'.ms', xaxis='time', yaxis='amp',
		        ydatacolumn='corrected', field='',
			averagedata=True, avgchannel='3840', avgtime='',
			avgscan=F, avgbaseline=F, coloraxis='corr')
	dummy_string = raw_input("Next plot for "+name+" . Hit <Enter> to continue.")
	plotms(vis = name+'.ms', xaxis='time', yaxis='phase',
		        ydatacolumn='corrected', field='',
			avgchannel='3840', avgscan=F, avgbaseline=F, coloraxis='corr')

#  Check calibrator image to check calibration
clean(vis = name+'.ms',
      cell = '0.8arcsec', imsize = 128,
      field = 'nrao530', interactive = T,
      imagename = 'SgrA_B3Spw0_image_calibrator',
      spw='0', mode='channel',
      interpolation='nearest',
      width=10, niter=1000, npercycle=10)

#----------------------------------------------------------------------------------
#----- Split out calibrated data on SgrA* -----------------------------------------
#----------------------------------------------------------------------------------

for spwnum in ['0','1','2','3']:
		split(vis=name+'.ms',
		      datacolumn='corrected',
		      outputvis=name+'.mssplitSpw'+spwnum,
		      field='SgrA', spw=spwnum)


#----------------------------------------------------------------------------------
#----- End of calibration script.
#----- To continue, see imaging script "SgrA_Band3_Imaging.py"
#----------------------------------------------------------------------------------