from AIPS import AIPS, AIPSDisk
from AIPSTask import AIPSTask, AIPSList
from AIPSData import AIPSUVData, AIPSImage
import LocalProxy
from xmlrpclib import ServerProxy
import copy, optparse, os, sys, re, time
#
# AIPS configuration
#
AIPS.proxies = [ LocalProxy ]
AIPSTask.version = '31DEC04'
AIPS.userno = 1522
#AIPS.disks = [ None, AIPSDisk(None, 1), AIPSDisk(None, 2),
#               AIPSDisk(None, 3), AIPSDisk(None, 4) ]

#######################################################################
#   My routines
#
def summary(uvdata,name):
    dtsum = AIPSTask('dtsum')
    dtsum.indata = uvdata
    dtsum.docrt = 0
    dtsum.aparm[1] = 0
    dtsum.outprint = './' + name.lower() + '_sum.txt'
    dtsum()
    listr = AIPSTask('listr')
    listr = AIPSTask('listr')
    listr.indata = uvdata
    listr.optype = 'SCAN'
    listr.outprint = './' + name.lower() +  '_scan.txt'
    listr.docrt = 0
    listr()


def loadmode(imode,uvdata,control):
    dirname = control['scratch']
    expname = control['expname']
    nfits = int(control[imode]['nidi'])
    for i in xrange(1,nfits+1):
        fitld = AIPSTask('fitld') 
        fitld.infile = dirname + expname + '_' + str(imode) + '_1.IDI' + str(i)
        if os.path.exists(fitld.infile) and\
               os.path.getsize(fitld.infile) > 0:
            fitld.outdata = uvdata
            fitld.douvcomp = 1
            fitld.doconcat = 1
            fitld()
        else:
            print 'no such file:',fitld.infile
            raise IOError, "no such file"

def writefits(uvdata,outfits):
    fittp = AIPSTask('fittp')
    fittp.outfile = outfits
    fittp.indata = uvdata
    fittp()

def time_average(uvdata,timavg,dtin,dtout):
    # do a time average and attach the NX again
    uvavg = AIPSTask('uvavg')
    uvavg.indata = uvdata
    uvavg.outdata = timavg
    uvavg.yinc = dtout
    uvavg.zinc = dtin
    uvavg()
    indxr = AIPSTask('indxr')
    indxr.indata = timavg
    indxr.cparm[3] = 1
    indxr()

def spectral_average(uvdata,spavg,navg):
    avspc = AIPSTask('avspc')
    avspc.indata = uvdata
    avspc.outdata = spavg
    avspc.avoption = 'SUBS'
    avspc.channel = navg
    avspc()

def external_cal(uvdata,refant):
    if uvdata.table_highver('AIPS TY')==0:
        antab = AIPSTask('antab')
        antab.indata = uvdata
        antab.infile = 'EL032.ANTAB'
        antab()
    if uvdata.table_highver('AIPS FG')==0:
        uvflg = AIPSTask('uvflg')
        uvflg.indata = uvdata
        uvflg.infile = 'EL032.UVFLG'
        uvflg()
    if uvdata.table_highver('AIPS SN')==0:        
        apcal = AIPSTask('apcal')
        apcal.indata = uvdata
        apcal.tyver = 1
        apcal.gcver = 1
        apcal.snver = 1
        apcal.prtlev = 1
        apcal()
    if uvdata.table_highver('AIPS CL')==1:                
        clcal = AIPSTask('clcal')
        clcal.indata = uvdata
        clcal.opcode = 'CALI'
        clcal.interpol = 'SELF'
        clcal.doblank = 1
        clcal.refant = refant
        clcal.snver = 1
        clcal.gainver = 1
        clcal.gainuse = 2
        clcal()

def direct_fringe(uvdata,refant,cals):
    if uvdata.table_highver('AIPS SN')==1:
        fring = AIPSTask('fring')
        fring.indata = uvdata
        fring.gainuse = 2
        fring.flagver = 1
        fring.docalib = 1
        fring.refant = refant
        fring.solint = 4
        fring.aparm[1] = 4
        fring.aparm[6] = 2
        fring.aparm[7] = 6
        fring.aparm[9] = 0
        fring.dparm[1] = 1
        fring.dparm[4] = 3.99
        fring.calsour = AIPSList(cals)
        fring()
    if uvdata.table_highver('AIPS CL')==2:
        clcal = AIPSTask('clcal')
        clcal.indata = uvdata
        clcal.opcode = 'CALI'
        clcal.interpol = 'AMBG'
        clcal.doblank = 1
        clcal.refant = refant
        clcal.snver = 2
        clcal.gainver = 2
        clcal.gainuse = 3
        clcal()

def update_cal(indata,outdata,ext,ver):
    #copies 1-1 extension files
    extkey = 'AIPS ' + ext
    if outdata.table_highver(extkey)==ver-1:
        tacop = AIPSTask('tacop')
        tacop.indata = indata
        tacop.outdata = outdata
        tacop.inext = ext
        tacop.invers = ver
        tacop.outver = ver
        tacop.ncount = 1
        tacop()
    else:
        print 'SKIP tacop',outdata.name,ext,ver-1

