##############################################################################################
                     #################################### General Information #####################################
                     ##############################################################################################

##### The following script has been designed to calibrate one execution (ASDM: SDP81_B4_uncalibrated) of the Band 4 observations of SDP81 obtained during the long baseline campaign: https://casaguides.nrao.edu/index.php/ALMA2014_LBC_SVDATA#SDP.81_-_Lensed_Galaxy

#####  This script has been validated for execution in CASA 6.4. It makes use of the analysisutils environment (aU), which can be downloaded separately from CASA (see https://casaguides.nrao.edu/index.php?title=Analysis_Utilities).

##### This script can be either executed at once (execfile('calibration.py'), or the different steps to be executed one by one, by copying them in the CASA terminal

##### This script starts from where data_prep.py left off.

##### At any time, if you think you modified SDP81_B4_uncalibrated.ms.split incorrectly, you can erase it from your work area and copy it 'fresh' from the common repository



                     ##############################################################################################
                     ########################################## Startup ###########################################
                     ##############################################################################################


##### start CASA in your terminal: type casa


##### importing packages and aliases

import re
import os
import analysisUtils as aU
es = aU.stuffForScienceDataReduction() 



                     ##############################################################################################
                     ######################## Getting Oriented and Initial Flagging ###############################
                     ##############################################################################################



##### listobs: prints the metadata associated to the measurement set to the logger

listobs("SDP81_B4_uncalibrated.ms.split")

##### we recommend recording the listobs output to a text file, with the following command:
os.system('rm -fr SDP81_B4_uncalibrated.ms.split.listobs')   
listobs(vis = 'SDP81_B4_uncalibrated.ms.split',
  listfile = 'SDP81_B4_uncalibrated.ms.split.listobs')

##### plot the antenna distribution
os.system('rm -fr plotants.png')   
plotants("SDP81_B4_uncalibrated.ms.split", figfile="plotants.png")

##### plot the uncalibrated data
plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="time", yaxis="amp", averagedata=True,avgchannel="1e3", coloraxis="field")
plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="u", yaxis="v", averagedata=True,avgchannel="1e3", coloraxis="field",customsymbol=True,symbolshape="circle",symbolsize=5)

##### This command will restore the initial state of flags of SDP81_B4_uncalibrated.ms.split (only useful if the this is not the first time you are running the script).

if os.path.exists('SDP81_B4_uncalibrated.ms.split.flagversions/')==True:
    flagmanager(vis = 'SDP81_B4_uncalibrated.ms.split',
      mode = 'restore',
      versionname = 'Original')
  
##### This command will take a snapshot of the state of flags of SDP81_B4_uncalibrated.ms.split . The saved flag table, Original.flags, can be restored at any time using mode='restore'

if os.path.exists('SDP81_B4_uncalibrated.ms.split.flagversions/')==False:
    flagmanager(vis = 'SDP81_B4_uncalibrated.ms.split',
      mode = 'save',
      versionname = 'Original')

##### We proceed with a few initial flags on data which we know not to be usable even before inspection
##### Flagging shadowed data and edge channels

flagdata(vis = 'SDP81_B4_uncalibrated.ms.split',
    mode = 'shadow',
    flagbackup = False)

flagdata(vis = 'SDP81_B4_uncalibrated.ms.split',
    mode = 'manual',
    spw = '0:0~3;60~63, 1:0~3;60~63,2:0~3;60~63,3:0~7;1912~1919',
    flagbackup = False)
  
############### ATTENTION: If this is your first time running the script you do not need to fill the flagging section below. When you re-run the script after a first pass at calibration and inspection, this is where you will put the flags that you have chosen to apply based on your data inspection.


############### Begin Flagging section ###############


############### End Flagging section ###############




                     ##############################################################################################
                     #################################### Bandpass calibration ####################################
                     ##############################################################################################


##### verify in the listobs which source is the bandpass calibrator

##### solve for phase time variations on the bandpass calibrator
##### We use the integration timescale as the interval to solve for time variations (solint = 'int'), but a larger solution interval can be used if the solutions look noisy.


