##########################################################################
# This casapy script was prepared for the EU ARC Cycle 2 PI CASA tutorial.
# It was first made to work correctly.
# Then, a number of problems were introduced.
# It is your task to fix them.
# You will find a list  of the problems further below.
# Look for the string 'PROBLEM' to find all problems. 
#
# The script uses some infrastructure to make it easier for you
# to run and repeat the script step by step:
# Further below you find the "step_title" Python dictionary which
# shows you what is happening in each analysis step.
# You can start this script by starting casapy and entering
#     execfile('reduce-iras16293-casa4-with-problems.py')
# By default it will run all steps.
# (Of course it will only run correctly when you have fixed all problems.)
# If you would like to run individual steps, you can set the
# variable mysteps before starting the script, e.g.:
#
#     mysteps = [4,5,6]
#     execfile('reduce-iras16293-casa4-with-problems.py')
#
# This example would only execute steps 4, 5, and 6.
# Setting mysteps to an empty list [] will execute all steps.
#
#
# List of problems in the individual analysis steps of this script:


#
##########################################################################

# EU ALMA Regional Centre CASA Tutorial for ALMA Cycle 2
# Sample ALMA data PI analysis script
# "IRAS16293" 

step_title = { 0: 'Preliminary \n'\
               ' (listobs to check data and delete the pointing table)',
	       1: 'Do a number of plotms to see which fields, spws and channels have emission \n'\
               ' (to prepare for continuum self-cal)',
	       2: 'Flagging \n'\
               ' (flag the line emission for now...)',
               3: 'Spliiting \n'\
               ' (Split out a pseudo-continuum ms )',
               4: 'First cleaning before self-cal \n'\
               ' (make the first continuum map before any self-cal)',
               5: 'First self-cal round \n'\
               ' (starting with a phase-only self-cal )',
               6: 'Cleaning after first self-cal round \n'\
               ' (make the continuum map after one round of phase-only self-cal)',
               7: 'Second self-cal round \n'\
               ' (another phase-only self-cal )',
               8: 'Cleaning after second self-cal round \n'\
               ' (make the continuum map after two rounds of phase-only self-cal)',
               9: 'Third self-cal round \n'\
               ' (first a&p self-cal )',
               10: 'Cleaning after third self-cal round \n'\
               ' (make the continuum map after the first a&p self-cal)',
               11: 'Fourth self-cal round \n'\
               ' (second a&p self-cal )',
               12: 'Cleaning after fourth self-cal round \n'\
               ' (make the continuum map after the second a&p self-cal)',
               13: 'Correct for primary beam \n'\
               ' (primary beam correction and write out png image)',
               14: 'Calibrate line data round \n'\
               ' (Apply calibration to the original line data and prepare for imaging)',
               15: 'Subtract the continuum\n'\
               ' WARNING: Only do this if you have some time! This will take a while.',

               }

    
############################################################################

# Some infrastructure to make repeating individual parts
#   of this workflow more convenient.

totaltime = 0
inittime = time.time()
ttime = inittime
steptime = []
thesteps = []

def timing():
    global totaltime
    global inittime
    global ttime
    global steptime
    global step_title
    global mystep
    global thesteps
    thetime = time.time()
    dtime = thetime - ttime
    steptime.append(dtime)
    totaltime += dtime
    ttime = thetime
    casalog.origin('TIMING')
    casalog.post( 'Step '+str(mystep)+': '+step_title[mystep].splitlines()[0], 'WARN')
    casalog.post( 'Time now: '+str(ttime), 'WARN')
    casalog.post( 'Time used this step: '+str(dtime), 'WARN')
    casalog.post( 'Total time used so far: ' + str(totaltime), 'WARN')
    casalog.post( 'Step  Time used (s)  Fraction of total time (percent) [description]', 'WARN')
    for i in range(0, len(steptime)):
        casalog.post( '  '+str(thesteps[i])+'   '+str(steptime[i])+'  '+str(steptime[i]/totaltime*100.)
                      +' ['+step_title[thesteps[i]].splitlines()[0]+']', 'WARN')
        
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 'mysteps empty. Executing all steps: ', thesteps

