#!/usr/bin/env python
# Permute indices on Fortran arrays

# Written by @gradylemoine

import re, sys

def permute_file(infile, outfile, permute_list, write_header=True):
    '''permute_file(INFILE, OUTFILE, PERMUTE_LIST):
Permute array indices in a Fortran source file.  INFILE is the name
of the inputfile, and OUTFILE is the name of the output file.  It's
an error if INFILE == OUTFILE.  See permute_indices for documentation
on PERMUTE_LIST.'''

    if (infile == outfile):
        print('Identical filenames specified for input and output; not doing anything')
        return

    f = open(infile, 'r')    # No error checking
    code_in = f.read()
    f.close()

    err_list, code_out = permute_indices(code_in, permute_list)

    f = open(outfile, 'w')    # No error checking
    if (write_header):
        f.write("! *** DO NOT EDIT THIS FILE ***\n")
        f.write("! This file was generated automatically from %s using\n"%infile)
        f.write("! the permute.py Python script to permute the array\n")
        f.write("! indices of the following variables:\n")
        f.write("! %s\n"%(', '.join([str(x[0]) for x in permute_list])))
        f.write("! If this file is wrong, edit the parent file (for\n")
        f.write("! algorithmic errors) or the permute.py script\n")
        f.write("! (for array permutation errors)\n\n")
    f.write(code_out)
    f.close()
    for lineno, bad_vars, line in err_list:
        print(('Line %4d: confused about '%lineno)+', '.join(bad_vars)+' in: '+line)
    

def permute_indices(code, permute_list):
    '''permute_indices(CODE, PERMUTE_LIST):
Permute array indices on a single line.  CODE is the Fortran code to
be operated on; PERMUTE_LIST is a list of tuples (NAME, PLIST), where
NAME is the name of an array to be permuted, and PLIST is a list of
integers saying which indices of the original array go where in the
final list.  (This list is indexed starting from 1, because it's a
list of regex backrefences.)

Returns a tuple (ERR_LIST, CODE_OUT), where ERR_LIST is a list tuples
of the lines that weren't successfully processed, the variables that
weren't successfully processed on them, and the original line, and
CODE_OUT is the modified code.'''

    regex_match = []    # To be used with re.findall()
    regex_replace = []  # To be used with re.subn()
    str_replace = []    # String to replace with in re.subn()
    for name, plist in permute_list:
        matchstr = r'\b'+name+r'\b'
        repstr = (r'\b'+name+r'\s*\('
                  +r','.join([r'\s*([^,\)]+?)\s*']*len(plist))
                  +r'\)')
        regex_match.append(re.compile(matchstr, re.IGNORECASE))
        regex_replace.append(re.compile(repstr, re.IGNORECASE))
        str_replace.append(name+'('+','.join([r'\%d'%n for n in plist])+')')

    iscomment = re.compile('^(c|\s*!)', re.IGNORECASE)

    lines = code.split('\n')
    code_out = ''
    err_list = []
    for i in range(len(lines)):
        bad_vars = []
        line = lines[i]
        #if (not iscomment.match(line)):
        if True:
            
            line0 = line
            for j in range(len(permute_list)):
                nmatch = len(regex_match[j].findall(line))
                line_new, nrep = regex_replace[j].subn(str_replace[j], line)
                if (nrep == nmatch):
                    line = line_new
                else:
                    bad_vars.append(permute_list[j][0])

            if (len(bad_vars) > 0):
                err_list.append((i+1, bad_vars, line0))

        code_out = code_out + line + '\n'

    return (err_list, code_out)


if (__name__ == '__main__'):
    if (len(sys.argv) < 3):
        print('permute.py: need source and destination file specified on command line')
        
    else:
        # assuming 2d:
        plist = [('q',     [3,1,2]),
                 ('aux',   [3,1,2])]

        permute_file(sys.argv[1], sys.argv[2], plist, write_header=False)
