#
#   Usage instructions:
#   ###############################
#   Invocation:
#   parseltongue megapossm.py INNA INCL SOURCE_NAME AIPS_NO docalib 
#                             stokes disk doband lofar
#      the last four are optional and default to -1, LL, 1, -1, no
#   Output: a pdf file called megapossm_plot.pdf.
#
from math import *
from AIPS import AIPS, AIPSDisk
from AIPSTask import AIPSTask, AIPSList
from AIPSData import AIPSUVData, AIPSImage, AIPSCat
from optparse import OptionParser
import re, sys, numpy as np, os

def tconvert (t):
    t_d = int(t)
    t = (t-float(t_d))*24.0
    t_h = int(t)
    t = (t-float(t_h))*60.0
    t_m = int(t)
    t = (t-float(t_m))*60.0
    t_s = int(t)
    return ([t_d,t_h,t_m,t_s])

def reorganise (fname,ant1,ant2):
    ant1=ant1.rstrip('BA')
    ant2=ant2.rstrip('BA')
    dowrite=False
    f = open(fname)
    for i in f:
        if i[1:10]=='Timerange':
            timer=(i.split()[1]+'-'+i.split()[3]).rstrip(')')
    f.close()
    f = open(fname)
    fo = open('temp','w')
    for i in f:
        if not i[0]=='(':
            fo.write('%s'%i)
            continue
        if i[1:5]=='Plot':
            dowrite=True
            fo.write('/Helvetica-Bold findfont  134 scalefont setfont\n')
            fo.write('(%s %s-%s)\n' % (timer,ant1,ant2))
        else:
            if dowrite:
                fo.write('(.)\n')
            else:
                fo.write('%s'%i)
    f.close()
    fo.close()
    os.system('mv temp %s'%fname)

default_ver = '31DEC12'

##############  BEGINNING OF SCRIPT #####################

if len(sys.argv)<3:
    print 'syntax: parseltongue megapossm.py inna incl source aips_no'
    sys.exit()

AIPS.userno = int(sys.argv[4])
print 'Using AIPS user id ',AIPS.userno

#from Wizardry.AIPSData import AIPSUVData
#from Wizardry.AIPSData import AIPSImage

docalib = -1
if len(sys.argv)>5:
    docalib=int(sys.argv[5])
stokes='LL'
if len(sys.argv)>6:
    stokes=sys.argv[6]
disk=1
if len(sys.argv)>7:
    disk=int(sys.argv[7])
doband=-1
if len(sys.argv)>8:
    doband=int(sys.argv[8])

uvdata = AIPSUVData (sys.argv[1], sys.argv[2], disk, 1)
source = sys.argv[3]; source_id=-1

cl = uvdata.table('CL',1)
nx = uvdata.table('NX',1)
su = uvdata.table('SU',1)
an = uvdata.table('AN',1)

if len(sys.argv)>9 and sys.argv[9] == 'lofar':
    inewan = 0
    newan = dict()
    for ia in an:
        if ia['anname'][0] not in ['C','R']:
            newan[inewan] = ia
            inewan+=1
    an=newan

try:
    uvdata.zap_table('PL',-1)
except:
    pass

print 'Locating sources...',
for i in su:
    if source==i.source.rstrip():
        source_id=i.id__no

if source_id==-1:
    print 'Source not found'
    sys.exit()

print 'found at location ',source_id
os.system('rm POL_*.ps')

tsamp = []
if len(cl)<100:
    for i in range(len(cl)):
        if cl[i].source_id==source_id:
            tsamp.append(cl[i].time)
else:
    for i in range(len(nx)):
        if nx[i].source_id==source_id:
            tsamp.append(nx[i].time)

tsamp=np.unique(tsamp)

# get baseline pairs
antennas = np.array([an[0].nosta,an[1].nosta])
annames = np.array([an[0].anname.strip(),an[1].anname.strip()])
for i in range(1,len(an)-1):
    antennas = np.vstack((antennas,np.array([an[i].nosta,an[i+1].nosta])))
    annames = np.vstack((annames,np.array([an[i].anname.rstrip(),an[i+1].anname.rstrip()])))
antennas = np.vstack((antennas,np.array([an[len(an)-1].nosta,an[0].nosta])))
annames = np.vstack((annames,np.array([an[len(an)-1].anname.rstrip(),an[0].anname.rstrip()])))

print antennas
print annames
ntsamp = 7

for i in range(ntsamp):
    tidx = i*len(tsamp)/ntsamp
    tstart = tsamp[tidx]
    timer1 = tconvert (tstart)
    timer2 = tconvert (tstart+5.0/(60.*24.))
    for j in range(len(antennas)):
        print 'Processing time ',timer1,' antenna ',j
        possm=AIPSTask('possm')
        possm.timer[1:]=timer1+timer2
        possm.indata=uvdata
        possm.bif=1
        possm.flagver=0
        possm.sources[1]=source
        possm.antennas[1]=float(antennas[j,0])
        possm.baseline[1]=float(antennas[j,1])
        possm.stokes=stokes
        possm.docalib=docalib
        possm.doband=doband
        possm.solint=5
        possm.aparm[1:]=[0,0,0,0,0,0,0,0,1,0]
        possm.inp()
        possm.go()
        lwpla=AIPSTask('lwpla')
        lwpla.indata=uvdata
        lwpla.plver=0
        lwpla.invers=0
        lwpla.outfile='./POL_%02d_%02d.ps'%(j,i)
        lwpla.go()
        reorganise('./POL_%02d_%02d.ps'%(j,i),annames[j,0],annames[j,1])

os.system('montage POL_*.ps -tile %dx%d -geometry 600x%d -rotate 90 megapossm_plot.png'%(ntsamp,len(antennas),600*len(antennas)/ntsamp))