# The Python variable 'mysteps' will control which steps
# are executed when you start the script using
#   execfile('reduce-m100.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.

##########################################################################

# Begin analysis 


# listobs to check data and delete the pointing table
mystep = 0
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # first import analysisUtils
    import analysisUtils as aU

    listobs(vis='IRAS16293_Band9.fixed.rebin.ms')

    tb.open('IRAS16293_Band9.fixed.rebin.ms/POINTING', nomodify = False)
    a = tb.rownumbers()
    tb.removerows(a)
    tb.close()

    # Make a plot of the mosaic pattern with field ids identified
    aU.plotMosaic(vis='IRAS16293_Band9.fixed.rebin.ms',sourceid='0',
              doplot=True,figfile='mosaic_pattern.png')

# do a number of plotms to see which fields, spws and channels have emission
mystep = 1
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # See which mosaic fields have the brightest line emission
    plotms(vis='IRAS16293_Band9.fixed.rebin.ms', 
           field='',xaxis='uvdist', yaxis='amp',
           spw='1', avgchannel='3840',coloraxis='field',
           iteraxis='field')

    user_check=raw_input('....Hit Return to continue script\n')

    # Iterate over spws of all fields
    plotms(vis='IRAS16293_Band9.fixed.rebin.ms', 
           field='',xaxis='uvdist', yaxis='amp',
           spw='', avgchannel='3840',coloraxis='field',
           iteraxis='PROBLEM')

    user_check=raw_input('....Hit Return to continue script\n')

    # Select channels that contain line emission. Average every 3 channels for
    # increased S/N.
    plotms(vis='IRAS16293_Band9.fixed.rebin.ms', 
           xaxis='channel', yaxis='amp',field='3',
           spw='', avgtime='1e8',avgscan=T,coloraxis='corr',
           avgchannel='3',iteraxis='spw',ydatacolumn='corrected',
           xselfscale=True,yselfscale=True)


# flagging
mystep = 2
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Save the flagging state for later.
    flagmanager(vis='IRAS16293_Band9.fixed.rebin.ms',mode='save',
                versionname='Original')

    # Flagging the line channels
    # Don't worry about the exact channel ranges here. This is just copied from the casa guide
    flagdata(vis='IRAS16293_Band9.fixed.rebin.ms',
             spw='0:1700~2500,1:400~500;1100~1900;3000~3450,2:1700~2200;3100~3839,3:0~800;1600~2900;3300~3839',
             flagbackup=F)



# split and average the data
mystep = 3
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]


    # Average 200 channels, as you do not need to have a high spectral resolution for continuum imaging.
    split(vis='IRAS16293_Band9.fixed.rebin.ms',
          outputvis='IRAS16293_Band9.cont.ms',field='',
          datacolumn='data',width=PROBLEM,spw='0~3:20~3819')


# First cleaning before self-cal
mystep = 4
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]


    # First cleaning before self-cal
    os.system('rm -rf IRAS16293.cont.*')
    clean(vis='IRAS16293_Band9.cont.ms',
          imagename='IRAS16293.cont',
          spw='0~3',field='',phasecenter='3',
          mode='PROBLEM',imagermode='PROBLEM',
          imsize=500,cell='0.04arcsec',minpb=0.25,
          interactive=T,
          weighting='PROBLEM',robust=0.5,
          niter=10000,usescratch=F)

    # checking the model data for all the fields (are there clean components
    # in all fields?)
    plotms(vis='IRAS16293_Band9.cont.ms', 
           xaxis='time', yaxis='amp',field='',iteraxis='PROBLEM',
           spw='', avgchannel='4000',coloraxis='field',
           ydatacolumn='PROBLEM',xselfscale=True,yselfscale=True)


# First self-cal round
mystep = 5
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # find first phase-only self-cal solutions
    # Combine both polarizations to increase signal to noise
    os.system('rm -rf cont.SC.p1')
    gaincal(vis='IRAS16293_Band9.cont.ms',
            caltable='cont.SC.p1',
            gaintype='PROBLEM',
            refant='DV14',calmode='PROBLEM',combine='',spw='',field='',
            solint='PROBLEM',minsnr=3.0,minblperant=4)

    # Check the solutions
    plotcal(caltable='cont.SC.p1',xaxis='time',yaxis='phase',
            spw='',iteration='antenna',subplot=411,plotrange=[0,0,-180,180],
            antenna='',timerange='')
 
    # application of the first phase-only self-cal table
    # what is the correct calwt for self calibration?
    applycal(vis='IRAS16293_Band9.cont.ms',
             gaintable=['cont.SC.p1'],calwt=PROBLEM,flagbackup=F)

