import math, sys, os, numpy
import v190_utils

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

################################################################################
# Use splat twice (doweight on, then off) to fix the weights of a single-source file
def fixweights(rootdir, experiment, source, uvklass, seqno,
               v190_num, otherklass):
    #Hardcoded table of applied tsys values
    if experiment == 'v190a':
        tsys = [43.0, 86.0, 430.0, 560.0, 600.0]
        starttime = 0.915
        endtime = 1.42
    elif experiment == 'v190b':
        tsys = [40.0, 68.0, 340.0, 470.0]
        starttime = 0.375
        endtime = 1.0
    elif experiment == 'v190c':
        tsys = [42.0, 68.0, 340.0, 420.0, 23.0]
        starttime = 0.0
        endtime = 1.0
    elif experiment == 'v190d':
        tsys = [42.0, 68.0, 340.0, 420.0, 23.0]
        starttime = 0.45
        endtime = 1.46
    elif experiment == 'v190e':
        tsys = [86.0, 430.0, 43.0, 560.0, 600.0, 25.0]
        starttime = 0.375
        endtime = 0.875
    elif experiment == 'v190f':
        tsys = [42.0, 68.0, 340.0, 420.0, 23.0]
        starttime = 0.04
        endtime = 1.09
    elif experiment == 'v190g':
        tsys = [86.0, 430.0, 43.0, 560.0, 600.0, 25.0]
        starttime = 0.04
        endtime = 0.6
    elif experiment == 'v190h':
        tsys = [42.0, 68.0, 340.0, 420.0, 23.0]
        starttime = 0.66
        endtime = 1.67
    elif experiment == 'v190i':
        tsys = [42.0, 68.0, 340.0, 420.0, 23.0]
        starttime = 0.958
        endtime = 1.96
    elif experiment == 'v190j':
        tsys = [42.0, 68.0, 340.0, 420.0]
        starttime = 0.08
        endtime = 1.125
    elif experiment == 'v190k':
        tsys = [86.0, 430.0, 43.0, 560.0, 600.0, 113.0]
        starttime = 0.375
        endtime = 0.875
    elif experiment == 'v190l':
        tsys = [42.0, 68.0, 340.0, 420.0, 23.0]
        starttime = 0.04
        endtime = 1.085
    elif experiment == 'TECJNK':
        tsys = [42.0, 68.0, 340.0, 420.0, 23.0]
        starttime = 0.0
        endtime = 2.0
    else:
        print "Couldn't locate experiment (got " + experiment + \
              ") - aborting!!!"
        sys.exit(1)

    aips_version_std = '31DEC05'
    uvdata = normUVData(experiment.upper(), uvklass, 1, seqno)
    if v190_num > 0:
        normpsrdata = normUVData(source, otherklass, 1, v190_num)
        psrdata = AIPSUVData(source, otherklass, 1, v190_num)
        normsplatdata = normUVData('TSYSJUNK', 'TSYS', 1, 1)
        if normsplatdata.exists():
            normsplatdata.zap()
    else:
        print "Can't load from disk - must already be in catalogue! Aborting"
        sys.exit()

    #Get an example SN table line, attach a new SN table
    wizuvdata = AIPSUVData(experiment.upper(), uvklass, 1, seqno)
    exampleSNtable = wizuvdata.table('SN',1)
    samplerow = exampleSNtable[0]
    print "Keywords are:"
    print exampleSNtable.keywords
    print "samplerow is:"
    print samplerow
    sntable = psrdata.attach_table('SN', 1)
    sntable.keywords['NO_ANT'] = len(psrdata.antennas)
    sntable.keywords['NO_IF'] = 4
    sntable.keywords['NO_POL'] = 2
    num_freqs = 4
    num_stokes = 2
    gainlength = len(samplerow.real1)
    newdelay1 = []
    newdelay2 = []
    newrate1 = []
    newrate2 = []
    newimag1 = []
    newimag2 = []
    for i in range(gainlength):
        newdelay1.append(0.0)
        newdelay2.append(0.0)
        newrate1.append(0.0)
        newrate2.append(0.0)
        newimag1.append(0.0)
        newimag2.append(0.0)
    samplerow.delay_1 = newdelay1
    samplerow.delay_2 = newdelay2
    samplerow.rate_1 = newrate1
    samplerow.rate_2 = newrate2
    samplerow.imag1 = newimag1
    samplerow.imag2 = newimag2

    #set all the stuff in sample row to zero, except the time, antenna, IF
    #and real/imag
    samplerow.i_far_rot = 0
    samplerow.time_interval = 0.001041667 #One entry every 90 secs
    samplerow.mbdelay2 = 0
    samplerow.source_id = 1
    samplerow.mbdelay1 = 0
    samplerow.node_no = 0
    samplerow.source_id = 1
    samplerow.freq_id = 1

    #For a ridiculous 3 day time window, make the SN table
    time = starttime
    while time < endtime:
        samplerow.time = time
        for i in range(len(tsys)):
            samplerow.antenna_no = i+1
            newreal1 = []
            newreal2 = []
            newweight1 = []
            newweight2 = []
            for k in range(gainlength):
                newreal1.append(math.sqrt(tsys[i]/1.0))
                newreal2.append(math.sqrt(tsys[i]/1.0))
                newweight1.append(1.0/tsys[i])
                newweight2.append(1.0/tsys[i])
            samplerow.real1 = newreal1
            samplerow.real2 = newreal2
            samplerow.weight_1 = newweight1
            samplerow.weight_2 = newweight2
            sntable.append(samplerow)
        time = time + samplerow.time_interval
    sntable.close()
    
    #Use splat to apply the SN table, with weights
    splat = AIPSTask('splat', version = aips_version_std)
    splat.indata = normpsrdata
    splat.outdata = normsplatdata
    splat.sources[1:] = ''
    splat.qual = -1
    splat.calcode = ''
    splat.timer[1:] = [0]
    splat.stokes = ''
    splat.selband = -1
    splat.selfreq = -1
    splat.freqid = -1
    splat.bif = 0
    splat.eif = 0
    splat.bchan = 0
    splat.echan = 0
    splat.subarray = 0
    splat.docalib = 2
    splat.gainuse = 1
    splat.dopol = -1
    splat.blver = -1
    splat.flagver = 0
    splat.doband = -1
    splat.bpver = 0
    splat.smooth[1:] = [0]
    splat.douvcomp = -1
    splat.aparm[1] = 1
    splat.aparm[2:] = [0]
    splat.aparm[4] = 1
    splat.ichansel[1:][1:] = [0]
    splat.channel = 0
    splat.solint = 0
    splat.baddisk[1:] = [0]
    splat()

    #print "Exiting so data can be inspected..."
    #sys.exit()

    #Delete the original psr data, create the pointer to the Wiz-style splatdata
    normpsrdata.zap()
    splatdata = AIPSUVData('TSYSJUNK', 'TSYS', 1, 1)

    #Create an SN table to undo the effect on amplitudes
    sntable2 = splatdata.attach_table('SN', 1)
    sntable2.keywords['NO_ANT'] = len(psrdata.antennas)
    sntable2.keywords['NO_IF'] = 4
    sntable2.keywords['NO_POL'] = 2
    time = starttime
    #newweight1 = []
    #newweight2 = []
    #for i in range(gainlength):
    #    newweight1.append(0.0)
    #    newweight2.append(0.0)
    #samplerow.weight_1 = newweight1
    #samplerow.weight_2 = newweight2
    while time < endtime:
        samplerow.time = time
        for i in range(len(tsys)):
            samplerow.antenna_no = i+1
            newreal1 = []
            newreal2 = []
            newweight1 = []
            newweight2 = []
            for k in range(gainlength):
                newreal1.append(math.sqrt(1.0/tsys[i]))
                newreal2.append(math.sqrt(1.0/tsys[i]))
                newweight1.append(tsys[i]/1.0)
                newweight2.append(tsys[i]/1.0)
            samplerow.real1 = newreal1
            samplerow.real2 = newreal2
            samplerow.weight_1 = newweight1
            samplerow.weight_2 = newweight2
            sntable2.append(samplerow)
        time = time + samplerow.time_interval
    sntable2.close()

    #Use splat to apply, this time without doing weights
    splat.indata = normsplatdata
    splat.outdata = normpsrdata
    splat.docalib = 1
    splat.aparm[4] = 0
    splat()

    #Zap the splatdata catalogue entry
    normsplatdata.zap()