os.system('rm -rf phase_int_bpass.cal')   
gaincal(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'phase_int_bpass.cal',
    field = '0', # J0825+0309
    spw = '0:22~42, 1:22~42,2:22~42,3:800~1200',
    scan = '3',
    solint = 'int',
    refant = 'DA56',
    calmode = 'p')

##### plot the phase solutions

plotms(vis="phase_int_bpass.cal", xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw="0", coloraxis='corr', plotrange=[0,0,-180,180])
plotms(vis="phase_int_bpass.cal", xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw="1", coloraxis='corr', plotrange=[0,0,-180,180])
plotms(vis="phase_int_bpass.cal", xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw="2", coloraxis='corr', plotrange=[0,0,-180,180])
plotms(vis="phase_int_bpass.cal", xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw="3", coloraxis='corr', plotrange=[0,0,-180,180])

##### create the bandpass solution
  
os.system('rm -rf bandpass.cal') 
bandpass(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'bandpass.cal',
    field = '0', # J0825+0309
    solint = 'inf',
    scan = '3',
    combine = 'scan',
    refant = 'DA56',
    solnorm = True,
    bandtype = 'B',
    gaintable = 'phase_int_bpass.cal')

##### plot the bandpass solutions

plotbandpass(caltable="bandpass.cal", xaxis="chan", yaxis="both",  subplot=42)

##### create a smoother bandpass solution for spw 3 which has narrower channels

os.system('rm -rf bandpass_smooth.cal') 
bandpass(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'bandpass_smooth.cal',
    field = '0', # J0825+0309 
    spw = '0,1,2',
    scan = '3',
    solint = 'inf',
    combine = 'scan',
    refant = 'DA56',
    solnorm = True,
    bandtype = 'B',
    append = False,
    gaintable = 'phase_int_bpass.cal')

bandpass(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'bandpass_smooth.cal',
    field = '0', # J0825+0309 
    spw = '3',
    scan = '3',
    solint = 'inf,5ch',
    combine = 'scan',
    refant = 'DA56',
    solnorm = True,
    bandtype = 'B',
    append = True,
    gaintable = 'phase_int_bpass.cal')


##### plot the smooth bandpass solutions

plotbandpass(caltable="bandpass_smooth.cal", xaxis="chan", yaxis="both",  subplot=42)
 

##### apply the bandpass solutions to check if they correct well the bandpass calibrator

applycal(vis="SDP81_B4_uncalibrated.ms.split",field="0",gaintable=["bandpass_smooth.cal", "phase_int_bpass.cal"],interp=["linear","linear"],gainfield=["0","0"],applymode='calonly')


##### Check if the application worked well.
##### Before:
plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="chan", yaxis="phase",  ydatacolumn="data", field="0", averagedata=True, avgtime="1e3", coloraxis="corr")
##### After:
plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="chan", yaxis="phase",  ydatacolumn="corrected", field="0", averagedata=True, avgtime="1e3", coloraxis="corr")
##### Before:
plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="chan", yaxis="amp",  ydatacolumn="data", field="0", averagedata=True, avgtime="1e3", coloraxis="corr")
##### After:
plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="chan", yaxis="amp",  ydatacolumn="corrected", field="0", averagedata=True, avgtime="1e3", coloraxis="corr")

##### you can also plot only a selection of spectral windows by setting the spw parameter, i.e., spw='0,1,2' or spw = '3'

                     ##############################################################################################
                     ###################################### Gain calibration ######################################
                     ##############################################################################################
##### set model for reference quasar
##### query the calibrator catalog for all sources
aU.getALMAFluxForMS('SDP81_B4_uncalibrated.ms.split')

##### The outputs are printed in the terminal
##### Our flux reference source is J0854+2006 (field 1). So we use the results of aU.getALMAFluxForMS for field 1 in the following setjy command.

##### apply the model for source 1 (flux calibrator) to the data
setjy(vis = 'SDP81_B4_uncalibrated.ms.split',
    standard = 'manual',
    field = '1', #J0854+2006'
    fluxdensity = [3.986837, 0, 0, 0],
    spix = -0.456158813,
    reffreq = '149.593012274GHz')