# Cleaning after first self-cal round
mystep = 6
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Second cleaning after one round of phase-only self-cal
    os.system('rm -rf IRAS16293.cont.p1.*')
    clean(vis='IRAS16293_Band9.cont.ms',
          imagename='IRAS16293.cont.p1',
          spw='0~3',field='',phasecenter='3',
          mode='PROBLEM',imagermode='PROBLEM',
          imsize=500,cell='0.04arcsec',minpb=0.25,
          interactive=T,
          weighting='briggs',robust=0.5,
          mask='IRAS16293.cont.mask',
          niter=10000)
 
    # visualize the pre-self-cal and post-self-cal images
    imview (raster=[{'file':'IRAS16293.cont.image',
                     'range': [-0.01,0.5], 'scaling': -1},
                    {'file': 'IRAS16293.cont.p1.image',
                     'range': [-0.01,0.5], 'scaling': -1}])


# Second self-cal round
mystep = 7
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # find phase-only solution for second self-cal round
    os.system('rm -rf cont.SC.p2')
    gaincal(vis='IRAS16293_Band9.cont.ms',
            caltable='cont.SC.p2',
            gaintype='PROBLEM',
            refant='DV14',calmode='PROBLEM',combine='',spw='',field='',
            solint='inf',minsnr=3.0,minblperant=4)
 
    # Checking the table
    plotcal(caltable='cont.SC.p2',xaxis='time',yaxis='phase',
            spw='',iteration='antenna',subplot=411,plotrange=[0,0,-180,180],
            antenna='',timerange='')
    
    # Applying results to the data
    applycal(vis='IRAS16293_Band9.cont.ms',
             gaintable=['cont.SC.p2'],calwt=F,flagbackup=F)
    
# Cleaning after second self-cal round
mystep = 8
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Cleaning after second phase-only self-cal
    os.system('rm -rf IRAS16293.cont.p2.*')
    clean(vis='IRAS16293_Band9.cont.ms',
          imagename='IRAS16293.cont.p2',
          spw='0~3',field='',phasecenter='3',
          mode='PROBLEM',imagermode='PROBLEM',
          imsize=500,cell='0.04arcsec',minpb=0.25,
          interactive=T,
          weighting='briggs',robust=0.5,
          mask='IRAS16293.cont.p1.mask',
          niter=10000)
  
    # visualize the images after one round and two rounds of self-cal
    imview (raster=[{'file':'IRAS16293.cont.p1.image',
                     'range': [-0.01,0.5], 'scaling': -1},
                    {'file': 'IRAS16293.cont.p2.image',
                     'range': [-0.01,0.5], 'scaling': -1}])


    # There is very little change, the bold could try another iteration using a shorter
    # solint phase self-cal. For now, moving on to amplitude.

# Third self-cal round
mystep = 9
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # For amplitude self-cal we only want to get one solution per spw, per dataset
    os.system('rm -rf cont.SC.ap1')
    gaincal(vis='IRAS16293_Band9.cont.ms',
            caltable='cont.SC.ap1',
            gaintype='PROBLEM',
            refant='DV14',calmode='PROBLEM',combine='scan,field',spw='',field='',
            solint='8000s',minsnr=3.0,minblperant=4,
            gaintable='cont.SC.p2')

    # plotting the table
    plotcal(caltable='cont.SC.ap1',xaxis='time',yaxis='amp',
            spw='',iteration='antenna',subplot=411,plotrange=[0,0,0.8,1.8],
            antenna='',timerange='')
    
    # and applying to the data
    applycal(vis='IRAS16293_Band9.cont.ms',
             gaintable=[PROBLEM],calwt=F,flagbackup=F)
    

