#!/usr/bin/env python
# ALMA Data Reduction Script

# Calibration

thesteps = []
step_title = {0: 'Import of the ASDM',
              1: 'listobs',
              2: 'Split by antenna',
              3: 'sdlist',
              4: 'Generation of the Tsys cal table',
              5: 'Do initial flagging',
              6: 'Calibration of the data into Kelvins',
              7: 'Application of non-linearity correction factor',
              8: 'Subtracting the baseline',
              9: 'Converting ASAP -> MS',
              10: 'Split and concatenation'}

if 'applyonly' not in globals(): applyonly = False
try:
  print 'List of steps to be executed ...', mysteps
  thesteps = mysteps
except:
  print 'global variable mysteps not set.'
if (thesteps==[]):
  thesteps = range(0,len(step_title))
  print 'Executing all steps: ', thesteps

# The Python variable 'mysteps' will control which steps
# are executed when you start the script using
#   execfile('scriptForCalibration.py')
# e.g. setting
#   mysteps = [2,3,4]# before starting the script will make the script execute
# only steps 2, 3, and 4
# Setting mysteps = [] will make it execute all steps.

import re

import filltsys

es = aU.stuffForScienceDataReduction()

if re.search('^4.3.0', casadef.casa_version) == None:
 sys.exit('ERROR: PLEASE USE THE SAME VERSION OF CASA THAT YOU USED FOR GENERATING THE SCRIPT: 4.3.0')


# CALIBRATE_AMPLI: 
# CALIBRATE_ATMOSPHERE: J1215+1654,M100
# CALIBRATE_BANDPASS: 
# CALIBRATE_FLUX: 
# CALIBRATE_FOCUS: 
# CALIBRATE_PHASE: 
# CALIBRATE_POINTING: J1215+1654
# OBSERVE_TARGET: M100


# Import of the ASDM
mystep = 0
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  if os.path.exists('uid___A002_X86fcfa_Xd9.ms') == False:
    importasdm('uid___A002_X86fcfa_Xd9', asis='Antenna Station Receiver Source CalAtmosphere CalWVR', bdfflags=False)
    os.system(os.environ['CASAPATH'].split()[0]+'/bin/bdflags2MS -f "COR DELA INT MIS SIG SYN TFB WVR ZER" uid___A002_X86fcfa_Xd9 uid___A002_X86fcfa_Xd9.ms')
  if applyonly != True: es.fixForCSV2555('uid___A002_X86fcfa_Xd9.ms')

# listobs
mystep = 1
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -rf uid___A002_X86fcfa_Xd9.ms.listobs')
  listobs(vis = 'uid___A002_X86fcfa_Xd9.ms',
    listfile = 'uid___A002_X86fcfa_Xd9.ms.listobs')
  
  aU.getTPSampling(vis = 'uid___A002_X86fcfa_Xd9.ms', showplot = True, plotfile = 'uid___A002_X86fcfa_Xd9.ms.sampling.png')

# Split by antenna
mystep = 2
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  for i in ['DV10', 'PM03', 'PM04']:
    os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.'+i+'*')
  
  sdsave(infile = 'uid___A002_X86fcfa_Xd9.ms',
    splitant = True,
    outfile = 'uid___A002_X86fcfa_Xd9.ms.asap',
    overwrite = True)
  
  

