import datetime, os, re, string

########################################################################################
# Available functions intended for outside use
########################################################################################

def get_logs(prefix, experiment):
    #go away to the various ftp sites and get the logs
    #If you can't find one make a fake ANTAB one, with one line
    print 'Not yet implemented - you\'d better get them yourself :)'


def create_antab(prefix, experiment):
    logfiledir = prefix + '/' + experiment + '/logs/tsys/'

    #Get the Tsys values that were applied at the correlator [Pk, At, Mp, Ho, Cd, Ti]
    #The get_applied_tsys call also edits the antenna names and gets the timerange
    ant_names = ['PKS','CAT','MOP','HOB','CED','DSS']
    timerange = [0.0,0.0]
    applied_tsys, ant_names = get_applied_tsys(prefix, experiment, ant_names, timerange)

    #Remove any existing concatenated file
    if os.path.isfile(logfiledir + experiment + '.antab'):
        os.remove(logfiledir + experiment + '.antab')
    
    #Get a listing of all the log files in the tsys directory
    logfiles = os.listdir(logfiledir)
    
    #Create the single antab file we will loop through
    found = False
    antab_out = open(logfiledir + experiment + '.antab', 'w')
    line = '!' + experiment + '.' + 'antab generated by create_antab on ' + str(datetime.date.today()) + '\n'
    antab_out.write(line)
    line = 'INDEX = \'L1\', \'R1\', \'L2\', \'R2\', \'L3\', \'R3\', \'L4\', \'R4\' /\n'
    antab_out.write(line)

    #Do any Parkes files
    for logfile in logfiles:
        if (re.search('parkes', logfile) or re.search('PARKES', logfiledir)) and not re.search('antab', logfile): #its a PARKES logfile
            add_parkes_file(logfiledir+logfile, applied_tsys, ant_names, timerange, antab_out)
            found = True
    if found: antab_out.write('/\n')
    found = False

    #Do any ATCA files
    for logfile in logfiles:
        if re.search('atca', logfile) or re.search('ATCA', logfiledir): #its an ATCA logfile
            add_atca_file(logfiledir+logfile, applied_tsys, ant_names, timerange, antab_out)
            found = True
    if found: antab_out.write('/\n')
    found = False
    
    #Do any Mopra files
    for logfile in logfiles:
        if re.search('mopra', logfile) or re.search('MOPRA', logfiledir): #its a MOPRA logfile
            add_mopra_file(logfiledir+logfile, applied_tsys, ant_names, timerange, antab_out)
            found = True
    if found: antab_out.write('/\n')
    found = False
    
    #Do any ANTAB format files (should include Hob, Ced and Tid)
    for logfile in logfiles:
        if re.search('antab', logfile) or re.search('ANTAB', logfiledir): #its an antab logfile
            add_antab_file(logfiledir+logfile, applied_tsys, ant_names, timerange, antab_out)

    antab_out.close()

########################################################################################
# Private functions
########################################################################################

def add_parkes_file(logfile, applied_tsys, ant_names, timerange, antab_out):
    if applied_tsys[0] > 0.0:
        line = 'GAIN ' + ant_names[0] + ' ALTAZ DPFU = 1.0 POLY = 1.0 /\n'
        antab_out.write(line)
        line = 'TSYS ' + ant_names[0] + ' TIMEOFF = 0.0 FT = ' + str(1.0/applied_tsys[0]) + ' /\n'
        antab_out.write(line)
        #FOR NOW, DUMMY DATA FOR PARKES (LEAVE A PRIORI UNCHANGED)
        doy = mjd2doy(int(timerange[0]))
        hour = int((timerange[0] - int(timerange[0]))*24)
        minute = 0
        runminutes = int((timerange[1] - timerange[0])*1440)
        for i in range(runminutes/20 + 3):
            line = "%03d %02d:%02d.00 %.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n" % (doy,hour,minute,applied_tsys[0],applied_tsys[0],applied_tsys[0],applied_tsys[0],applied_tsys[0],applied_tsys[0],applied_tsys[0],applied_tsys[0])
            minute = minute + 20
            hour = hour + minute/60
            minute = minute % 60
            doy = doy + hour/24
            hour = hour % 24
            antab_out.write(line)
        #END DUMMY PARKES DATA
        applied_tsys[0] = -1.0 #Only write the header once
        
    infile = open(logfile, 'r')
    lines = infile.readlines()
    infile.close()
    #Actually do the parkes style files!!!

