########################################################################
# Data reduction script for Orion Spectral Sweep in Band 6:
# - Imaging script: "Orion_SpectralSweep_Calbration.py" -
# Tested in CASA Version 3.3.0 (r16856)
########################################################################
"""
See accompanying README file for details of this script, 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=T

#----- 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 

#----------------------------------------------------------------------------------
#----- Reduction for dataset 1
#----------------------------------------------------------------------------------

#----- Input MS names 
#-----  (the data provided have already been coverted to CASA measurement set
#-----  format (.ms) using the task importasdm)
basename_ORIG=['uid___A002_X37b127_X2f','uid___A002_X37b127_X17d','uid___A002_X37b127_X2cb']

#----- Re-name the MS to something shorter
basename_all=['X2f','X17d','X2cb']

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

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

#----- RUN LISTOBS TO GET INFORMATION
for basename in basename_all:
	os.system('rm -rf '+basename+'.listobs.txt')
	listobs(vis=basename+'.ms', verbose=True, 
            listfile=basename+'.listobs.txt')
	print "<< Printed listobs output to "+basename+'.listobs.txt'

#----- RUN PLOTANTS TO GET ANTENNA POSITION INFORMATION
for basename in basename_all:
	print "<< Plotting antenna configuration for dataset "+basename+' to '+basename+'.antpos.png'
	os.system('rm -rf '+basename+'.antpos.png')
	plotants(vis=basename+'.ms', figfile=basename+'.antpos.png')

#----------------------------------------------------------------------------------
#----- A- INITIAL CALIBRATIONS
#----------------------------------------------------------------------------------

print "<< Doing flagging... "
#----- Remove all the previous flags
for basename in basename_all:
	flagdata(vis= basename+'.ms',mode='manualflag', unflag=T, flagbackup=F)
#----- Save flags so you can always go back to unflagged data
	flagmanager(vis = basename+'.ms', mode = 'save', versionname = 'Original')
#----- Flagging shadowed data
	flagdata(vis= basename+'.ms',mode = 'shadow', diameter=12.0, flagbackup = F)
#----- Flag autocorr
	flagdata(vis = basename+'.ms',mode = 'manualflag', autocorr = T,flagbackup = F)
#----- Flag pointing etc
	flagdata(vis = basename+'.ms',mode = 'manualflag',
         intent = '*POINTING*,*SIDEBAND*,*ATMOSPHERE*',flagbackup = F)
#----- Saving 'a priori' flags
	flagmanager(vis = basename+'.ms', mode = 'save', versionname = 'Apriori')
#----- To restore 'a priori' flags:
#        flagmanager(vis = basename+'.ms', mode = 'restore', versionname = 'Apriori')


#----- Apply to all SCIENCE (i.e. fdm: Callisto, J0607, Orion) data:
		
#----- For first dataset:
print "<< Doing applycal... "
for field in ['1','2','3']:
	applycal(vis=basename_all[0]+'.ms',field=field,
          spw='33,35,37,39,65,67,69,71,73,75,77,79,81,83,85,87,89,91,93,95',
          gaintable = ['X2f.ms.wvr.smooth'],
          gainfield = [field,'',''],interp = 'linear',calwt = T,flagbackup=F)

#----- For second dataset:
print "<< Doing applycal... "
applycal (vis=basename_all[1]+'.ms',field='2,3',
          spw='33,35,37,39,65,67,69,71,73,75,77,79,81,83,85,87,89,91,93,95',
          gaintable = ['X17d.ms.wvr.smooth'],
          gainfield = ['2,3'],interp = 'linear',calwt = T,flagbackup=F)

#----- For third dataset:
print "<< Doing applycal... "
applycal (vis=basename_all[2]+'.ms',field='2,3',
          spw='33,35,37,39,65,67,69,71,73,75,77,79,81,83,85,87,89,91,93,95',
          gaintable = ['X2cb.ms.wvr.smooth'],
          gainfield = ['2,3'],interp = 'linear',calwt = F,flagbackup=F)


#----------------------------------------------------------------------------------
#----- B- GET FLUX OF PHASECAL - This step uses FIRST dataset only as Callisto
#----- was at a higher elevation (the other two datasets were at low elevation)
#----------------------------------------------------------------------------------

#----- Fix planet/satellite position
fixplanets(vis=basename_all[0]+'.ms', field='Callisto', fixuvw=True)

#----- Split out WVR-corrected Callisto data, keep ONLY relevant spws (i.e. TDM for atmCal, FDM for ampcal)
os.system('rm -rf callisto_data')
print "<< Doing split... "
split (vis=basename_all[0]+'.ms',outputvis='callisto_data',
       datacolumn='corrected',field='1',
       spw='25,27,29,31,33,35,37,39')

#----- Inspect Tsys table
if(calplots):
	plotcal(caltable='tsys_callisto.fdm', field='0',
		iteration='antenna', xaxis='freq', yaxis='amp',
		subplot=11)
	dummy_string = raw_input("Hit <Enter> to continue...")

#----- We have Tsys measurements only for T3 on the ampcal and T5 for the science/phase data
#----- We therefore apply available Tsys corrections to both the ampcal (T3) and phasecal (T5)

#----- Apply Tsys table to Callisto data
print "<< Doing applycal... "
applycal (vis='callisto_data',spw='4,5,6,7',
        gaintable=['tsys_callisto.fdm'],gainfield=[''],
        interp = 'linear', calwt = T,flagbackup = F)

#----- Follow the same procedure for the T5 phasecal data (the one that has correct associated atmcal)
os.system('rm -rf phase_t5_data')
print "<< Doing split... "
split (vis=basename_all[0]+'.ms',outputvis='phase_t5_data',
        datacolumn='corrected',field='2',
        spw='57,59,61,63,89,91,93,95')

print "<< Doing applycal... "
applycal (vis='phase_t5_data',spw='4,5,6,7',gaintable=['tsys_phase.fdm'],
        gainfield=[''],interp = 'linear', calwt = T,flagbackup = F)

#----- Split out relevant (tsys-corrected) spws from phase_t5_data and callisto_data
os.system('rm -rf phase_t5_tmp')
print "<< Doing split... "
split (vis='phase_t5_data',outputvis='phase_t5_tmp',
       datacolumn='corrected',field='0',spw='4,5,6,7')
os.system('rm -rf callisto_tmp')
print "<< Doing split... "
split (vis='callisto_data',outputvis='callisto_tmp',
       datacolumn='corrected',field='0',spw='4,5,6,7')

#----- Concatenate ampcal and phasecal (tsys corrected) data
os.system('rm -rf fluxcal_data')
print "<< Doing concatenate... "
concat(vis=['callisto_tmp','phase_t5_tmp'],concatvis='fluxcal_data')

# ----- Concatenated data set has:
# Callisto - Field 0
# J0607-085 - Field 1
# T3 has spw 0,1,2,3
# T5 has spw 4,5,6,7
#----- We have now Tsys-corrected Field 0 (spw 0,1,2,3) and Field 1 (spw 4,5,6,7)

#----- Flagging for shadowing
print "<< Doing flagging... "
flagdata(vis = 'fluxcal_data',mode = 'shadow',flagbackup = F)

#----- Fill model column for Solar System Ampcal (Callisto - Field 0)
print "<< Doing setjy... "
setjy(vis = 'fluxcal_data', field='0',standard = 'Butler-JPL-Horizons 2010')

#----- Do amplitude calibration for the phase calibrator in T3 
#----- Selfcal in phase only
os.system('rm -rf '+basename_all[0]+'.a2_corr')
print "<< Doing gaincal... "
gaincal(vis = 'fluxcal_data',caltable = basename_all[0]+'.a2_corr',
        field = '0,1',solint = 'int',refant = 'DA43',
        gaintype = 'G',calmode = 'p')

#----- Selfcal in amp only
os.system('rm -rf '+basename_all[0]+'.a3_corr')
print "<< Doing gaincal... "
gaincal(vis = 'fluxcal_data', caltable = basename_all[0]+'.a3_corr',
        field = '0,1',solint = 'inf',refant = 'DA43',
        gaintype = 'T',calmode = 'a',gaintable = [basename_all[0]+'.a2_corr'])

#----- Scale gain solution for phasecal relative to that of the ampcal
#----- Transfer spws from T3 ampcal to T5 phasecal (both have Tsys applied)
os.system('rm -rf '+basename_all[0]+'.a4_corr')
print "<< Doing fluxscale... "
fluxscale(vis = 'fluxcal_data',caltable = basename_all[0]+'.a3_corr',
        fluxtable = basename_all[0]+'.a4_corr',reference = '0',
        transfer='1',refspwmap = [0,1,2,3,0,1,0,1])

# This gives the following fluxes for the phasecal J0607-085: spw0 = 1.38 Jy, spw1 = 1.37 Jy, spw2 = 1.47 Jy, spw3 = 1.46 Jy. SNR~70-80.The SMA flux measurement from Jan 9 ~ 1.3 Jy, which is quite close (to within 10 %). So adopt the flux of the phasecal to be 1.4 Jy. In comparison, dataset 3 where Mars was used as an ampcal gives a phasecal flux ~1.1 Jy, somewhat lower. But SNR~20, so much lower. Mars was observed at low elevation and may actually be partially resolved out at B6, yielding a flux that is too low. Therefore, I ignore this measurement.

#----------------------------------------------------------------------------------
#----- C- Calibrate the phasecal/bandpasscal and science data
#----------------------------------------------------------------------------------

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

#----- Split out the data that we'll need (WVR-corrected phase and science FDM spws)
for basename in basename_all:
	os.system('rm -rf '+basename+'.ms.wvr')
	print "<< Doing split... "
	split (vis=basename+'.ms',outputvis=basename+'.ms.wvr',
	       datacolumn='corrected',field='2,3',
	       spw='33,35,37,39,65,67,69,71,73,75,77,79,81,83,85,87,89,91,93,95',
	       keepflags=F)

#----- Clear pointing table
for basename in basename_all:
	tb.open(basename+'.ms.wvr/POINTING', nomodify = False)
	a = tb.rownumbers()
	tb.removerows(a)
	tb.close()

#----- RUN ANOTHER LISTOBS TO GET INFORMATION
for basename in basename_all:
	os.system('rm -rf '+basename+'.wvr.listobs.txt')
	listobs(vis=basename+'.ms.wvr', verbose=True, 
            listfile=basename+'.wvr.listobs.txt')
	print "<< Printed listobs output to "+basename+'.wvr.listobs.txt'
#----- Calibrator is now field '0' and science target is field '1'

#----- Further data flagging: shadowing
for basename in basename_all:
	print "<< Doing flagging... "
	flagdata(vis = basename+'.ms.wvr',mode = 'shadow',flagbackup = F)

#----- Further data flagging: for first dataset only flag scan8, DV16
#----- for both phase and science as it shows very low amplitudes
print "<< Doing flagging... "
flagdata(vis = basename_all[0]+'.ms.wvr',mode = 'manualflag',
	 selectdata= T, antenna = 'DV16',scan = '8',flagbackup = F)

#----- Bandpass calibration
for basename in basename_all:
	os.system('rm -rf '+basename+'.g1_corr')
	print "<< Doing gaincal... "
	gaincal(vis = basename+'.ms.wvr',caltable = basename+'.g1_corr',
		field ='0',spw = '*:1800~2200',solint = 'int',
		refant = 'DA43',calmode = 'ap')

for basename in basename_all:
	os.system('rm -rf '+basename+'.b1_corr')
	print "<< Doing bandpass... "
	bandpass(vis = basename+'.ms.wvr',caltable = basename+'.b1_corr',
		 field ='0',solint = 'inf', refant = 'DA43',
		 solnorm = T,bandtype = 'B',gaintable = basename+'.g1_corr')

#----- To make figures of bandpass plots: the bandpass just computed and
#----- a smoothed version that is already provided
for basename in basename_all:
	if(calplots):
		os.system('rm -rf bandpass_plots/')
		os.system('mkdir bandpass_plots/')
		print "<< Plotting bandpass to file... "
		plotcal(caltable=basename+'.b1_corr', xaxis='chan', yaxis='amp',
			subplot=11, figfile='bandpass_plots/'+basename+'.b1_corr.png')
		plotcal(caltable=basename+'.b1_corr_smooth', xaxis='freq', yaxis='amp',
			subplot=11, figfile='bandpass_plots/'+basename+'.b1_corr_smooth.png')

#----- Set the flux of the phasecal to the value determined previously
for basename in basename_all:
	print "<< Doing setjy... "
	setjy(vis=basename+'.ms.wvr',field='0',fluxdensity=[1.4,0,0,0])

#----- Do gain calibration (self-cal on phasecal)
for basename in basename_all:
	os.system('rm -rf '+basename+'.g2_corr')
	print "<< Doing gaincal... "
	gaincal(vis = basename+'.ms.wvr',caltable = basename+'.g2_corr',
		field = '0',solint = 'int',refant = 'DA43',
		gaintype = 'G',calmode = 'p',gaintable=basename+'.b1_corr')

for basename in basename_all:
	os.system('rm -rf '+basename+'.g3_corr')
	print "<< Doing gaincal... "
	gaincal(vis = basename+'.ms.wvr', caltable = basename+'.g3_corr',
		field = '0',solint = 'inf',refant = 'DA43',
		gaintype = 'T',calmode = 'a',gaintable = [basename+'.b1_corr',basename+'.g2_corr'])

#----- Determine phase gain solution
for basename in basename_all:
	os.system('rm -rf '+basename+'.g5_corr')
	print "<< Doing gaincal... "
	gaincal(vis = basename+'.ms.wvr', caltable = basename+'.g5_corr',
		field = '0',solint = 'inf',refant = 'DA43',gaintype = 'G',
		calmode = 'p',gaintable = basename+'.b1_corr')

#----- Check the gaincal tables
for basename in basename_all:
	if(calplots):
		plotcal(caltable = basename+'.g2_corr',xaxis = 'time',yaxis = 'phase',
			subplot = 221,iteration = 'spw,antenna',
			plotrange = [0, 0, -180, 180])
		dummy_string = raw_input("Hit <Enter> to continue...")
		plotcal(caltable = basename+'.g3_corr',xaxis = 'time',yaxis = 'amp',
			subplot = 221,iteration = 'spw,antenna')
		dummy_string = raw_input("Hit <Enter> to continue...")
		plotcal(caltable = basename+'.g5_corr',xaxis = 'time',yaxis = 'phase',
			subplot = 221,iteration = 'spw,antenna',
			plotrange = [0, 0, -180, 180])
		dummy_string = raw_input("Hit <Enter> to continue...")

#----- apply the smoothed bandpass and the gaincal to calibrator
#----- (the smoothed bandpass is provided)
for basename in basename_all:
	print "<< Doing applycal... "
	applycal(vis = basename+'.ms.wvr',field = '0',
		gaintable = [basename+'.b1_corr_smooth', basename+'.g2_corr',basename+'.g3_corr'],
		interp = 'linear',calwt = F,flagbackup = F)

#----- apply the smoothed bandpass and the gaincal to science target
for basename in basename_all:
	print "<< Doing applycal... "
	applycal(vis = basename+'.ms.wvr',field = '1',
		gaintable = [basename+'.b1_corr_smooth', basename+'.g3_corr', basename+'.g5_corr'],
		gainfield = ['', '0', '0'],interp = 'linear',calwt = F,flagbackup = F)

#======================================
# All three datasets are now calibrated
#======================================

#----- Combine the three datasets
print "<< Concatenating the three datasets... "
cal_vis = [name+'.ms.wvr' for name in basename_all]
os.system('rm -rf comb_data')
concat(vis=cal_vis, concatvis='comb_data', freqtol='500kHz')

#----- Split out corrected data column
print "<< Splitting out the calibrated combined data... "
split(vis = 'comb_data',outputvis = 'Orion_all.ms', datacolumn = 'corrected',keepflags = T)

#----- This dataset has 20 spectral windows for each the phasecal and the science target

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