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

################################################################################
# Correct for a poor correlator model (wrong source or antenna positions)
def correct_model(rootdir, experiment, klass, correct_bestlinear, skipcheck):
    print "Correcting model for experiment " + experiment
    uvdata = AIPSUVData(experiment.upper(), klass, 1, 1)
    applieddir = rootdir + '/' + experiment + '/applied/'
    #If the applied directory already exists, ask whether to proceed with
    #what's in there and a new vex2model, otherwise abort
    if os.path.exists(applieddir):
        if (not skipcheck) and (not v190_utils.yesno(False, 'Is the existing model data in applied/ directory that which was applied?  Yes to continue (option to rerun vex2model for most up-to-date model), difference between new model and whats in applied/ will be corrected, no to abort pipeline: ')):
            sys.exit()
        if not os.path.exists(applieddir + experiment + '.uvw') or not os.path.exists(applieddir + experiment + '.delay'):
            print "Applied directory exists but either .delay or .uvw file is missing - aborting!!!"
            sys.exit()
        if (not skipcheck) and v190_utils.yesno(False, 'Rerun vex2model?'):
            vex2modelcommand = 'vex2model ' + experiment + '.skd'
            os.system(vex2modelcommand)
    else:
        #create applied directory and move existing model files there
        os.mkdir(applieddir)
        os.system('mv ' + experiment + '.uvw applied/')
        os.system('mv ' + experiment + '.delay applied/')
        vex2modelcommand = 'vex2model ' + experiment + '.skd'
        os.system(vex2modelcommand)

    #Open the .uvw file, read the antenna values and update the AN table if required
    antable = uvdata.table('AN', 1)
    uvwin = open(rootdir + '/' + experiment + '/' + experiment + '.uvw')
    uvwyear = int(uvwin.readline()[20:-1])
    uvwmonth = int(uvwin.readline()[20:-1])
    uvwday = int(uvwin.readline()[20:-1])
    uvwhour = int(uvwin.readline()[20:-1])
    uvwminute = int(uvwin.readline()[20:-1])
    uvwsecond = float(uvwin.readline()[20:-1])
    uvwincrement = int(uvwin.readline()[20:-1])
    uvwfiletelescopes = int(uvwin.readline()[20:-1])
    antennas_updated = []
    antenna_indices  = []
    #rows = []
    #rowindices = []
    for row in antable:
        antennas_updated.append(False)
        antenna_indices.append(-1)
    
    for i in range(uvwfiletelescopes):
        name = uvwin.readline()[20:-1]
        mount = uvwin.readline()[20:-1]
        x = float(uvwin.readline()[20:-1])
        y = float(uvwin.readline()[20:-1])
        z = float(uvwin.readline()[20:-1])
        count = 0
        for row in antable:
            if row.anname.strip() == name.strip():
                row.stabxyz = [x,y,z]
                #rows.append(row)
                #rowindices.append(count)
                row.update()
                antennas_updated[count] = True
            count = count+1
    count = 0
    for val in antennas_updated:
        if not val:
            print "Didn't find antenna " + antable[count].anname + " in model file - aborting!!"
            sys.exit()
        count = count + 1
    antable.close()

    #Delete the antable and recreate
    #rdate  = antable.keywords['RDATE']
    #timsys = antable.keywords['TIMSYS']
    #datutc = antable.keywords['DATUTC']
    #ut1utc = antable.keywords['UT1UTC']
    #arrnam = antable.keywords['ARRNAM']
    #freq   = antable.keywords['FREQ']
    #nopcal = antable.keywords['NOPCAL']
    #NormUVData(experiment.upper(), klass, 1, 1).table('AN', 1).zap()
    #newantable = uvdata.attach_table('AN', 1)
    #newantable.keywords['RDATE']  = rdate
    #newantable.keywords['TIMSYS'] = timsys
    #newantable.keywords['DATUTC'] = datutc
    #newantable.keywords['UT1UTC'] = ut1utc
    #newantable.keywords['ARRNAM'] = arrnam
    #newantable.keywords['FREQ']   = freq
    #newantable.keywords['NOPCAL'] = nopcal
    #for i in range(count):
    #    rcount = 0
    #    for rindex in rowindices:
    #        if rindex == i:
    #            newantable.append(rows[rcount])
    #            break
    #        rcount = rcount + 1
    
    #Read through the .uvw file, find 1 instance of each source, update SU table if required
    numsources = len(uvdata.sources)
    sutable = uvdata.table('SU', 1)
    line = uvwin.readline()
    for row in sutable:
        while line and not line[20:-1].rstrip() == row.source.rstrip():
            line = uvwin.readline()
        if not line:
            print "Reached end of UVW file without finding antenna " + row.source.rstrip() + ", aborting!!!"
            sys.exit()
        print "Found " + row.source.rstrip() + ", ra is " + str(row.raepo) + ", dec is " + str(row.decepo)
        newraepo =  float(uvwin.readline()[20:-1])*180/math.pi
        newdecepo = float(uvwin.readline()[20:-1])*180/math.pi
        row.raapp = newraepo - row.raepo + row.raapp
        row.decapp = newdecepo - row.decepo + row.decapp
        row.raepo = newraepo
        row.decepo = newdecepo
        row.update()
    sutable.close()
    uvwin.close()

    #Open the .delay file and skip the header
    delayin = open(rootdir + '/' + experiment + '/' + experiment + '.delay')
    olddelayin = open(rootdir + '/' + experiment + '/applied/' + experiment + '.delay')
    year = int(delayin.readline()[20:-1])
    olddelayin.readline()
    month = int(delayin.readline()[20:-1])
    olddelayin.readline()
    day = int(delayin.readline()[20:-1])
    olddelayin.readline()
    hour = int(delayin.readline()[20:-1])
    olddelayin.readline()
    minute = int(delayin.readline()[20:-1])
    olddelayin.readline()
    second = float(delayin.readline()[20:-1])
    olddelayin.readline()
    increment = float(delayin.readline()[20:-1])
    olddelayin.readline()
    filetelescopes = int(delayin.readline()[20:-1])
    olddelayin.readline()
    antable = uvdata.table('AN', 1)
    foundcount = 0
    for i in range(filetelescopes):
        telescopename = delayin.readline()[20:-1]
        oldtelescopename = olddelayin.readline()[20:-1]
        if telescopename != oldtelescopename:
            print "Telescope " + str(i) + " does not match between delay files: old has (" + oldtelescopename + "), new has (" + telescopename + ") - aborting!!!"
            sys.exit(1)
        count = 0
        for row in antable:
            if row.anname.rstrip() == telescopename.rstrip():
                print "Antenna index for antenna " + str(count) + " (" + row.anname.rstrip() + ") is " + str(i)
                antenna_indices[count] = i
                foundcount = foundcount + 1
            count = count + 1
    if foundcount != len(antable):
        print "Only found " + str(foundcount) + " antennas out of " + str(len(antable)) + " in the AN table: aborting!!"
        sys.exit()
    numscans = int(delayin.readline()[20:-1])
    olddelayin.readline()

    #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']
    numpol = oldcl.keywords['NO_POL']

    #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"
    fqtable = uvdata.table('FQ', 1)
    for row in fqtable:
        try:
            for iffreq in row.if_freq:
                freqentry = float(iffreq) + float(uvdata.header.crval[2])
                reffreqs.append(float(iffreq) + float(uvdata.header.crval[2]))
                print "Next reffreq is " + str(freqentry) + " Hz"
        except (AttributeError, TypeError):
                freqentry = float(row.if_freq) + float(uvdata.header.crval[2])
                reffreqs.append(float(row.if_freq) + float(uvdata.header.crval[2]))
                print "Next reffreq is " + str(freqentry) + " Hz"
    refwavelength = 299792458.0/reffreqs[0]
    print "Reference wavelength is " + str(refwavelength)

    #Open the CL table and go through it line by line
    delayarray = numpy.zeros((len(oldcl), len(antable)+1), numpy.float64)
    a = numpy.zeros(len(antable), numpy.float64)
    b = numpy.zeros(len(antable), numpy.float64)
    oldline = ''
    lastline = ''
    lastoldline = ''
    scan = 0
    scanstart = 0
    scancount = 0
    scanpoints = 0
    timeoffset = float(hour)/24.0 + float(minute)/1440.0 + float(second)/86400.0
    time = timeoffset - increment/86400.0
    clcount = -1
    lasttime = -999.9
    i = 1
    while oldcl[i].time == oldcl[0].time:
        i = i + 1
    cl_interval_seconds = 86400.0*(oldcl[i].time - oldcl[0].time)
    print "CL Table interval in seconds is " + str(cl_interval_seconds)
    first_end = True
    for row in oldcl:
        if not row.time == lasttime:
            clcount = clcount + 1
            lasttime = row.time
            if clcount > 0:
                #Average all the offsets for that time and subtract average
                for i in range(len(antable)):
                    delayarray[clcount-1][len(antable)] = delayarray[clcount-1][len(antable)] + delayarray[clcount-1][i]
                for i in range(len(antable)):
                    delayarray[clcount-1][i] = delayarray[clcount-1][i] - delayarray[clcount-1][len(antable)]/len(antable)
        while time < row.time:
            if scancount == scanpoints: #Time to move on to the next scan
                if scan == numscans:
                    if first_end:
                        print "Run out of data in the .delay file - will see if this is the last row..."
                        time = row.time	
                        first_end = False
                        break
                    else:
                        print "Run out of data in the .delay file - aborting!!!"
                        sys.exit()
                if scan != 0:
                    # Skip past the end of scan interpolation point
                    line = delayin.readline()
                    olddelayin.readline()
                scan = scan + 1
                scanpoints = int(delayin.readline()[20:-1])
                olddelayin.readline()
                scanstart = int(delayin.readline()[20:-1])
                olddelayin.readline()
                time = timeoffset + float(scanstart-1)*increment/86400.0
                line = delayin.readline() #Skip the source name
                olddelayin.readline()
                lastline = delayin.readline()[20:-1] # The -1 index
                lastoldline = olddelayin.readline()[20:-1] # The -1 index
                line = delayin.readline()[20:-1]
                oldline = olddelayin.readline()[20:-1]
                scancount = 0
            lastline = line
            lastoldline = oldline
            line = delayin.readline()[20:-1]
            oldline = olddelayin.readline()[20:-1]
            scancount = scancount+1
            time = timeoffset + float(scanstart+scancount)*increment/86400.0
        #Put in the correct delay and phase for this antenna - ignore clock and clock rate for now,
        #would need to open the .input file
        delays = line.split()
        olddelays = oldline.split()
        lastdelays = lastline.split()
        lastolddelays = lastoldline.split()
        currenterror = float(delays[antenna_indices[row.antenna_no-1]]) - float(olddelays[antenna_indices[row.antenna_no-1]])
        lasterror = float(lastdelays[antenna_indices[row.antenna_no-1]]) - float(lastolddelays[antenna_indices[row.antenna_no-1]])
        delayarray[clcount][row.antenna_no-1] = currenterror + (row.time - time)*(currenterror-lasterror)*increment/86400.0

    if correct_bestlinear:
        for j in range(len(antable)):
            sumx = 0
            sumy = 0
            sumxx = 0
            sumxy = 0
            sumyy = 0
            for i in range(clcount):
                #Work out a best-fit linear approximation for this antenna
                sumx = sumx + i
                sumy = sumy + delayarray[i][j]
                sumxx = sumxx + i*i
                sumxy = sumxy + i*delayarray[i][j]
                sumyy = sumyy + delayarray[i][j]*delayarray[i][j]
            Sxx = sumxx - sumx*sumx/(clcount)
            yy  = sumyy - sumy*sumy/(clcount)
	    Sxy = sumxy - sumx*sumy/(clcount)
            b = Sxy/Sxx
            a = (sumy - b*sumx)/(clcount)
            for i in range(clcount + 1):
                #Subtract best-fit linear (clock rate error estimate) for this antenna
                delayarray[i][j] = delayarray[i][j] - (a + b*i)
            

    #Actually write the new cl table
    clcount = -1
    lasttime = -999.9
    lastsource = -1
    for row in oldcl:
        if not row.time == lasttime:
            clcount = clcount + 1
            lasttime = row.time
        delay = delayarray[clcount][row.antenna_no-1]
        if row.source_id == lastsource:
            lastdelay = delayarray[clcount-1][row.antenna_no-1]
        else:
            if row.antenna_no == len(antable):
                lastsource = row.source_id
            lastdelay =  2*delayarray[clcount][row.antenna_no-1] - delayarray[clcount+1][row.antenna_no-1]
        deltadelay = (delay-lastdelay)
        secpersec = deltadelay/(1000000.0*cl_interval_seconds)
        #print "Lastdelay is " + str(lastdelay) + ", deltadelay is " + str(deltadelay) + ", secpersec is " + str(secpersec)
        newdelay1 = []
        newrate1 = []
        newreal1 = []
        newimag1 = []
        newdelay2 = []
        newrate2 = []
        newreal2 = []
        newimag2 = []
        for i in range(len(reffreqs)):
            newdelay1.append(delay/1000000.0)
            newrate1.append(secpersec)
            phase = delay*reffreqs[i]*2.0*math.pi/1000000.0
            newreal1.append(math.cos(phase))
            newimag1.append(math.sin(phase))
            if numpol > 1:
                newdelay2.append(delay/1000000.0)
                newrate2.append(secpersec)
                newreal2.append(math.cos(phase))
                newimag2.append(math.sin(phase))
        row.delay_1 = newdelay1
        row.rate_1 = newrate1
        row.real1 = newreal1
        row.imag1 = newimag1
        if numpol > 1:
            row.delay_2 = newdelay2
            row.rate_2 = newrate2
            row.real2 = newreal2
            row.imag2 = newimag2
        #for i in range(len(reffreqs)):
        #    row.delay_1[i] = delay/1000000.0
        #    row.rate_1[i] = secpersec
        #    phase = delay*reffreqs[i]*2.0*math.pi/1000000.0
        #    #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.imag1[i] = math.sin(phase)
        #    if numpol > 1:
        #        row.delay_2[i] = delay/1000000.0
        #        row.rate_2[i] = secpersec
        #        row.real2[i] = row.real1[i]
        #        row.imag2[i] = row.imag1[i]
        newcl.append(row)       

    newcl.close()
    oldcl.close()

    #Finally, update the uvw values of the actual visibilities - need to reopen uvw file
    lastuvws = numpy.zeros((len(antable), 3), numpy.float64)
    lastolduvws = numpy.zeros((len(antable), 3), numpy.float64)
    currentuvws = numpy.zeros((len(antable), 3), numpy.float64)
    currentolduvws = numpy.zeros((len(antable), 3), numpy.float64)
    duvw =  numpy.zeros(3, numpy.float64)
    newuvws = numpy.zeros(3, numpy.float64)
    uvwin = open(rootdir + '/' + experiment + '/' + experiment + '.uvw')
    olduvwin = open(rootdir + '/' + experiment + '/applied/' + experiment + '.uvw')
    while not line[0:20] == "NUM SCANS:          ":
        line = uvwin.readline()
        olduvwin.readline()
    numscans = int(line[20:-1])
    scan = -1
    scanstart = 0
    scancount = 0
    scanpoints = 0
    timeoffset = float(hour)/24.0 + float(minute)/1440.0 + float(second)/86400.0
    scanstarttime = timeoffset
    uvwtime = scanstarttime
    first_end = True
    for vis in uvdata:
        while uvwtime < vis.time - 0.000006:
            #Skip through the vis file
            if scancount == scanpoints:
                if scan+1 == numscans:
                    if first_end:
                        print "Run out of data in the .uvw file - will see if this is the last visibility..."
                        uvwtime = vis.time
                        first_end = False
                        break
                    else:
                        print "Run out of data in the .uvw file - aborting!!!"
                        sys.exit()
                if not scan < 0:
                    olduvwin.readline()
                    olduvwin.readline()
                    uvwin.readline()
                    uvwin.readline() #Skip the last two points
                scanpoints = int(uvwin.readline()[20:-1])
                olduvwin.readline()
                scanstart = int(uvwin.readline()[20:-1])
                olduvwin.readline()
                scanstarttime = timeoffset + float(scanstart*uvwincrement)/86400.0
                scancount = 0
                olduvwin.readline()
                olduvwin.readline()
                olduvwin.readline()
                uvwin.readline()
                uvwin.readline()
                uvwin.readline() #Skip the name, ra and dec
                line = olduvwin.readline()
                extract_uvws(line[20:-1], len(antable), antenna_indices, currentolduvws)
                line = uvwin.readline()
                extract_uvws(line[20:-1], len(antable), antenna_indices, currentuvws)
                scan = scan + 1
                print "Now reading scan " + str(scan) + " which starts at " + str(scanstarttime) + ", while we're looking for visibility time " + str(vis.time)
                
            lastuvws = currentuvws.copy()
            lastolduvws = currentolduvws.copy()
            line = olduvwin.readline()
            extract_uvws(line[20:-1], len(antable), antenna_indices, currentolduvws)
            line = uvwin.readline()
            extract_uvws(line[20:-1], len(antable), antenna_indices, currentuvws)
            uvwtime = scanstarttime + (scancount*uvwincrement)/86400.0
            scancount = scancount + 1
        #Print out what was there
        #print "For baseline " + str(vis.baseline[0]) + "->" + str(vis.baseline[1]) + " at time " + str(vis.time) + ", we are replacing uvw (" + str(vis.uvw[0]) + "," + str(vis.uvw[1]) + "," + str(vis.uvw[2]) + ") with:"
        #Now we've found the right uvws, store it
        t1 = vis.baseline[0]-1
        t2 = vis.baseline[1]-1
        for i in range(3):
            duvw[i] = (currentuvws[t1][i] - currentuvws[t2][i] - lastuvws[t1][i] + lastuvws[t2][i])*86400.0/uvwincrement
            newuvws[i] = (currentuvws[t1][i] - currentuvws[t2][i] + (vis.time - uvwtime)*duvw[i])/refwavelength
            expecteduvw = currentolduvws[t1][i] - currentolduvws[t2][i] + (currentolduvws[t1][i] - currentolduvws[t2][i] - lastolduvws[t1][i] + lastolduvws[t2][i])*86400.0*(vis.time - uvwtime)/uvwincrement
            if(abs(expecteduvw - vis.uvw[i]*refwavelength) > 0.501):
                print "Warning - I don't think the applied/ directory and the data match for time " + str(vis.time) + "on baseline " + str(vis.baseline[0]) + "-" +  str(vis.baseline[1]) + "!"
                print "Data uvw[" + str(i) + "] is " + str(vis.uvw[i]*refwavelength) + ", while applied/ uvw[" + str(i) +"] is " + str(expecteduvw) + "!!! For reference, we will be putting in " + str(newuvws[i]*refwavelength)
                print "t1 uvw is " + str(currentuvws[t1][i]) + ", t2 uvw is " + str(currentuvws[t2][i])
        vis.uvw = newuvws
        vis.update()

        #Print out what we've replaced it with
        #print "                     time " + str(vis.time) + ",                      (" + str(vis.uvw[0]) + "," + str(vis.uvw[1]) + "," + str(vis.uvw[2]) + ")"
        #raw_input("Press any key to continue ")

################################################################################
# Split a line containing uvws and store them in the provided array
def extract_uvws(line, numantennas, antenna_offsets, uvwarray):
    splitline = line.split()
    for i in range(numantennas):
        for j in range(3):
            uvwarray[i][j] = float(splitline[antenna_offsets[i]*3 + j])