def cross_bpass(uvdata,refant,cals):
    if uvdata.table_highver('AIPS BP')==0:
        bpass = AIPSTask('bpass')
        bpass.indata = uvdata
        bpass.gainuse = 3
        bpass.flagver = 1
        bpass.docalib = 1
        bpass.refant = refant
        bpass.solint = -1
        bpass.bpassprm = AIPSList([ 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0])
        bpass.calsour = AIPSList(cals)
        bpass()

def hc(data, type, istart, istop):
    #make a hardcopy, remove existing file first 
    lwpla = AIPSTask('lwpla')
    outname = './' + rootname.lower() + '_' + type + '.ps'
    if os.path.isfile(outname):
        backup = outname + '.old'
        if os.path.isfile(backup):
            os.remove(backup)
        print 'renaming outname backup'
        os.rename(outname,backup)
    lwpla.indata = data
    lwpla.outfile = outname
    lwpla.dparm[6] = 4
    lwpla.dparm[7] = 11
    lwpla.dparm[8] = 11
    lwpla.plver = istart
    lwpla.inver = istop
    lwpla()
    #data.zap_table('AIPS PL', -1)

def calplot(uvdata,inext, invers, optype):
    istart = uvdata.table_highver('AIPS PL')+1
    snplt = AIPSTask('snplt')
    snplt.indata = uvdata
    snplt.inext = inext
    snplt.invers = invers
    snplt.optype = optype
    snplt.nplots = 8
    snplt()
    istop = timavg.table_highver('AIPS PL')
    type = inext.lower()+str(invers)+'_'+optype.lower()
    hc(uvdata,type,istart,istop)

def specplot(uvdata,type,sources=None,refant=1):
    istart = uvdata.table_highver('AIPS PL')+1
    possm = AIPSTask('possm')
    possm.indata = uvdata
    possm.nplots = 6
    possm.docal = 1
    possm.doband = 1
    possm.bchan = 5
    if type == 'bandp':
        possm.aparm = AIPSList([ 0, 0, 0, 0, 0, 0, 1, 2])
    elif type == 'xspec':
        possm.aparm = AIPSList([ 1, 0, 0, 0, 0, 0, 1])
        possm.antenna[1] = refant
    elif type == 'auto':
        possm.aparm = AIPSList([ 0, 0, 0, 0, 0, 0, 1, 1])
        possm.antenna[1] = refant
    else:
        raise RuntimeError, "Unknown option type",type
    if type == 'bandp':
        possm()
    else:
        for isu in sources:
            possm.sources[1] = isu
            possm()
        
    istop = timavg.table_highver('AIPS PL')
    prtype = type.lower()
    hc(uvdata,prtype,istart,istop)

def timplot(uvdata,type,refant=1):
    istart = timavg.table_highver('AIPS PL')+1
    vplot = AIPSTask('vplot')
    vplot.indata = uvdata
    vplot.nplots = 8
    #need to know how many channels there are
    header = uvdata.header()
    iax = header.get('ctype').index('FREQ    ')
    nchan = header.get('inaxes')[iax]
    vplot.bparm[2] = -1
    vplot.bchan = int(0.02*nchan)
    vplot.echan = int(0.98*nchan)
    vplot.antenna[1] = refant
    if type == 'raw':
        vplot.docal = 0
        vplot.doband = 0
    elif type == 'cal':
        vplot.docal = 1
        vplot.doband = 1
    else:
        raise RuntimeError, "Unknown option type",type
    vplot.stokes = 'RR'
    vplot()
    vplot.stokes = 'LL'
    vplot()
    istop = timavg.table_highver('AIPS PL')
    prtype = 'tim'+type.lower()
    hc(uvdata,prtype,istart,istop)

def uvcal_plots(uvdata,root):
    #print a lot of plots
    calplot(uvdata,'TY',1,'TSYS')
    calplot(uvdata,'CL',3,'AMP')
    calplot(uvdata,'CL',3,'DELA')
    calplot(uvdata,'CL',3,'RATE')
    calplot(uvdata,'CL',3,'PHAS')
    timplot(uvdata,'raw',refant)
    timplot(uvdata,'cal',refant)
    specplot(uvdata,'bandp',refant)
    specplot(uvdata,'xspec',cals_fringe,refant)

def parse_com(name):

    repair = re.compile(r'^\s*(\w+)\s*=\s*(\S+)\s*$')
    relist = re.compile(r'^\s*(\w+)\[\]\s*=\s*(\S+(?:\s*\,\s*\S+)*)\s*$')
    resplit = re.compile(r'\s*\,\s*')
    remode = re.compile(r'^\s*mode\s*=\s*(\d+)\s*$')
    recom = re.compile(r'^\#')

    control = {}
    imode = 0
    general = True
    file = open(name,'r')
    nmode = 0
    for line in file:
        yescom = recom.match(line)
        yespair = repair.match(line)
        yeslist = relist.match(line)
        yesmode = remode.match(line)
        if yescom:
            continue
        elif yespair:
            input = yespair.groups()
            if yesmode:
                if general:
                    general = False
                nmode += 1
                imode = int(input[1])
                if imode != nmode:
                    raise ValueError, "Not a simple mode list"
                control[nmode] = {}

            if general:
                control[input[0]]=input[1]
            else:
                control[nmode][input[0]]=input[1]            
        elif yeslist:
            input = yeslist.groups()
            yessplit = resplit.split(input[1])
            if general:
                control[input[0]]=yessplit
            else:
                control[nmode][input[0]]=yessplit
        else:
            print 'crap, no match:', line
    control['nmode'] = nmode
    return control

