""" Parse the output of listr to get information on scan times. This information
is used to drive AIPS to save the text files needed for total power plotting

wmax, Nov 8, 2014
"""

import os
import sys
import imp
import scipy

# load script configuration parameters from PWD
# this will ensure that only the configredrfi.py file on current directory is 
# loaded
PWD = os.environ['PWD']

try:
    configredrfi = imp.load_source('configredrfi',
                             os.path.join(PWD, 'configredrfi.py'))
    from configredrfi import *
except IOError:
    error_message = """No configredrfi.py file was found in %s
use 'redrfi --help' for help\n""" % (PWD)
    sys.exit(error_message)


#-------------------------------------------------------------------------------
def parse_listr_scan_list(filename):
    """ Parse listr with scan list to get information on scan times, numbers,
    sources.
    
    input
    -----
    filename

    output
    ------
    A dictionary with the data for the scans. The form is 
    data = {'scan': scipy.array(SCAN),
            'source': scipy.array(SOURCE),
            'timerang': scipy.array(TIMERANG)
            }
    
    notes
    -----


    wmax, Nov 8, 2014
    """
    
    fileObj = open(filename)

    # Arrays to save data points
    SCAN = []
    SOURCE = []
    TIMERANG = []

    # Reading state. Goes to True after column description
    reading_data = False

    for line in fileObj:

        # Matches the line with column names, just before the data start
        if line[0:16] == 'Scan      Source':
            reading_data = True
            continue

        # Identifies the first line of the next data black and exits 
        elif line[0:14] == 'Source summary':
            reading_data = False
            break

        elif reading_data == True:
            # Parse the line and save to arrays
            fields = line.strip().split()
            
            # Check it is a valid data line, or simply skip it
            # This is required because long schedules have some information lines
            # inserted in between the list of scans
            # The check is very simple but it should work
            if len(fields) == 12 and fields[7] == '-':
                SCAN.append(fields[0])
                SOURCE.append(fields[1])

                # Timerang needs some cleaning to be usable with AIPS
                timerang = '%s,%s' % (fields[6], fields[8])
                timerang = timerang.replace('/', ',')
                timerang = timerang.replace(':', ',')

                TIMERANG.append(timerang)

                
    # Save in dictionary as scipy.arrays for easier manipulation
    data = {'scan': scipy.array(SCAN),
            'source': scipy.array(SOURCE),
            'timerang': scipy.array(TIMERANG)
            }
    
    return data


#-------------------------------------------------------------------------------
def write_aips_script(data,
                      aips_filename='data',
                      antennas=['BR', 'FD', 'HN', 'KP', 'LA',
                                'MK', 'NL', 'OV', 'PT', 'SC'],
                      stokes=['RR', 'LL'],
                      aips_id='2768',
                      aips_ehexid='24w',
                      antenna_number={'BR': 1, 'FD': 2, 'HN': 3, 'KP': 4,
                                      'LA': 5, 'MK': 6, 'NL': 7, 'OV': 8,
                                      'PT': 9, 'SC': 10},
                      flagver=0):
    """ Write an AIPS script to save the total power data

    input
    -----

    output
    ------

    notes
    -----

    wmax, Nov 8, 2014
    """

    # To configure the parameters of possm for plotting, variable parameters
    # are set later. One per antenna and scan
    possm = """%s
* Starts with user ID to run it from the command line

default 'possm'
task 'possm'
inname '%s'; inclass 'uvdata'
sources ''
stokes ''
* Display all channels
bchan 0; echan 0; bif 0; eif 0
docalib 1; gainuse 0
flagver %d
doband -1
aparm 0
aparm(1)=1
aparm(2)=1
aparm(5)=-180
aparm(6)=180
aparm(7)=1
aparm(8)=1
aparm(9)=1
codetype 'a&p'
nplot 0
bparm 0
dotv -1
inp

""" % (aips_id, aips_filename, flagver)

    aips_script = possm
    
    for antenna in antennas:
        
        for pol in stokes:

            for scan, timerang in zip(data['scan'],
                                      data['timerang']):
                
                command = """
antenna %i
stokes '%s'
outtext 'OUTPUT:%s-%s-%s-%s.txt
timerang %s
go
wait

""" % (antenna_number[antenna], pol, aips_filename, antenna, pol, scan, timerang)

                aips_script += command


    # Delete the file from AIPS, to leave the catalog clean
    command = """
inname '%s'; inclass 'uvdata'; zap
""" % (aips_filename)

    aips_script += command


    # Save AIPS script to text file
    output_filename = ('%sp.%s' % (aips_filename, aips_ehexid)).upper()
    fileObj = open(output_filename, 'w')
    fileObj.writelines(aips_script)
    fileObj.close()

    # Notify the user of the script name
    print 'Creating the AIPS run file "%s"\n' % (output_filename)

    return output_filename


#-------------------------------------------------------------------------------
# Run
#-------------------------------------------------------------------------------

for name_prefix in name_prefixes:

    aips_filename = name_prefix

    filename = 'output/%s_listr_scan.txt' % (aips_filename)
    data = parse_listr_scan_list(filename)

    # Write AIPS run file
    output_filename = write_aips_script(data,
                                        aips_filename=aips_filename,
                                        antennas=antennas,
                                        stokes=stokes,
                                        aips_id=aips_id,
                                        aips_ehexid=aips_ehexid,
                                        flagver=flagver)

    # Run the AIPS script from the command line
    # Define input and output AIPS folder
    os.putenv('INPUT', aips_input)
    os.putenv('OUTPUT', aips_output)
    
    aips_command = 'aips notv pr=1 > %s-aipslog.txt < %s' % (output_filename,
                                                             output_filename) 
    os.system(aips_command)