import math, sys, os, numpy
import v190_utils

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

################################################################################
# Correct for a poor correlator model (wrong source or antenna positions)
def correct_propermotion(rootdir, experiment, klass, seqno):
    uvdata = AIPSUVData(experiment.upper(), klass, 1, seqno)
    pmdir = rootdir + '/' + experiment + '/propermotion/'

    #Check the required directory and delay files exist
    if not os.path.exists(pmdir):
      print "There is no propermotion/ subdirectory - aborting!!"
      sys.exit()

    file1path = os.path.join(pmdir, experiment + '.start.delay')
    file2path = os.path.join(pmdir, experiment + '.end.delay')

    if not os.path.exists(file1path):
      print "There is no .start.delay file - aborting!!"
      sys.exit()
    if not os.path.exists(file2path):
      print "There is no .end.delay file - aborting!!"
      sys.exit()

    #Open both the .delay files, skip past the headers
    startdelayin = open(file1path)
    enddelayin = open(file2path)
    year = int(startdelayin.readline()[20:-1])
    enddelayin.readline()
    month = int(startdelayin.readline()[20:-1])
    enddelayin.readline()
    day = int(startdelayin.readline()[20:-1])
    enddelayin.readline()
    hour = int(startdelayin.readline()[20:-1])
    enddelayin.readline()
    minute = int(startdelayin.readline()[20:-1])
    enddelayin.readline()
    second = float(startdelayin.readline()[20:-1])
    enddelayin.readline()
    increment = float(startdelayin.readline()[20:-1])
    enddelayin.readline()
    filetelescopes = int(startdelayin.readline()[20:-1])
    enddelayin.readline()
    antable = uvdata.table('AN', 1)
    antenna_indices  = []
    for row in antable:
        antenna_indices.append(-1)
    for i in range(filetelescopes):
        telescopename = startdelayin.readline()[20:-1]
        enddelayin.readline()
        count = 0
        for row in antable:
            if row.anname == telescopename:
                antenna_indices[count] = i
            count = count + 1
    if count != len(antable):
        print "Only found " + str(count) + " antennas out of " + str(len(antable)) + " in the AN table: aborting!!"
        sys.exit()
    numscans = int(startdelayin.readline()[20:-1])
    enddelayin.readline()
    
    startlines = startdelayin.readlines()
    endlines = enddelayin.readlines()
    
    #Work out the number of scans and the length of the file in seconds
    numscans = 0
    scanpoints = []
    scanstarts = []
    numpoints = 0
    i = 0
    while i < len(startlines):
      numscans = numscans+1
      scanpoints.append(int(startlines[i][20:-1]))
      scanstarts.append(int(startlines[i+1][20:-1]))
      i = i+6+scanpoints[numscans-1] #Skip the scan start line, source name, and numpoints + 3 delay entries
      numpoints = numpoints + scanpoints[numscans-1]

    #Create and attach a new CL table
    oldcl = uvdata.table('CL', 1)
    newcl = uvdata.attach_table('CL', 2, no_term=oldcl.keywords['NO_TERM'])
    newcl.keywords['NO_ANT'] = oldcl.keywords['NO_ANT']
    newcl.keywords['NO_POL'] = oldcl.keywords['NO_POL']
    newcl.keywords['NO_IF'] = oldcl.keywords['NO_IF']

    #Work out the reference frequencies
    reffreqs = []
    fqtableentry = uvdata.table('FQ', 1)[0]
    numfreqs = oldcl.keywords['NO_IF']
    print uvdata.header
    for i in range(numfreqs):
        reffreqs.append(float(fqtableentry.if_freq[i] + uvdata.header.crval[2]))
        print "Reffreq[" + str(i) + "] is " + str(reffreqs[i]) + " Hz"
    refwavelength = 299792458.0/reffreqs[0]
    print "Reference wavelength is " + str(refwavelength)

    #Open the CL table and go through it line by line
    timeoffset = float(hour)/24.0 + float(minute)/1440.0 + float(second)/86400.0
    lasttime = -999.9
    cl_interval_seconds = 86400.0*(oldcl[len(antable)].time - oldcl[0].time)
    print "CL Table interval in seconds is " + str(cl_interval_seconds)
    for row in oldcl:
        if not row.time == lasttime:
            #Work out which scan we are up to
            atscan = numscans - 1
            while atscan > 0 and timeoffset + scanstarts[atscan]*increment/86400.0 > row.time:
              atscan = atscan - 1
            time = timeoffset + float(scanstarts[atscan])/86400.0
            nearestindex = int(((row.time - time)*86400.0)/increment + 0.5)
            lasttime = row.time
            print "atscan is " + str(atscan) + ", row.time is " + str(row.time) + ", time is " + str(time) + ", nearestindex is " + str(nearestindex) + ", scanpoints[atscan] is " + str(scanpoints[atscan]) + ", index will be " + str(nearestindex + scanstarts[atscan] + atscan*6 + 4) + ", len(lines) is " + str(len(startlines))
            startdelays = startlines[nearestindex + scanstarts[atscan] + atscan*6 + 4][20:-1].split()
            enddelays = endlines[nearestindex + scanstarts[atscan] + atscan*6 + 4][20:-1].split()
            nextstartdelays = startlines[nearestindex + scanstarts[atscan] + atscan*6 + 5][20:-1].split()
            nextenddelays = endlines[nearestindex + scanstarts[atscan] + atscan*6 + 5][20:-1].split()

        #Find which antenna we want, go locate the appropriate value, and put it in.
        currenterror = ((float(enddelays[antenna_indices[row.antenna_no-1]]) - float(startdelays[antenna_indices[row.antenna_no-1]]))*float(nearestindex+scanstarts[atscan]))/float(1000000*numpoints)
        nexterror = ((float(nextenddelays[antenna_indices[row.antenna_no-1]]) - float(nextstartdelays[antenna_indices[row.antenna_no-1]]))*float(nearestindex+scanstarts[atscan]))/float(1000000*numpoints)
        secpersec = (nexterror-currenterror)/float(cl_interval_seconds)
        print "Current error is " + str(currenterror) + ", rate is " + str(secpersec)
        for i in range(len(reffreqs)):
          row.delay_1[i] = currenterror
          row.delay_2[i] = currenterror
          row.rate_1[i] = secpersec
          row.rate_2[i] = secpersec
          phase = currenterror*reffreqs[i]*2.0*math.pi
          #print "rate of freq " + str(i) + "(" + str(reffreqs[i]) + ") is " + str(row.rate_1[i]) + ", phase is " + str(phase)
          row.real1[i] = math.cos(phase)
          row.real2[i] = row.real1[i]
          row.imag1[i] = math.sin(phase)
          row.imag2[i] = row.imag1[i]
        newcl.append(row)
    newcl.close()
    oldcl.close()