###############################################################################
#                                                                             #
# Use Python Script for testing CASA (4.0 and newer versions)                 #
# Jun-Hui Zhao (SAO):                                                         #
# created 2012 - for correcting variable point source of Sgr A*               #
# updated 2014 - for A-array data                                             #
# added comments September 2015 -                                             #  
###############################################################################
#
#Baseline corrections: http://www.vla.nrao.edu/astro/archive/baselines/
#
#!/usr/bin/python
from sys import exit
import time
import os, sys
from gencal      import gencal
from plotweather import plotweather
from listobs     import listobs
from plotants    import plotants
from plotms      import plotms
from flagdata    import flagdata
from setjy       import setjy
from gaincal     import gaincal
from plotcal     import plotcal
from bandpass    import bandpass
from applycal    import applycal
from fluxscale   import fluxscale
from split       import split
from clearcal    import clearcal
from uvcontsub   import uvcontsub
from vishead     import vishead
from cvel        import cvel
from clean       import clean
from viewer      import viewer
from immoments   import immoments
from imstat      import imstat
from concat      import concat
from uvsub       import uvsub
from exportuvfits import exportuvfits
from ft          import ft
from taskinit    import *
from hanningsmooth import hanningsmooth
from widebandpbcor import widebandpbcor

T = True
F = False
#
#=====================================================================
#
########################################
#Original data: D
########################################
#
datao   = 'sgra_ad2'
imagr   = 'SGRA_A'
datapfx = 'SGRA.TINTERVAL'
#a total number of bin
ntotal  = str(32)
datain  = [datapfx+'1',datapfx+'2',datapfx+'3',datapfx+'4',datapfx+'5',datapfx+'6',
datapfx+'7',datapfx+'8',datapfx+'9',datapfx+'10',datapfx+'11',datapfx+'12',datapfx+'13',
datapfx+'14',datapfx+'15',datapfx+'16',datapfx+'17',datapfx+'18',datapfx+'19',datapfx+'20',
datapfx+'21',datapfx+'22',datapfx+'23',datapfx+'24',datapfx+'25',datapfx+'26',datapfx+'27',
datapfx+'28',datapfx+'29',datapfx+'30',datapfx+'31',datapfx+'32']
dataall  =  'sgra_d2nstar'
#
########################################
#Corrected data: D*
########################################
#
datafinal = 'sgra_d2star'
########################################


def step01():
 print 'export UVFIT: SGRA.A1.UVFITS' 
 exportuvfits(vis=datao, fitsfile='SGRA.A2.UVFITS',
 datacolumn='corrected', field='SgrAEast', avgchan=8,
 writesyscal=F, multisource=F, combinespw=T,
 padwithflags=T, writestation=T)
########################################
# go to AIPS to determine the time variability of Sgr A* -
# get individual value of flux density for a given time interval 
# with DFTPL, a AIPS task
########################################



########################################
#Before step02() please update the values 
#of following variables
#ID     (1, 2, .... N), 
#trange (08:11:59.5~08:21:59.5) &
#flux   (SgrA* flux density)
#for each of the minor steps in this loop
#each ID number corresponds to a relevant time range "trange"
#and a value of flux density 
########################################
#ID=str(1)
#trange  = '08:11:59.5~08:21:59.5'
#flux    = str(0.760)
#ID=str(2)
#trange  = '08:21:59.5~08:31:59.5'
#flux    = str(0.757)
ID=str(3)
trange  = '08:31:59.5~08:36:59.5'
flux    = str(0.756)
#
#reading DFTPL output to fill in the variables ID,trange and flux
#... until ntotal
#
datatm  = 'SGRA.TINTERVAL'+ID

def step02():
 print 'ID     = '+ID
 print 'datatm = '+datatm
 print 'trange = '+trange
 print 'flux   = '+flux
 print ' '
 print '#######################################'
 print '#step 1. SPLIT'
 print '#######################################'
 print ' '
 print 'bin the original data'
 os.system('rm -rf '+datatm)
 split(vis=datao, datacolumn='corrected', 
 outputvis=datatm, timerange=trange)
 user_check=raw_input('done with split')
 print ' '
 print '#######################################'
 print '#step 2. Addcomponent'
 print '#######################################'
 print ' '
 print 'create SgrA * model for a given time interval'
 os.system('rm -rf sgra_comp.cl')
 print 'addcomponent:'