##### Gain calibration: long-term phase solutions 
os.system('rm -rf phase_inf.cal') 
gaincal(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'phase_inf.cal',
    field = '0~2', # J0825+0309,J0854+2006,J0909+0121
    solint = 'inf',
    refant = 'DA56',
    gaintype = 'G',
    calmode = 'p',
    gaintable = 'bandpass_smooth.cal')

##### plot the phase calibration table
plotms(vis="phase_inf.cal",xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw='0', coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10, plotfile="ss20_phase_scan.png")
plotms(vis="phase_inf.cal",xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw='1', coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10)#, plotfile="ss20_phase_scan.png")
plotms(vis="phase_inf.cal",xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw='2', coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10)#, plotfile="ss20_phase_scan.png")
plotms(vis="phase_inf.cal",xaxis="time",yaxis="phase", gridrows=3, gridcols=3, iteraxis="antenna", spw='3', coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10)#, plotfile="ss20_phase_scan.png")

##### Gain calibration: short-term phase solutions 
os.system('rm -rf phase_int.cal') 
gaincal(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'phase_int.cal',
    field = '0~2', # J0825+0309,J0854+2006,J0909+0121
    solint = 'int',
    refant = 'DA56',
    gaintype = 'G',
    calmode = 'p',
    gaintable = 'bandpass_smooth.cal')

##### plot the phase calibration table
plotms(vis="phase_int.cal",xaxis="time",yaxis="phase",gridrows=3, gridcols=3, iteraxis="antenna", spw="0", coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10, plotfile="ss20_phase_int.png")
plotms(vis="phase_int.cal",xaxis="time",yaxis="phase",gridrows=3, gridcols=3, iteraxis="antenna", spw="1", coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10)
plotms(vis="phase_int.cal",xaxis="time",yaxis="phase",gridrows=3, gridcols=3, iteraxis="antenna", spw="2", coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10)
plotms(vis="phase_int.cal",xaxis="time",yaxis="phase",gridrows=3, gridcols=3, iteraxis="antenna", spw="3", coloraxis='corr', plotrange=[0,0,-180,180], symbolsize=10)

##### Gain calibration: long-term amplitude solutions 
os.system('rm -rf ampli_inf.cal') 
gaincal(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'ampli_inf.cal',
    field = '0~2', # J0825+0309,J0854+2006,J0909+0121
    solint = 'inf',
    refant = 'DA56',
    gaintype = 'T',
    calmode = 'a',
    gaintable = ['bandpass_smooth.cal', 'phase_int.cal'])

##### plot the amplitude calibration table
plotms(vis="ampli_inf.cal",xaxis="time",yaxis="amp",gridrows=3,gridcols=3,iteraxis="antenna",spw='0', coloraxis='corr', plotrange=[0,0, 0.125,0.15],symbolsize=10, field='2',plotfile="ss20_ampli_scan.png")
plotms(vis="ampli_inf.cal",xaxis="time",yaxis="amp",gridrows=3,gridcols=3,iteraxis="antenna",spw='1', coloraxis='corr', plotrange=[0,0, 0.125,0.15],symbolsize=10, field='2')
plotms(vis="ampli_inf.cal",xaxis="time",yaxis="amp",gridrows=3,gridcols=3,iteraxis="antenna",spw='2', coloraxis='corr', plotrange=[0,0, 0.125,0.15],symbolsize=10, field='2')
plotms(vis="ampli_inf.cal",xaxis="time",yaxis="amp",gridrows=3,gridcols=3,iteraxis="antenna",spw='3', coloraxis='corr', plotrange=[0,0, 0.125,0.15],symbolsize=10, field='2')

                     ##############################################################################################
                     ################################## Setting the flux scale ####################################
                     ##############################################################################################

##### The next 4 lines are designed to create a text file with the derived fluxes for the bandpass and phase calibrators (for future reference). 
os.system('rm -rf flux_inf.cal') 
os.system('rm -rf fluxscale.txt') 
mylogfile = casalog.logfile()
casalog.setlogfile('fluxscale.txt')