def add_atca_file(logfile, applied_tsys, ant_names, timerange, antab_out):
    if applied_tsys[1] > 0.0:
        line = 'GAIN ' + ant_names[1] + ' ALTAZ DPFU = 1.0 POLY = 1.0 /\n'
        antab_out.write(line)
        line = 'TSYS ' + ant_names[1] + ' TIMEOFF = 0.0 FT = ' + str(10.0/applied_tsys[1]) + ' /\n'
        antab_out.write(line)
        applied_tsys[1] = -1.0 #Only write the header once

    infile = open(logfile, 'r')
    lines = infile.readlines()
    infile.close()
    for line in lines:
        splitline = line.split()
        if len(splitline) > 1:
            doy = atca_datecontained(splitline[0], splitline[1], timerange) #0 if not contained
            if splitline[2] == 'TSYS' and doy:
                day = int(doy)
                hh = int(24.0*(doy-day))
                mm = int(1440.0*(doy - day - hh/24.0))
                ff = int(144000*(doy - day - hh/24.0 - mm/1440.0) + 0.5)
                line = "%03d %02d:%02d.%02d %.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n" % (day,hh,mm,ff,float(splitline[4]),float(splitline[5]),float(splitline[4]),float(splitline[5]),float(splitline[6]),float(splitline[7]),float(splitline[6]),float(splitline[7]))
                antab_out.write(line)

def add_mopra_file(logfile, applied_tsys, ant_names, timerange, antab_out):
    if applied_tsys[2] > 0.0:
        line = 'GAIN ' + ant_names[2] + ' ALTAZ DPFU = 1.0 POLY = 1.0 /\n'
        antab_out.write(line)
        line = 'TSYS ' + ant_names[2] + ' TIMEOFF = 0.0 FT = ' + str(10.0/applied_tsys[2]) + ' /\n'
        antab_out.write(line)
        applied_tsys[2] = -1.0 #Only write the header once

    datawritten = False
    infile = open(logfile, 'r')
    tsys = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    instanttsys = [0.0, 0.0, 0.0, 0.0]
    numaveraged = 0
    lastminute = 0
    secondssum = 0
    lasthour = 0
    lines = infile.readlines()
    infile.close()
    year = string.atoi((lines[0][23:27]))
    print 'Mopra year is ' + str(year)
    for line in lines:
        if line[0] != '!' and len(line) > 1: #Not a comment
            #print 'Working with line ' + line
            splitline = line.split()
            doy = string.atoi(splitline[0].split('/')[0])
            hour = string.atoi(splitline[0].split('/')[1].split(':')[0])
            minute = string.atoi(splitline[0].split('/')[1].split(':')[1])
            second = int(string.atof(splitline[0].split('/')[1].split(':')[2]))
            if mopra_datecontained(year, doy, hour, minute, second, timerange):
                if datawritten and (minute != lastminute or hour != lasthour):
                    line = "%03d %02d:%02d.%02d %.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n" % (doy,hour,minute,int(float(secondssum)*10.0/(6.0*numaveraged) + 0.5),tsys[0]/numaveraged,tsys[1]/numaveraged,tsys[2]/numaveraged,tsys[3]/numaveraged,tsys[4]/numaveraged,tsys[5]/numaveraged,tsys[6]/numaveraged,tsys[7]/numaveraged)
                    antab_out.write(line)
                    lastminute = minute
                    lasthour = hour
                    tsys = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] #clear sum
                    numaveraged = 0
                    secondssum = 0
                    datawritten = False
                for i in range(4):
                    instanttsys[i] = string.atof(splitline[i+3])
                for i in range(4):
                    if instanttsys[i] < 0.00000001:
                        for j in range(4):
                            if instanttsys[3-j] > 0.00000001:
                                instanttsys[i] = instanttsys[3-j]
                if(instanttsys[0] > 0.0000001):
                    datawritten = True
                    numaveraged = numaveraged + 1
                    secondssum = secondssum + second
                    tsys[0] = tsys[0] + instanttsys[0]
                    tsys[1] = tsys[1] + instanttsys[1]
                    tsys[2] = tsys[2] + instanttsys[0]
                    tsys[3] = tsys[3] + instanttsys[1]
                    tsys[4] = tsys[4] + instanttsys[2]
                    tsys[5] = tsys[5] + instanttsys[3]
                    tsys[6] = tsys[6] + instanttsys[2]
                    tsys[7] = tsys[7] + instanttsys[3]
                    