# Cleaning after third self-cal round
mystep = 10
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Fourth cleaning after the a&p self-cal
    os.system('rm -rf IRAS16293.cont.ap1.*')
    clean(vis='IRAS16293_Band9.cont.ms',
          imagename='IRAS16293.cont.ap1',
          spw='0~3',field='',phasecenter='3',
          mode='PROBLEM',imagermode='PROBLEM',
          imsize=500,cell='0.04arcsec',minpb=0.25,
          interactive=T,
          weighting='briggs',robust=0.5,
          mask='IRAS16293.cont.p2.mask',
          niter=10000)

    # visualize the images after phase self-cal and a&p self-cal
    imview (raster=[{'file':'IRAS16293.cont.p2.image',
                     'range': [-0.01,0.5], 'scaling': -1},
                    {'file': 'IRAS16293.cont.ap1.image',
                     'range': [-0.01,0.5], 'scaling': -1}])
    

# Fourth self-cal round
mystep = 11
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Now try amplitude self-cal with higher time resolution
    os.system('rm -rf cont.SC.ap2')
    gaincal(vis='IRAS16293_Band9.cont.ms',
            caltable='cont.SC.ap2',
            gaintype='PROBLEM',
            refant='DV14',calmode='PROBLEM',combine='scan,field',spw='',field='',
            solint='PROBLEM',minsnr=3.0,minblperant=4,
            gaintable='cont.SC.p2')

    # plotting the table
    plotcal(caltable='cont.SC.ap2',xaxis='time',yaxis='amp',
            spw='',iteration='antenna',subplot=411,plotrange=[0,0,0.8,1.8],
            antenna='',timerange='')
    
    # and applying to the data
    applycal(vis='IRAS16293_Band9.cont.ms',
             gaintable=[PROBLEM],calwt=F,flagbackup=F)
    
    
# Cleaning after fourth self-cal round
mystep = 12
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Fifth cleaning after the second a&p self-cal
    os.system('rm -rf IRAS16293.cont.ap2.*')
    clean(vis='IRAS16293_Band9.cont.ms',
          imagename='IRAS16293.cont.ap2',
          spw='0~3',field='',phasecenter='3',
          mode='PROBLEM',imagermode='PROBLEM',
          imsize=500,cell='0.04arcsec',minpb=0.25,
          interactive=T,
          weighting='briggs',robust=0.5,
          mask='IRAS16293.cont.ap1.mask',
          niter=10000)

    # visualize the images after first and second a&p self-cal
    imview (raster=[{'file':'IRAS16293.cont.ap1.image',
                     'range': [-0.01,0.5], 'scaling': -1},
                    {'file': 'IRAS16293.cont.ap2.image',
                     'range': [-0.01,0.5], 'scaling': -1}])
    


# Correct for primary beam and write out image
mystep = 13
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]


    # correcting for primary beam
    impbcor(imagename='IRAS16293.cont.ap2.image',
            pbimage='IRAS16293.cont.ap2.flux',
            outfile='IRAS16293.cont.ap2.image.pbcor',
            box='150,104,345,299')
 
    # and checking the image
    imview (raster=[{'file': 'IRAS16293.cont.ap2.image.pbcor',
                     'range': [-0.01,1.5],'colorwedge':True}],
            out='IRAS16293.cont.ap2.image.pbcor.png')

# Apply calibration to the original line data
mystep = 14
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Restore the flags made for the continuum imaging
    flagmanager(vis='IRAS16293_Band9.fixed.rebin.ms',mode='restore',versionname='Original')

    # Apply final self-calibration to the line data
    applycal(vis='IRAS16293_Band9.fixed.rebin.ms',
             gaintable=[PROBLEM],calwt=F,flagbackup=F)

# Subtract the continuum from the line data
mystep = 15
if(mystep in thesteps):
    print 'Step ', mystep, step_title[mystep]

    # Subtract the continuum from the line data
    contspw='0:20~1700;2500~3820,1:20~400;500~1100;1900~3000;3450~3800,2:20~1700;2200~3100,3:800~1600;2900~3300'
    uvcontsub(vis='IRAS16293_Band9.fixed.rebin.ms',fitspw=contspw,fitorder=1,combine='spw')