#!/usr/bin/env python

"""
nbody_fix_ttree_assignments.py
       
This is an ugly little script which was not intended to be run by end users.

Typical usage:

nbody_fix_ttree_assignments.py "angles" new_Angles.template \
  < ttree_assignments.txt > ttree_assigmnents_new.txt

What it does:

In this example, this program extracts the first column from 
"new_Angles.template", and appends to it the first column from 
the lines in ttree_assignments.txt containing "@angle:".
Then it adds a second column which is just a sequence of integers
counting upwards.
Finally it inserts this 2-column text into the appropriate place in a 
ttree_assignments.txt file, replacing the original @angle variables
with the new ones (followed by the renumbered original).


   AWK/GREP equivalent 
This program is roughly equivalent to the following lines of awk/grep:

awk 'BEGIN{i=-1} {if(substr($0,0,8)=="$/angle") i=NR; if (i==-1){print $0}}'\
    < ttree_assignments.txt > ttree_assignments_new.txt

awk '{print $1}' < new_Angles.template > Angles_column1.txt
grep '$/angle:' ttree_assignments.txt | awk '{print $1}' >> Angles_column1.txt
awk '{print $1 "  " NR}' <  Angles_column1.txt >> ttree_assignments_new.txt

awk 'BEGIN{found=0;passed=0} {if(substr($0,0,8)=="$/angle") found=1; 
                              else {if (found) {passed=1}} 
                              if (passed) print $0}' \
          < ttree_assignments.txt >> ttree_assignments_new.txt

I wrote this python script (instead of using awk) just to handle quoted stings
(and strings with other fancy characters and escape sequences).

"""

import sys
from ttree_lex import SplitQuotedString, EscCharStrToChar, SafelyEncodeString, InputError

try:
    if (len(sys.argv) != 3):
        raise InputError('Error running  \"'+__file__+'\"\n'
                         '   Wrong number of arguments.\n'
                         '   (This is likely a programmer error.\n'
                         '    This script was not intended to be run by end users.)\n')

    cat_name = sys.argv[1] 
    f = open(sys.argv[2])
    lines_generated = f.readlines()
    f.close()

    # Selections are simply lists of 2-tuples (pairs) 
    #f = open('ttree_assignments.txt','r')
    #lines_bindings = f.readlines()
    #f.close()
    lines_bindings = sys.stdin.readlines()

    # Figure out which lines in the 'ttree_assignments.txt' file
    # contain the variables of the type you are looking for.
    # Make note of the relevant line numbers
    i_preexisting_begin = -1
    i_preexisting_end = -1
    in_section = False
    possible_cat_names = set(['$'+cat_name, '$/'+cat_name, '${'+cat_name, '${/'+cat_name])

    preexisting_interaction_list = []
    for i in range(0, len(lines_bindings)):
        line = lines_bindings[i].strip()
        tokens = SplitQuotedString(line)  #strip comments, handle quotes
        if len(tokens) == 2:
            before_colon = tokens[0].split(':')[0]
            if before_colon in possible_cat_names:
                if i_preexisting_begin == -1:
                    i_preexisting_begin = i
                    in_section = True
            else:
                if in_section:
                    i_preexisting_end = i
                in_section = False

    if i_preexisting_end == -1:
        i_preexisting_end = len(lines_bindings)

    if i_preexisting_begin == -1:
        for line in lines_bindings:
            sys.stdout.write(line)
    else:
        # write out all the lines in the original file up until the point where
        # the variables in the category we are looking for were encountered
        for i in range(0, i_preexisting_begin):
            sys.stdout.write(lines_bindings[i])

    sys.stderr.write('  (adding new lines)\n')

    # Now add some new lines (2-column format).
    # As with any ttree_assignment.txt file:
    #   The first column has our generated variable names
    #   The second column has the counter assigned to that variable
    new_counter = 1
    for line_orig in lines_generated:
        line = line_orig.strip()
        if len(line) > 0:
            tokens = SplitQuotedString(line)  #strip comments, handle quotes
            sys.stdout.write(tokens[0]+'  '+str(new_counter)+'\n')
            new_counter += 1

    sys.stderr.write('  (adding pre-exisiting lines)\n')
    if i_preexisting_begin != -1:
        # Append the original pre-existing interactions of that type, but assign
        # them to higher numbers.  (Hopefully this helps to make sure that these
        # assignments will override any of the automatic/generated assignments.)
        # As with any ttree_assignment.txt file:
        #   The first column has our generated variable names
        #   The second column has the counter assigned to that variable

        #sys.stderr.write('  i_preexisting_begin='+
        #                 str(i_preexisting_begin)+
        #                 ' i_preexisting_end='+str(i_preexisting_end)+'\n')

        for i in range(i_preexisting_begin, i_preexisting_end):
            line = lines_bindings[i].strip()
            tokens = SplitQuotedString(line)  #strip comments, handle quotes
            if len(tokens) == 2:
                sys.stdout.write(tokens[0]+'  '+str(new_counter)+'\n')
                new_counter += 1

        #sys.stderr.write('  (writing pre-exisiting lines)\n')

        # write out all the lines in the original file after this point.
        for i in range(i_preexisting_end, len(lines_bindings)):
            sys.stdout.write(lines_bindings[i])

    sys.exit(0)


except (ValueError, InputError) as err:
    sys.stderr.write('\n'+str(err)+'\n')
    sys.exit(-1)

