import math, sys, os, numpy
import v190_utils

from math import sin,cos,floor,pi,sqrt
from Wizardry.AIPSData import AIPSUVData

################################################################################
# Duplicate the solutions for the non-closing baselines and copy to actual uv dataset
def wizardry_selfcalcopy(experiment, uvklass, source, isll, calno, basesntabno):
    srcuvdata = AIPSUVData(source, 'SPLIT', 1, 1)
    uvdata = AIPSUVData(experiment.upper(), uvklass, 1, 1)

    oldsn = srcuvdata.table("SN", 1)
    print oldsn.keywords
    #print "number of terms is " + str(oldsn.keywords['NO_TERM'])
    #newsn = uvdata.attach_table('SN', basesntabno+calno+1, no_term=oldsn.keywords['NO_TERM'])
    newsn = uvdata.attach_table('SN', basesntabno+calno+1)
    newsn.keywords['NO_ANT'] = oldsn.keywords['NO_ANT']
    newsn.keywords['NO_POL'] = 2
    newsn.keywords['NO_IF'] = 4

    #Work out the average value for each IF of each telescope
    numvals = len(srcuvdata)
    antable = srcuvdata.table('AN', 1)
    ancount = 0
    pksindex = -1
    atcaindex = -1
    mopraindex = -1
    for ant in antable:
        if ant.anname == 'PKS':
            pksindex = ancount + 1
        if ant.anname[0:3] == 'CAT':
            atcaindex = ancount + 1
        if ant.anname == 'MOPRA':
            mopraindex = ancount + 1
        ancount = ancount + 1
    ntelescopes = ancount
    sums = numpy.zeros((ntelescopes, 4, 2), numpy.float64)
    counts = numpy.zeros((ntelescopes), numpy.float64)
    factors = numpy.zeros((ntelescopes, 4, 2), numpy.float64)
    for vis in srcuvdata:
        #If its a AT-MP-PA baseline
        if (not vis.baseline[0] == vis.baseline[1]) and \
        (vis.baseline[0] == pksindex or vis.baseline[0] == atcaindex or vis.baseline[0] == mopraindex) and \
        (vis.baseline[1] == pksindex or vis.baseline[1] == atcaindex or vis.baseline[1] == mopraindex):
            for i in range(4):
                for j in range(2):
                    if visibility.visibility[i][0][j][0] == visibility.visibility[i][0][j][0]:
                        #Its not a NaN so use it
                        toadd = math.sqrt(visibility.visibility[i][0][j][0]*visibility.visibility[i][0][j][0]
                                          + visibility.visibility[i][0][j][1]*visibility.visibility[i][0][j][1])
                        sums[vis.baseline[0]][i][j] = sums[vis.baseline[0]][i][j] + toadd
                        sums[vis.baseline[1]][i][j] = sums[vis.baseline[0]][i][j] + toadd
            counts[vis.baseline[0]] = counts[vis.baseline[0]] + 1
            counts[vis.baseline[1]] = counts[vis.baseline[1]] + 1

    if isll:
        for i in range(ntelescopes):
            if i==pksindex or i ==atcaindex or i==mopraindex:
                for j in range(4):
                    factors[i][j][0] = sums[i][j][1]/sums[i][j][0]
                    factors[i][j][1] = 1.0
            else:
                for j in range(4):
                    factors[i][j][0] = 1.0
                    factors[i][j][1] = 1.0
    #else:
    #    for i in range(ntelescopes):
    #        for i in range(4):
    #            factors[i][j][0] = 1.0
    #            factors[i][j][1] = 1.0
    #        if i==pksindex or i ==atcaindex or i==mopraindex:
    #            num = (sums[i][0][0] + sums[i][0][1] + sums[i][1][0] + sums[i][1][1])/4.0
    #            factors[i][2][0] = num/sums[i][2][0]
    #            factors[i][2][1] = num/sums[i][2][1]
    #            factors[i][3][0] = num/sums[i][3][0]
    #            factors[i][3][1] = num/sums[i][3][1]
    
    for row in oldsn:
        use = True
        if isll:
            for i in range(4):
                amp = math.sqrt(row.real2[i]*row.real2[i] + row.imag2[i]*row.imag2[i])
                if not(amp > 0.2 and amp < 5.0 and amp == amp): #Sensible values only
                    use = False
                #print "Antenna " + str(row.antenna_no) + "'s amp is " + str(amp) + ", existing real/imag is " + str(row.real1[i]) + " + " + str(row.imag1[i]) + " i, while the real2/imag2 were " + str(row.real2[i]) + " + " + str(row.imag2[i])
                #if amp*factors[row.antenna_no-1][i][0] > 10000:
                #    print "Antenna " + str(row.antenna_no) + "'s amp is " + str(amp) + ", existing real/imag is " + str(row.real1[i]) + " + " + str(row.imag1[i]) + " i, while the real2/imag2 were " + str(row.real2[i]) + " + " + str(row.imag2[i]) + ", factors is " + factors[row.antenna_no-1][i][0]
                else:
                    row.real1[i] = row.real1[i]*factors[row.antenna_no-1][i][0]
                    row.imag1[i] = row.imag1[i]*factors[row.antenna_no-1][i][0]
        else:
            #amp = (math.sqrt(row.real1[0]*row.real1[0] + row.imag1[0]*row.imag1[0]) +
            #       math.sqrt(row.real2[0]*row.real2[0] + row.imag2[0]*row.imag2[0]) +
            #       math.sqrt(row.real1[1]*row.real1[1] + row.imag1[1]*row.imag1[1]) +
            #       math.sqrt(row.real2[1]*row.real2[1] + row.imag2[1]*row.imag2[1]))/4.0
            amp1 = (math.sqrt(row.real1[0]*row.real1[0] + row.imag1[0]*row.imag1[0]) + \
                    math.sqrt(row.real1[1]*row.real1[1] + row.imag1[1]*row.imag1[1]))/2.0
            amp2 = (math.sqrt(row.real2[0]*row.real2[0] + row.imag2[0]*row.imag2[0]) + \
                    math.sqrt(row.real2[1]*row.real2[1] + row.imag2[1]*row.imag2[1]))/2.0
            if not (amp1 > 0.2 and amp1 < 5.0 and amp1 == amp1 and \
                    amp2 > 0.2 and amp2 < 5.0 and amp2 == amp2): #Sensible values only
                use = False
            else:
                row.real1[2] = row.real1[0]
                row.imag1[2] = row.imag1[0]
                row.real1[3] = row.real1[1]
                row.imag1[3] = row.imag1[1]
                row.real2[2] = row.real2[0]
                row.imag2[2] = row.imag2[0]
                row.real2[3] = row.real2[1]
                row.imag2[3] = row.imag2[1]
        if use:
            newsn.append(row)
        else:
            print "Not using row: "
            print row

    newsn.close()
    oldsn.close()