#!/usr/bin/env python
import os

datadir = '.'
asdmlist = ['uid___A002_X85c183_X36f',
            'uid___A002_X85c183_X60b',
            'uid___A002_X8602fa_X2ab',
            'uid___A002_X8602fa_X577',
            'uid___A002_X864236_X2d4',
            'uid___A002_X864236_X693',
            'uid___A002_X86fcfa_X664',
            'uid___A002_X86fcfa_X96c',
            'uid___A002_X86fcfa_Xd9']
factors = [[41.37, 42.39, 43.45, 42.82],
           [41.37, 42.39, 43.45, 42.82],
           [40.99, 42.74, 40.08, 42.09],
           [40.99, 42.74, 40.08, 42.09],
           [39.42, 42.08, 40.18, 40.82],
           [39.42, 42.08, 40.18, 40.82],
           [43.49, 43.70, 42.01, 43.04],
           [43.49, 43.70, 42.01, 43.04],
           [43.49, 43.70, 42.01, 43.04]]

# 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)

for (asdm,jyperk) in zip(asdmlist,factors):
    os.system('rm -rf %s*'%(asdm))
    
    vis = asdm+'.ms'
    listfile = vis + '.listobs'
    calvis = vis + '.cal'
    splitvis = calvis + '.split'
    blvis = vis + '.bl'
    nlccaltable = vis + '.nlc'
    jypkcaltable = vis + '.jy'
    jypkvis = vis + '.jypk'
    
    # importasdm
    importasdm(os.path.join(datadir, asdm), vis=vis, asis='Antenna Station Receiver Source CalAtmosphere CalWVR', bdfflags=False, overwrite=True)

    # 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
    tsdcal(infile=vis, calmode='ps,tsys,apply', spwmap={9:[17], 11:[19], 13:[21],15:[23]}, spw='9,11,13,15,17,19,21,23')
    
    # 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=map(to_amp_factor, [1.25]))
    applycal(vis=splitvis, gaintable=nlccaltable)

    # baseline subtraction
    tsdbaseline(infile=splitvis, datacolumn='corrected', spw='', blfunc='poly', order=1, outfile=blvis, overwrite=True)

    os.system('rm -rf %s'%(splitvis))
    
    # Jy/K factor application
    if os.path.exists(jypkcaltable):
        os.system('rm -rf %s'%(jypkcaltable))

    gencal(vis=blvis, caltable=jypkcaltable, caltype='amp', spw='0,1,2,3', parameter=map(to_amp_factor, jyperk))
    applycal(vis=blvis, gaintable=jypkcaltable)

#Flag out PM02 in two datasets because of instability in bandpasses.
flagdata(vis='uid___A002_X8602fa_X2ab.ms.bl', antenna='PM02')
flagdata(vis='uid___A002_X8602fa_X577.ms.bl',antenna='PM02')