#!/usr/bin/python
# The following assumes standard calibration of the amplitude calibrator and
# science target via the script generator.
# Examples of each are provided in:
# Amplitude Calibrator: uid___A002_X86fcfa_X3ae.ms.scriptForSDCalibration.py
# Science target: uid___A002_X86fcfa_X664.ms.scriptForSDCalibration.py
#
# Below is the procedure to achieve a fully calibrated TP Science image,
# in Jansky/beam units, taking as input a set of calibrated Science and
# Amplitude Calibrator datasets (in units of Kelvin).
#
# 1. Image the amplitude calibrator and measure the value of Jy/K from the image
# (values should be obtained for each night data is taken and for each spw and antenna)
#
# 2. Convert the science target units from Kelvin to Jansky
# (treating each night's data and each spw separately)
#
# 3. Image the science target
# (combining data from all antennas)
#
# 4. Optionally subtract a residual background from the image
# (currently, this is likely necessary in most cases; in future, improvements
# in the sdbaseline task should remove the need for this step)
#
# 5. Apply the restoring beam to the science image

### scriptForImagingAmpCalAndDerivingJyPerK.py deals with STEP 1. ###
### Tested in CASA Version 4.3.0                                  ###

### 1. Image the amplitude calibrator and measure the value of Jy/K from the image ###
# (values should be obtained for each night data is taken and for each spw and antenna)
# One ampcal per night is assumed
#
# In this case, data taken on:
# 2014-07-01  uid___A002_X85c183_X895
# 2014-07-05  uid___A002_X8602fa_Xc3
# 2014-07-07  uid___A002_X864236_Xe1
# 2014-07-17  uid___A002_X86fcfa_X3ae


#----------------------------------------------------------------------------------
# First set variables for the name of a calibrated dataset and a couple of
# parameters to predict the intrinsic beam size. The predicted beam size is used to
# determine the appropriate grid spacing and gridding convolution function to image
# the source.
#----------------------------------------------------------------------------------

# Determine some parameters of the data for each dataset
# (comment in as relevant)
ampcaldata = ['uid___A002_X85c183_X895.ms.cal',
              'uid___A002_X8602fa_Xc3.ms.cal',
              'uid___A002_X864236_Xe1.ms.cal',
              'uid___A002_X86fcfa_X3ae.ms.cal']

spwlist = [17, 19, 21, 23]  # the science spws

# Define some properties of the beam
fwhmfactor = 1.13
diameter = 12

# Determine the Jy/K values
for msname in ampcaldata:

   ### Open a file to which the derived Jy/K values will be written. 
    # Open a file to write the derived Jy/K values
    fout = open(msname+'.JyPerK.txt', 'w')

#----------------------------------------------------------------------------------
# Obtain the spatial sampling of the data (spacings along and perpendicular to the
# scan direction, and the largest dimension in arcsec) using the function
# getTPSampling of the Analysis Utilities.
#----------------------------------------------------------------------------------
    # Determine from the pointing table the sampling used to observe the data
    xSampling, ySampling, maxsize = aU.getTPSampling(msname, showplot=False)

    ### Obtain the list of antennas
    msmd.open(msname)
    antlist = msmd.antennanames()
    msmd.close()

#----------------------------------------------------------------------------------
# Specify a SPW for which the Jy/K value is derived. The center frequency of the
# SPW is obtained, and then the expected beam size is calculated using the function
# primaryBeamArcsec of the Analysis Utilities.
#----------------------------------------------------------------------------------
    for spw in spwlist:
        # Determine the frequency of each science spw
        msmd.open(msname)
        freq = msmd.meanfreq(spw)
        msmd.close()
        print "SPW %d: %.3f GHz" % (spw, freq*1e-9)

        # Determine cell size (using a Gaussian approximation)
        theorybeam = aU.primaryBeamArcsec(frequency=freq*1e-9,
                                          fwhmfactor=fwhmfactor,
                                          diameter=diameter)

#----------------------------------------------------------------------------------
# The cell spacing and size (i.e., pixel size and number of pixels) of the
# resulting map are determined from the estimated beam size and the dimension of
# the observed raster map pattern. The cell spacing is set to be 1/9 of the beam
# size, according to the NAASC Memo 114 (draft).
        # Set up the imaging parameters