# sdlist
mystep = 3
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.DV10.asap.sdlist')
  sdlist(infile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.sdlist')
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM03.asap.sdlist')
  sdlist(infile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.sdlist')
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM04.asap.sdlist')
  sdlist(infile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.sdlist')
  
  

# Generation of the Tsys cal table
mystep = 4
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.tsys')
  
  gencal(vis = 'uid___A002_X86fcfa_Xd9.ms',
    caltable = 'uid___A002_X86fcfa_Xd9.ms.tsys',
    caltype = 'tsys')
  
  plotbandpass(caltable='uid___A002_X86fcfa_Xd9.ms.tsys', overlay='time', 
    xaxis='freq', yaxis='amp', subplot=22, buildpdf=False, interactive=False,
    showatm=True,pwv='auto',chanrange='92.1875%',showfdm=True, 
    field='', figfile='uid___A002_X86fcfa_Xd9.ms.tsys.plots.overlayTime/uid___A002_X86fcfa_Xd9.ms.tsys') 
  
  es.checkCalTable('uid___A002_X86fcfa_Xd9.ms.tsys', msName='uid___A002_X86fcfa_Xd9.ms', interactive=False)
  

# Do initial flagging
mystep = 5
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  sdflag(infile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap',
    mode = 'manual',
    spw = '17:0~119;3960~4079,19:0~119;3960~4079,21:0~119;3960~4079,23:0~119;3960~4079',
    overwrite = True)
  
  sdflag(infile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap',
    mode = 'manual',
    spw = '17:0~119;3960~4079,19:0~119;3960~4079,21:0~119;3960~4079,23:0~119;3960~4079',
    overwrite = True)
  
  sdflag(infile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap',
    mode = 'manual',
    spw = '17:0~119;3960~4079,19:0~119;3960~4079,21:0~119;3960~4079,23:0~119;3960~4079',
    overwrite = True)
  
  

# Calibration of the data into Kelvins
mystep = 6
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  from recipes.almahelpers import tsysspwmap
  tsysmap = tsysspwmap(vis = 'uid___A002_X86fcfa_Xd9.ms', tsystable = 'uid___A002_X86fcfa_Xd9.ms.tsys')
  
  spwmap = {}
  for i in [17, 19, 21, 23]:
    if not tsysmap[i] in spwmap.keys():
      spwmap[tsysmap[i]] = []
    spwmap[tsysmap[i]].append(i)
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal')
  
  sdcal2(infile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap',
    calmode = 'ps,tsys,apply',
    spw = '9,11,13,15,17,19,21,23',
    tsysspw = '9,11,13,15',
    spwmap = spwmap,
    outfile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal',
    overwrite = True)
  
  es.SDcheckSpectra('uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal', spwIds='17,19,21,23', interactive=False)
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal')
  
  sdcal2(infile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap',
    calmode = 'ps,tsys,apply',
    spw = '9,11,13,15,17,19,21,23',
    tsysspw = '9,11,13,15',
    spwmap = spwmap,
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal',
    overwrite = True)
  
  es.SDcheckSpectra('uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal', spwIds='17,19,21,23', interactive=False)
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal')
  
  sdcal2(infile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap',
    calmode = 'ps,tsys,apply',
    spw = '9,11,13,15,17,19,21,23',
    tsysspw = '9,11,13,15',
    spwmap = spwmap,
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal',
    overwrite = True)
  
  es.SDcheckSpectra('uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal', spwIds='17,19,21,23', interactive=False)
  
  

# Application of non-linearity correction factor
mystep = 7
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc')
  
  sdscale(infile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc',
    factor = 1.25)
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc')
  
  sdscale(infile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc',
    factor = 1.25)
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc')
  
  sdscale(infile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc',
    factor = 1.25)
  
  

# Subtracting the baseline
mystep = 8
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  # The following is a temporary workaround until CASA 4.4 is released.
  tb.open('uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc', nomodify = False)
  subtb = tb.query('FLAGROW==1')
  flaggedrows = subtb.rownumbers()
  if len(flaggedrows) > 0: tb.removerows(flaggedrows)
  subtb.close()
  tb.flush()
  tb.close()
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc.bl')
  
  sdbaseline(infile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc',
    spw = '17,19,21,23',
    maskmode = 'auto',
    thresh = 5.0,
    avg_limit = 4,
    blfunc = 'poly',
    order = 1,
    outfile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc.bl',
    overwrite = True)
  
  es.SDcheckSpectra('uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc.bl', spwIds='17,19,21,23', interactive=False)
  
  # The following is a temporary workaround until CASA 4.4 is released.
  tb.open('uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc', nomodify = False)
  subtb = tb.query('FLAGROW==1')
  flaggedrows = subtb.rownumbers()
  if len(flaggedrows) > 0: tb.removerows(flaggedrows)
  subtb.close()
  tb.flush()
  tb.close()
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc.bl')
  
  sdbaseline(infile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc',
    spw = '17,19,21,23',
    maskmode = 'auto',
    thresh = 5.0,
    avg_limit = 4,
    blfunc = 'poly',
    order = 1,
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc.bl',
    overwrite = True)
  
  es.SDcheckSpectra('uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc.bl', spwIds='17,19,21,23', interactive=False)
  
  # The following is a temporary workaround until CASA 4.4 is released.
  tb.open('uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc', nomodify = False)
  subtb = tb.query('FLAGROW==1')
  flaggedrows = subtb.rownumbers()
  if len(flaggedrows) > 0: tb.removerows(flaggedrows)
  subtb.close()
  tb.flush()
  tb.close()
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc.bl')
  
  sdbaseline(infile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc',
    spw = '17,19,21,23',
    maskmode = 'auto',
    thresh = 5.0,
    avg_limit = 4,
    blfunc = 'poly',
    order = 1,
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc.bl',
    overwrite = True)
  
  es.SDcheckSpectra('uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc.bl', spwIds='17,19,21,23', interactive=False)
  
  

# Converting ASAP -> MS
mystep = 9
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc.bl.ms')
  
  sdsave(infile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc.bl',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc.bl.ms',
    spw = '17,19,21,23',
    outform = 'MS2')
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc.bl.ms')
  
  sdsave(infile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc.bl',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc.bl.ms',
    spw = '17,19,21,23',
    outform = 'MS2')
  
  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc.bl.ms')
  
  sdsave(infile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc.bl',
    outfile = 'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc.bl.ms',
    spw = '17,19,21,23',
    outform = 'MS2')
  
  

# Split and concatenation
mystep = 10
if(mystep in thesteps):
  casalog.post('Step '+str(mystep)+' '+step_title[mystep],'INFO')
  print 'Step ', mystep, step_title[mystep]

  os.system('rm -Rf uid___A002_X86fcfa_Xd9.ms.cal')
  
  concat(vis = [ \
      'uid___A002_X86fcfa_Xd9.ms.DV10.asap.cal.nlc.bl.ms', \
      'uid___A002_X86fcfa_Xd9.ms.PM03.asap.cal.nlc.bl.ms', \
      'uid___A002_X86fcfa_Xd9.ms.PM04.asap.cal.nlc.bl.ms' ], \
    concatvis = 'uid___A002_X86fcfa_Xd9.ms.cal')