#!/usr/bin/env python
import os
import sys

import analysisUtils as aU

__rethrow_casa_exceptions=True

datadir = '.'
asdmlist = ['uid___A002_X85c183_X895',
            'uid___A002_X8602fa_Xc3',
            'uid___A002_X864236_Xe1',
            'uid___A002_X86fcfa_X3ae']

# Application of the caltable created by gencal with
# caltype 'amp' appears to multiply x^-2 where x is
# the factor in the caltable. So, desired multiplication
# factor must be converted to its inverse square root
# before creating caltable.
to_amp_factor = lambda x: 1. / sqrt(x)

fwhmfactor = 1.13
diameter = 12

for asdm in asdmlist:
    os.system('rm -rf %s.*'%(asdm))
    
    vis = asdm+'.ms'
    listfile = vis + '.listobs'
    calvis = vis + '.cal'
    splitvis = calvis + '.split'
    nlccaltable = vis + '.nlc'
    
    # importasdm
    importasdm(os.path.join(datadir, asdm), vis=vis, asis='Antenna Station Receiver Source CalAtmosphere CalWVR', bdfflags=False, overwrite=True)
    os.system(os.environ['CASAPATH'].split()[0] + '/bin/bdflags2MS -f "COR DELA INT MIS SIG SYN TFB WVR ZER" ' + os.path.join(datadir, asdm) +' '+vis)

    # listobs
    listobs(vis=vis, listfile=listfile)

    # edge channel flagging
    flagdata(vis=vis, mode='manual', spw='17:0~119;3960~4079,19:0~119;3960~4079,21:0~119;3960~4079,23:0~119;3960~4079', action='apply', flagbackup=True)

    # calibration
    sdcal(infile=vis, calmode='otfraster,tsys,apply', spwmap={9:[17], 11:[19], 13:[21],15:[23]}, spw='17,19,21,23', scan='5,6,7,8')
    
    # non-linearity correction
    split(vis=vis, outputvis=splitvis, datacolumn='corrected', spw='17,19,21,23')
    
    if os.path.exists(nlccaltable):
        os.system('rm -rf %s'%(nlccaltable))

    gencal(vis=splitvis, caltable=nlccaltable, caltype='amp', spw='', parameter=[to_amp_factor(1.25)])
    applycal(vis=splitvis, gaintable=nlccaltable)

    # imaging
    refvis = asdmlist[0] + '.ms'
    xSampling, ySampling, maxsize = aU.getTPSampling(refvis, showplot=False)
    spwlist = [0, 1, 2, 3]
    for spw in spwlist:
        msmd.open(splitvis)
        freq = msmd.meanfreq(spw)
        nchan = msmd.nchan(3)
        fieldid = msmd.fieldsforname('3c279')[0]
        msmd.close()
        print "SPW %d: %.3f GHz" % (spw, freq*1e-9)
    
        theorybeam = aU.primaryBeamArcsec(frequency=freq*1e-9,
                                          fwhmfactor=fwhmfactor,
                                          diameter=diameter)
 
        cell = theorybeam/9.0
        imsize = int(round(maxsize/cell)*2)

        imagename = '%s.3c279.spw%s.image'%(splitvis,spw)
        sdimaging(infiles=splitvis,
                  field='3c279',
                  spw='%s'%(spw),
                  nchan=1,
                  mode='channel',
                  start=0,
                  width=nchan,
                  outframe='lsrk',
                  restfreq='115.271204GHz',
                  gridfunction='SF',
                  convsupport=6,
                  stokes='I',
                  phasecenter=fieldid,
                  ephemsrcname='',
                  imsize=imsize,
                  cell='%sarcsec'%(cell),
                  overwrite=True,
                  outfile=imagename)