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

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]
        antlist = msmd.antennanames()
        msmd.close()
        print "SPW %d: %.3f GHz" % (spw, freq*1e-9)
        # skip imaging of bad antenna instead of flagging
        if vis.startswith('uid___A002_X8602fa_Xc3'):
            antlist.pop(antlist.index('PM02'))

        theorybeam = aU.primaryBeamArcsec(frequency=freq*1e-9,
                                          fwhmfactor=fwhmfactor,
                                          diameter=diameter)
 
        cell = theorybeam/9.0
        imsize = int(round(maxsize/cell)*2)

        for ant in antlist:
            imagename = '%s.%s.spw%s.image'%(vis, ant, 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)

    # Determine the Jy/K values
    fout = open(vis+'.JyPerK.txt', 'w')
    try:
        for spw in spwlist:
            srcflux = aU.getALMAFluxForMS(vis,
                                          field='3c279',
                                          spw=str(spw))['3c279']['fluxDensity']
            jyperklist = []
            for ant in antlist:
                imagename = '%s.%s.spw%s.image'%(vis, ant, spw)
                peak = imstat(imagename)['max'][0]
                fwhm = aU.getfwhm2(imagename)
                print 'Apparent FWHM (inc. gridding convolution) is %.2f arcsec' % fwhm
                # non-ephemeris source
                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, numpy.mean(jyperklist))
            fout.write('%d\t%s\t%.2f\n' % (spw, 'mean', numpy.mean(jyperklist)))
    finally:
        fout.close()