##### Bootstrapping the flux of calibrators based on the reference source J0854+2006  
fluxscaleDict = fluxscale(vis = 'SDP81_B4_uncalibrated.ms.split',
    caltable = 'ampli_inf.cal',
    fluxtable = 'flux_inf.cal',
    reference = '1') # J0854+2006  

casalog.setlogfile(mylogfile)

##### plot the flux scale solutions
plotms(vis="flux_inf.cal",xaxis="time",yaxis="amp",gridcols=3,gridrows=3,iteraxis="antenna", plotrange=[0,0,0.13,0.18],symbolsize=10, plotfile="ss20_flux_scan.png")


                     ##############################################################################################
                     #################################### Applying calibrations ###################################
                     ##############################################################################################

flagmanager(vis = 'SDP81_B4_uncalibrated.ms.split',
    mode = 'save',
    versionname = 'BeforeApplycal')

for i in ['0', '1']: # J0825+0309,J0854+2006
    applycal(vis = 'SDP81_B4_uncalibrated.ms.split',
      field = str(i),
      gaintable = ['bandpass_smooth.cal', 'phase_int.cal', 'flux_inf.cal'],
      gainfield = ['', i, i],
      interp = 'linear,linear',
      calwt = True,
      flagbackup = False)
  
applycal(vis = 'SDP81_B4_uncalibrated.ms.split',
    field = '2,3', # SDP.81
    gaintable = ['bandpass_smooth.cal', 'phase_inf.cal', 'flux_inf.cal'],
    gainfield = ['', '2', '2'], # J0909+0121
    interp = 'linear,linear',
    calwt = True,
    flagbackup = False)
  

                     ##############################################################################################
                     #################################### Inspect the Results #####################################
                     ##############################################################################################

##### Below are a few suggestion of how to use plotms to inspect your data. You can modify averaging, data selection and data display directly in the GUI or enter a new call to plotms at the casa prompt as is done in the examples below.

##### As a function of uvdistance: identify an antenna which is bad at all/most times
#plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="uvdist", yaxis="amp", ydatacolumn="corrected", field="0,2,3", averagedata=True,   avgchannel="1e3", avgtime="1e3", iteraxis='field', coloraxis="corr")

#plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="uvdist", yaxis="phase", ydatacolumn="corrected", field="0,2,3", averagedata=True,   avgchannel="1e3", avgtime="1e3", iteraxis='field', coloraxis="corr")


##### As a function of time: look for scan to scan variations in amplitude
#plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="time", yaxis="amp", ydatacolumn="corrected", field="0,2,3", averagedata=True,   avgchannel="1e3", avgtime="1e3", coloraxis="field")


##### As a function of time: look for scan to scan variations in phase
#plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="time", yaxis="phase", ydatacolumn="corrected", field="0,2,3", averagedata=True,   avgchannel="1e3", avgtime="1e3", coloraxis="field")


##### As a function of spectral channel (look for spikes in the data)
##### spectral windows with wide channels
#plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="channel", yaxis="amp", ydatacolumn="corrected", field="0,1,2,3", averagedata=True,   avgchannel="", avgtime="1e6", coloraxis="spw",iteraxis='field',spw='0,1,2',avgantenna=True)

##### spectral window with narrow channels
#plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="channel", yaxis="amp", ydatacolumn="corrected", field="0,1,2,3", averagedata=True,   avgchannel="", avgtime="1e6", coloraxis="corr",iteraxis='field',spw='3',avgantenna=True)

#plotms(vis="SDP81_B4_uncalibrated.ms.split", xaxis="channel", yaxis="amp", ydatacolumn="corrected", field="0,1,2,3", averagedata=True,   avgchannel="", avgtime="1e6", coloraxis="corr",iteraxis='field',spw='3',avgbaseline=True)

#### the atmospheric line between channel 1400 and 1500 has affected the bandpass correction significantly. In addition, the data is noisier in that section. Unless the atmospheric line coincides with the targeted line, it is better to flag that section

######################## Define the list of flags