def printcontrol(control):

    for key in control.iterkeys():
        if not type(control[key]) == type({}):
            print key,'=',control[key]
    for imode in range(1,control['nmode']+1):
            print imode,':'
            for key2 in control[imode].iterkeys():
                print imode,':',key2,':',control[imode][key2]
    return

def getoutarchive(control,imode):

    exp = control['expname']
    date = control['obsdate']
    inurl = control['inurl']
    outdi = control['scratch']
    passwd = control['passwd']

    for idi in xrange(1,int(control[imode]['nidi'])+1):

        file = exp.lower() + '_'\
               +str(imode)\
               +'_1.IDI'+str(idi)
        outfile = outdi+file

        command = 'wget -nv -O '+outfile+' '\
                  +'--http-user='+exp+' '\
                  +'--http-passwd='+passwd+' '\
                  +inurl+exp.upper()+'_'+date+'/fits/'+file
        #print command

        if not os.path.exists(outfile):
            os.system(command)
        elif os.path.getsize(outfile) < 1:
            os.system(command)
        else:
            print outfile, 'exists!'
            
    return

def cleanfits(control,imode):

    exp = control['expname']
    date = control['obsdate']
    outdi = control['scratch']

    for idi in xrange(1,int(control[imode]['nidi'])+1):

        file = exp.lower() + '_'\
               +str(imode)\
               +'_1.IDI'+str(idi)
        outfile = outdi+file

        if not os.path.exists(outfile):
            raise IOError, "file not found on delete"
        else:
            os.remove(outfile)
            
    return

def makelogname():

    script = os.path.splitext(sys.argv[0])[0]
    year = time.localtime()[0]-2000
    month = time.localtime()[1]
    day = time.localtime()[2]
    hr = time.localtime()[3]
    min = time.localtime()[4]
    ext = 'log'
    file = '%s.%02i%02i%02i_%02i%02i.log'\
           % ( script, year, month, day, hr, min )

    return file

#########################
# Load and sort the data
#
# AIPSTask.isbatch = 0


logfile = makelogname()

AIPS.log = open(logfile, 'w')
AIPSTask.msgkill = -5
control = parse_com('el032com.txt')
printcontrol(control)

refant = int(control['refant'])
tsmooth = float(control['tsmooth'])
tcorr = float(control['tcorr'])

for imode in range(1,control['nmode']+1):
#for imode in range(1,2):

    rootname = control[imode]['source']

    print 'STARTING mode %i on %s at %s' % (imode,rootname,time.asctime())

    cals_bpass = control[imode]['calibrators']
    cals_ref = control[imode]['refs']
    cals_fringe = cals_bpass+cals_ref

    outfits = control['scratch']+control['expname'].upper() + '_'+\
              rootname.upper()+'.UVF'
    if os.path.exists(outfits):
        print 'SKIP ENTIRE mode', imode
        continue
    
    getoutarchive(control,imode)

    fulldata = AIPSUVData( rootname, 'UVDATA', 1, 1)
    timavg = AIPSUVData( rootname, 'TIMAVG', 1, 1)
    specavg = AIPSUVData( rootname, 'SPCAVG', 1, 1)

    if not fulldata.exists():
        loadmode(imode,fulldata,control)
    else:
        print 'SKIP loading'
        
    if not timavg.exists():
        time_average(fulldata,timavg,tcorr,tsmooth)
        summary(timavg,rootname)
    else:
        print 'SKIP time averaging'

    if timavg.table_highver('AIPS CL')==1:                
        external_cal(timavg,refant)
    else:
        print 'SKIP external cal'

    if not specavg.exists():
        spectral_average(timavg,specavg,32)
    else:
        print 'SKIP spectral average'
                    
    if specavg.table_highver('AIPS CL')==2:          
        direct_fringe(specavg,refant,cals_fringe)
        update_cal(specavg,timavg,'SN',2)
        update_cal(specavg,timavg,'CL',3)
    else:
        print 'SKIP fringe fit'
    
    if timavg.table_highver('AIPS BP')==0:
        cross_bpass(timavg,refant,cals_bpass)
    else:
        print 'SKIP bandpass'

    if timavg.table_highver('AIPS PL')==0:
        uvcal_plots(timavg,rootname)
    else:
        print 'SKIP plots of uv calibration'
              
    if not os.path.exists(outfits):
        fulldata.zap()
        cleanfits(control,imode)
        writefits(timavg,outfits)
    else:
        print 'SKIP output fits'