def add_antab_file(logfile, applied_tsys, ant_names, timerange, antab_out):
    #Work out which antenna it is
    infile = open(logfile, 'r')
    print 'Checking antab style logfile ' + logfile
    lines = infile.readlines()
    infile.close()
    at = 0
    while len(lines[at]) > 1 and lines[at].strip()[0] == '!':
        at = at + 1
    splitline = lines[at].split()

    if splitline[1] == 'PARKES' or splitline[1] == 'PKS':
        index = 0
    elif splitline[1] == 'HOBART' or splitline[1] == 'HOB':
        index = 3
    elif splitline[1] == 'CEDUNA' or splitline[1] == 'CED':
        index = 4
    elif splitline[1] == 'TID' or splitline[1] == 'DSS43' or splitline[1] == 'DSS34':
        index = 5
    else:
        print 'UNKNOWN ANTENNA: ' + splitline[1]
        return

    #If its Parkes, work out the averaged value
    scale = 1.0
    if index == 0:
        pcount = 0
        psum = 0.0
        for line in lines[at+1:]:
            if line[0] != '!' and line[0] != '/': #Not a comment
                splitline = line.split()
                numvals = len(splitline) - 2
                if numvals > 2: numvals = 2
                for i in range(numvals):
                    psum = psum + string.atof(splitline[2+i])
                pcount = pcount + numvals
        average = psum/pcount
        scale = applied_tsys[index]/average

    #Write the header lines - copy from the input file
    line = 'GAIN ' + ant_names[index] + ' ALTAZ DPFU = 1.0 POLY = 1.0 /\n'
    antab_out.write(line)
    line = 'TSYS ' + ant_names[index]
    for i in range(len(splitline)):
        if splitline[i] == 'TIMEOFF':
            line = line + ' TIMEOFF = ' + splitline[i+2]
    line = line + ' FT = ' + str(scale/applied_tsys[index]) + '/\n'
    antab_out.write(line)

    #Write all the data lines
    averagedtsys = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    instanttsys = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    numaveraged = 0
    seconds = 0
    lasthour = -1
    lastminute = -1
    at = at + 1
    for line in lines[at:]:
        if line[0] != '!' and line[0] != '/': #Not a comment
            splitline = line.split()
            numvals = len(splitline) - 2
            if index == 0 and numvals > 2: #Parkes file
                numvals = 2
            if numvals <= 0:
                print "Read a blank line..."
            elif numvals == 1:
                instanttsys[0] = instanttsys[1] = instanttsys[2] = instanttsys[3] = string.atof(splitline[2])
                instanttsys[4] = instanttsys[5] = instanttsys[6] = instanttsys[7] = instanttsys[0]
            elif numvals == 2:
                instanttsys[0] = instanttsys[2] = instanttsys[4] = instanttsys[6] = string.atof(splitline[2])
                instanttsys[1] = instanttsys[3] = instanttsys[5] = instanttsys[7] = string.atof(splitline[3])
            else:
                instanttsys[0] = instanttsys[2] = string.atof(splitline[2])
                instanttsys[1] = instanttsys[3] = string.atof(splitline[3])
                instanttsys[4] = instanttsys[6] = string.atof(splitline[4])
                instanttsys[5] = instanttsys[7] = string.atof(splitline[5])
            timesplit = splitline[1].split(':')
            hour = string.atoi(timesplit[0])
            minute = string.atoi(timesplit[1])
            if numaveraged > 0 and (hour != lasthour or minute != lastminute):
                outline = "%03d %02d:%02d.%02d %.3f %.3f %.3f %.3f %.3f %.3f %.3f %.3f\n" % (int(day),int(lasthour),int(lastminute),int(float(seconds)/(numaveraged*60.0/100.0)),(averagedtsys[0]/numaveraged),(averagedtsys[1]/numaveraged),(averagedtsys[2]/numaveraged),(averagedtsys[3]/numaveraged),(averagedtsys[4]/numaveraged),(averagedtsys[5]/numaveraged),(averagedtsys[6]/numaveraged),(averagedtsys[7]/numaveraged))
                antab_out.write(outline)
                numaveraged = 0
                seconds = 0
                for i in range(8):
                    averagedtsys[i] = 0.0
                
            day = splitline[0]
            lasthour = hour
            lastminute = minute
            seconds = seconds + int(string.atof(timesplit[2]))
            numaveraged = numaveraged + 1
            for i in range(8):
                averagedtsys[i] = averagedtsys[i] + instanttsys[i]
    antab_out.write('/\n')
            