###### ATTENTION: this section is only used as a place to collect the flag definitions you determine based on your inspection of the data. The flag commands should stay commented out here. Once you are satisfied with your chosen flags, you will need to:
# 1) copy the flags you define below and enter them (without the "#" to comment them out) in the 'Flagging section' earlier in the script (before bandpass calibration)
# 2) re-run the entire calibration.py script from scratch using execfile('calibration.py') at the casa prompt

## Outlier antenna in spw2
#flagdata(vis = 'SDP81_B4_uncalibrated.ms.split',
#    mode = 'manual',
#    spw = '2',
#    antenna= 'PM04',
#    flagbackup = False)


## Feel free to add your own!


################   Re-run the script

### to accelerate this new script run, we recommend commenting out all call to plotms
### Note that the flagmanager command at the top of the script will restore the original flag state associated to SDP81_B4_uncalibrated.ms.split. This is necessary to make sure that we are running the calibration on the original SDP81_B4_uncalibrated.ms.split measurement set.

### Several iterations of inspection, defining flags and re-calibration can be made.


                     ##############################################################################################
                     ####################################### Applying Renorm ######################################
                     ##############################################################################################


### This dataset does not need renorm, however the example commands below are provided. A detailed description about renormalization
#and how this process is carried out can be found at the following link:
#https://help.almascience.org/kb/articles/what-are-the-amplitude-calibration-issues-caused-by-alma-s-normalization-strategy. 


#from pipeline.extern.almarenorm import ACreNorm

#applyRenorm = True

#if applyRenorm or applyonly!=True:

#  if not applyRenorm: os.system('rm -rf RN_plots')

#  rn=ACreNorm('SDP81_B4_uncalibrated.ms.split')

#  corrApplied = rn.checkApply()
#  if corrApplied:
#    casalog.post(' The renormalization has already been applied to these data', 'WARN')
#    if applyRenorm:
#        sys.exit('Do not rerun the renorm step again. Renorm was applied already.')

#  if applyRenorm or applyonly==True:
#    diagSpectra=False
#  else:
#    diagSpectra=True

#  rn.renormalize(docorr=applyRenorm, diagSpectra=diagSpectra, antHeuristicsSpectra=False, atmAutoExclude=True)

#  if applyRenorm or applyonly==True:
#    pass  # no plotting
#  else:
#    rn.plotSpectra()

#  if os.path.exists('RN_plots') and applyRenorm:
#    os.system('rm -f RN_plots/.pdf RN_plots/_scan*_field*.png')
#    os.system('mv RN_plots SDP81_B4_uncalibrated.ms.split.renorm.plots')

#  if not applyRenorm:
#    mystats = rn.rnpipestats
#    casalog.post('Maximum Renormalization scaling factors:\n'+str(mystats))

#    rnfactors = []
#    for field in mystats.keys():
#      print(' ************************')
#      print(field, ': max. renorm. scaling for each SPW')
#      for spw in mystats[field].keys():
#        print(spw,': ', mystats[field][spw])
#        rnfactors.append(mystats[field][spw]['max_rn'])
#    maxScaling = max(rnfactors)
#    print(' ************************')
#    print(' The maximum Renormalization scaling is '+str(maxScaling))
#    if maxScaling>1.02:
#      casalog.post('Renormalization should be applied before proceeding.','WARN')
#      rn.close()
#      sys.exit('Please edit the renorm step to set applyRenorm=True and rerun the step!')
#    else:
#      print(' No Renormalization needed.')
#    print(' ************************')

#  rn.close()

#### Typically after one is satisfied by the archived data quality, it is recommended to split out the corrected column of the data to a new measurement set. We will not do it in this tutorial in the interest of saving space. For future reference, this is how splitting out the correct column can be done:

#  split(vis = 'SDP81_B4_uncalibrated.ms.split',
#    outputvis = 'SDP81_B4_calibrated.ms.split',
#    datacolumn = 'corrected',
#    keepflags = True)


############# Wrap up:

####  To free space on your machine, you can remove SDP81_B4_uncalibrated.ms.split from your Calibration directory when you are done with the calibration.
#### If, at any time you want to get the original (uncalibrated) measurement set back, you can get it from the common repository