#----------------------------------------------------------------------------------
        cell = theorybeam/9.0
        imsize = int(round(maxsize/cell)*2)

       ### Obtain the ID and name of the field. 
        # Get the field ID/name
        msmd.open(msname)
        fieldid = msmd.fieldsforspw(spw, False)[0]
        fieldname = msmd.fieldsforspw(spw, True)[0]
        msmd.close()
        print "Field ID %d, name %s" % (fieldid, fieldname)
        
        # Set the center of the map and determine planet size if its a planet

        # Check if the object is a planet -- NB some of them are not good
        # calibrators (only Uranus is used for Cycle 1/2)
        if fieldname.lower() in ['mercury', 'venus', 'mars', 'jupiter',
            'saturn', 'uranus', 'neptune']:
            ephem = fieldname
            phasecenter = ''
            doPlanet = True
        else:
            ephem = ''
            phasecenter = fieldid
            doPlanet = False

#----------------------------------------------------------------------------------
# Now the image of the source is made for each antenna using the task sdimaging.
# The gridding convolution function used is 'SF' (prolate spheroidal wave function)
# with the support of 6 cell spacings, also according to the NAASC Memo 114 (draft).            
# Image the calibrator for the spectral line spw(s)
#----------------------------------------------------------------------------------
        print 'Imaging the amplitude calibrator...'

        for ant in antlist:
            sdimaging(infiles=msname,
                      field=str(fieldname),
                      spw='%d' % spw,  #sciencespw
                      antenna=ant,
                      nchan=1,
                      mode='channel',
                      width='4080',
                      gridfunction='SF',
                      convsupport= 6,
                      phasecenter=phasecenter,
                      ephemsrcname=ephem,
                      imsize=imsize,
                      cell=str(cell)+'arcsec',
                      overwrite=True,
                      outfile='%s.%s.spw%d.image' % (msname, ant, spw))
            
#----------------------------------------------------------------------------------
# The continuum flux of the source is obtained from the ALMA source catalog using
# the function getALMAFluxForMS of the Analysis Utilities.
# Obtain the Jy/K conversion factor.
#----------------------------------------------------------------------------------
        if doPlanet:
            planetInfo = aU.planetFlux(vis=msname, spw=spw)
            srcflux = planetInfo['fluxDensity']
            srcsize = (planetInfo['majorAxis']*planetInfo['minorAxis'])**0.5
        else:
            srcflux = aU.getALMAFluxForMS(msname,
                                          field=fieldname,
                                          spw=str(spw))[fieldname]['fluxDensity']

#----------------------------------------------------------------------------------
# Finally derive the Jy/K values by comparing the flux of the source (in Jy) with
# the observed brightness temperature (in K). In the case of this guide the amplitude
# calibrator is a point source (quasar 3C279), and hence the Jy/K value is simply
# estimated by the ratio between the source flux obtained above and the peak brightness
# temperature determined using the task imstat. Note that if the amplitude calibrator
# is a planet, the correction for the source size (multiplying the ratio between the
# areas of planet-deconvolved beam and apparent beam) is necessary. This process is
# not written in the code presented here, but implemented in the script
# "ScriptForImagingAmpCalAndDerivingJyPerK.py".
#---------------------------------------------------------------------------------- 
        jyperklist = []
        for ant in antlist:
            peak = imstat('%s.%s.spw%d.image' % (msname, ant, spw))['max'][0]
            fwhm = aU.getfwhm2('%s.%s.spw%d.image' % (msname, ant, spw))
            print 'Apparent FWHM (inc. gridding convolution) is %.2f arcsec' % fwhm
            if doPlanet:
                deconvfwhm = aU.deconvolveDiskFromBeam(fwhm, srcsize)
                print 'FWHM deconvolved for planet size is %.2f arcsec' % deconvfwhm
                # correction factor for dilution due to planet size
                srcsizecorr = (deconvfwhm/fwhm)**2
            else:
                srcsizecorr = 1.
            jyperk = srcflux/peak*srcsizecorr
            print 'SPW%d %s Jy/K = %.2f' % (spw, ant, jyperk)
            fout.write('%d\t%s\t%.2f\n' % (spw, ant, jyperk))
            jyperklist.append(jyperk)
        print '### SPW%d mean Jy/K = %.2f ###' % (spw, pl.mean(jyperklist))
        fout.write('%d\t%s\t%.2f\n' % (spw, 'mean', pl.mean(jyperklist)))

   ### Finally close the file in which the Jy/K values were written. 
    fout.close()