def get_applied_tsys(prefix, experiment, ant_names, timerange):
    try:
        corrinput = open(prefix + '/' + experiment + '/' + experiment + '.input', 'r')
    except IOError:
        files = os.listdir(prefix + '/' + experiment + '/')
        for filename in files:
            if re.search('input', filename) and (not re.search('log', filename)) and (not os.path.isdir(filename)):
                print 'Could not find ' + prefix + '/' + experiment + '/' + experiment + '.input - using ' + filename + ' instead'
                corrinput = open(prefix + '/' + experiment + '/' + filename, 'r')
                break
    lines = corrinput.readlines()
    corrinput.close()
    positions = [0, 1, 2, 3, 4, 5] #Positions of these stations in the input file telescope table
    tsys = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] #This is what we will fill and return
    new_names = ['', '', '', '', '', ''] #Also to be filled and returned

    #Work out the start and end MJD
    at = 0 
    while at < len(lines) and lines[at][0:12] != 'EXECUTE TIME':
        at = at+1
    timerange[1] = string.atof(getvalue(lines, at))/86400.0
    at = at+1
    timerange[0] = string.atof(getvalue(lines, at))
    at = at+1
    timerange[0] = timerange[0] + string.atof(getvalue(lines, at))/86400.0
    timerange[1] = timerange[1] + timerange[0]
    
    #Find the start of the telescope  table
    while at < len(lines) and lines[at][0:21] != '# TELESCOPE TABLE ##!':
        at = at + 1

    #Work out the order of the telescopes
    at = at+1
    numtelescopes = string.atoi(getvalue(lines, at))
    for offset in range(numtelescopes):
        at = at+1
        stationname = getvalue(lines, at)
        for i in range(6):
            if ant_names[i] == stationname[0:3]:
                positions[i] = offset
                new_names[i] = stationname
        at = at+1
        at = at+1

    #Find the start of the datastream table
    at = at + 1
    while at < len(lines) and lines[at][0:21] != '# DATASTREAM TABLE #!':
        at = at + 1

    #Grab and store the tsys values
    at = at+1
    numdatastreams = string.atoi(getvalue(lines, at))
    for offset in range(numdatastreams):
        at = at+1
        while lines[at][0:4] != 'TSYS':
            at = at+1
        for i in range(6):
            if positions[i] == offset:
                tsys[i] = string.atof(getvalue(lines, at))

    return tsys, new_names

def atca_datecontained(datestring, timestring, timerange):
    datecomponents = datestring.split('-')
    timecomponents = timestring.split(':')
    if len(datecomponents) < 3 or len(timecomponents) < 3:
        return False
    months = ['JAN','FEB','MAR','APR','MAY','JUN','JUL','AUG','SEP','OCT','NOV','DEC']
    year =string.atoi(datecomponents[2])
    month = 0
    day = string.atoi(datecomponents[0])
    for i in range(12):
        if months[i] == datecomponents[1].upper():
            month = i+1
    if month == 0:
        print 'Error locating ATCA month!!!\n'

    mjd = year*367 - int(7*(year + int((month + 9)/12))/4) + int(275*month/9) + day - 678987
    mjd = mjd + string.atof(timecomponents[0])/24.0 + string.atof(timecomponents[1])/1440.0 + string.atof(timecomponents[2])/86400.0
    if mjd < timerange[0] or mjd > timerange[1]:
        return 0
    return mjd + 678957 - year*367 + int(7*year/4) #Convert to day of year (with fractional component)

def mopra_datecontained(year, doy, hour, minute, second, timerange):
    mjd = year*367 - int(7*year/4) + doy - 678957
    mjd = float(mjd) + float(hour)/24.0 + float(minute)/1440.0 + float(second)/86400.0
    return mjd > timerange[0] and mjd < timerange[1]

def getvalue(filelines, linenumber):
    if len(filelines[linenumber]) < 21:
           return ''
    return filelines[linenumber].rstrip()[20:]

def mjd2doy(mjd):
    year = 2000
    doy = mjd + 678957 + int(7*year/4) - year*367
    while doy > 366 and (doy > 365 or year%4 == 0):
        year = year + 1
        doy = mjd + 678957 + int(7*year/4) - year*367

    return doy