#flux in cl.addcomponent must be updated
#with a value of SgrA* flux density corrseponding to
#the relevant time interval
 cl.done()
 cl.addcomponent(flux=0.757, fluxunit="Jy", 
     shape="Gaussian", dir="J2000 17h45m40.0409s -29d00m28.118s",
     majoraxis="0.042arcsec", minoraxis='0.023arcsec', 
     positionangle='80.0deg',freq="5.5GHz",
     spectrumtype="spectral index", index=0.11)
 cl.rename('sgra_comp.cl')
 cl.close()
 print ' '
 print '########################################'
 print '#step 3. FT'
 print '########################################'
 print ' '
 print "Fourier transform the model back to the relevant data"
 ft(vis=datatm,complist='sgra_comp.cl')
 user_check=raw_input('done with ft')
 print ' '
 print '########################################'
 print '#4. UVSUB'
 print '########################################'
 print ' '
 print "Taking out the variable source Sgr A*"
 uvsub(vis=datatm,reverse=F)
 user_check=raw_input('done with uvsub')
 print 'imaging the binned data & checking: '+imagr+'.*'
 os.system('rm -rf '+imagr+'.*')
 clean(vis=datatm,imagename=imagr,
      imagermode='csclean', uvrange='',
      imsize=1500,cell='0.075arcsec',spw='0~15:5~59',
      mode='mfs', multiscale = [0,5,15],
      weighting='briggs',robust=-.0,
      interactive=F, threshold='10mJy', niter=10)
 user_check=raw_input('hit Return to continue script for viewer\n')
 viewer(imagr+'.image')
 s1 = imstat(imagr+'.image', box='1197,548,1245,641')
 print 'Completei step '+ID+' of the SgrA* subtraction loop'
 print 'with a total number of step '+ntotal

def step03():
 print ' '
 print '########################################'
 print '#5. CONCAT'
 print '########################################'
 print ' '
 print 'combine data together into '+dataall+' from '+datapfx+'1'+' to '+datapfx+ntotal
 os.system('rm -rf '+dataall)
 concat(vis=datain,concatvis=dataall,timesort=F)
 user_check=raw_input('done with concat')
 os.system('rm -rf '+imagr+'.*')
 clean(vis=dataall,imagename=imagr,
      imagermode='csclean', uvrange='',
      imsize=1500,cell='0.075arcsec',spw='0~15:5~59',
      mode='mfs', multiscale = [0,5,15],
      weighting='briggs',robust=0, interactive=T,
      nterms=1,threshold='0.1mJy',niter=100)
 s1 = imstat(imagr+'.image', box='1097,448,1245,641')
 print 'assembled data without Sgr A*!'

def step04():
 print ' '
 print '########################################'
 print '#6. Addcomponent'
 print '########################################'
 print ' '
 print "create SgrA * model with a value of the mean flux density"
 os.system('rm -rf sgra_comp.cl')
 print "addcomponent:"
#
# update flux with a mean value of SgrA*'s flux in the period of observation
#
 cl.done()
 cl.addcomponent(flux=0.8, fluxunit="Jy", 
     shape="Gaussian", dir="J2000 17h45m40.0409s -29d00m28.118s",
     majoraxis="0.042arcsec", minoraxis='0.023arcsec', 
     positionangle='80.0deg',freq="5.5GHz",
     spectrumtype="spectral index", index=0.11)
 cl.rename('sgra_comp.cl')
 cl.close()
 print ' '
 print '########################################'
 print '#7. FT'
 print '########################################'
 print ' '
 ft(vis=dataall,complist='sgra_comp.cl')
 user_check=raw_input('done with ft')
 print ' '
 print '########################################'
 print '#8. UVSUB'
 print '########################################'
 print ' '
 print "Add back Sgr A* with a mean value of flux density"
 uvsub(vis=dataall,reverse=T)
 user_check=raw_input('done with uvsub')
 print 'imaging the binned data & checking: '+imagr+'.*'
 os.system('rm -rf '+imagr+'.*')
 clean(vis=dataall,imagename=imagr,
      imagermode='csclean', uvrange='',
      imsize=1500,cell='0.075arcsec',spw='0~15:5~59',
      mode='mfs', multiscale = [0,5,15],
      weighting='briggs',robust=-.0,
      interactive=F, threshold='10mJy', niter=10)
 s1 = imstat(imagr+'.image', box='1197,548,1245,641')
 print 'Complete the loop to add back Sgra* '

def step05():
 print ' '
 print '########################################'
 print '#Corrected Data: D*'
 print '########################################'
 print ' '
 os.system('rm -rf '+datafinal)
 split(vis=dataall, datacolumn='corrected',
 outputvis=datafinal)
 print 'created final corrected data: D* = '+datafinal
 
 print 'a superb job!'
 print ' '
######################################## 
#Done! 
########################################

def step00():
 print('a good try!')

def start():
    print "You are in casapy."
    print "Creat UVFITS for AIPS:                01"
    print "Removal of Sgr A*                     02"
    print "Assemble visdata with no SgrA*        03"
    print "Add back SgrA* with a mean flux       04"
    print "Create final corrected data           05"
    print "Which step do you take?                 "
    next = raw_input("> ")
    if "00" in next:
          step0()
          exit(0)
    if "01" in next:
          step01()
          exit(0)
    if "02" in next:
          step02()
          exit(0)
    if "03" in next:
          step03()
          exit(0)
    if "04" in next:
          step04()
          exit(0)
    if "05" in next:
          step05()
          exit(0)
    else:
      print "Ooops, a dead path."
      exit(0)
#
#=====================================================================
# Done
#
print "Done with the reduction"
start()