#!python

"""Simple script to renumber by-site output of ``phydms``.

Written by Jesse Bloom."""


import sys
import os
import re
import time
import math
import glob
import phydmslib.parsearguments
import phydmslib.file_io


def main():
    """Main body of script."""

    # Parse command line arguments
    parser = phydmslib.parsearguments.PhyDMSRenumberParser()
    args = vars(parser.parse_args())
    prog = parser.prog

    # print some basic information
    print('\nBeginning execution of %s in directory %s at time %s\n' % (prog, os.getcwd(), time.asctime()))
    print("%s\n" % phydmslib.file_io.Versions())
    print('Parsed the following command-line arguments:\n%s\n' % '\n'.join(['\t%s = %s' % tup for tup in args.iteritems()]))

    print("Reading the new numbering system from %s\n" % args['renumberfile'])
    renumber_d = {}
    with open(args['renumberfile']) as f:
        for line in f:
            if line[0] != '#' and not line.isspace():
                entries = line.split()
                assert len(entries) == 2, "Did not find exactly two entries on the following line in %s:\n%s" % (args['renumberfile'], line)
                (old_i, new_i) = (entries[0].strip(), entries[1].strip())
                if old_i in renumber_d:
                    raise ValueError("Duplicate site number for %d in first column of %s" % (old_i, args['renumberfile']))
                else:
                    renumber_d[old_i] = new_i
    assert len(renumber_d) == len(set(renumber_d.values())) + max(0, len([r for r in renumber_d.values() if r.lower() == 'none']) - 1), "Duplicate new site numbers in %s" % args['renumberfile']

    # identify files to renumber
    filesuffixes = ['_omegabysite.txt', '_stringencybysite.txt', '_diffprefsbysite.txt', '_diffprefsbysite_sumabs.txt']
    filestorenumber = []
    for outprefix in args['outprefixes']:
        files = []
        for suffix in filesuffixes:
            if re.search('%s$' % suffix.replace('.', '\.'), outprefix):
                files.append(outprefix)
            files += glob.glob('%s%s' % (outprefix, suffix))
        filestorenumber += files
        if not files:
            raise ValueError("Failed to find any files with one of the expected suffixes (%s) for outprefix %s" % (', '.join(filesuffixes), outprefix))

    # renumber 
    for oldfile in filestorenumber:
        (dirname, fname) = os.path.split(oldfile)
        if dirname:
            newfile = '%s/%s_%s' % (dirname, args['renumberedprefix'], fname)
        else:
            newfile = '%s_%s' % (args['renumberedprefix'], fname)
        print("Renumbering %s to create %s" % (oldfile, newfile))
        newlines = []
        foundsites = set([])
        with open(oldfile) as f:
            for line in f:
                if line[0] == '#' or line.isspace():
                    newlines.append(line)
                else:
                    (site, remainder) = line.split(None, 1)
                    if site in renumber_d:
                        if renumber_d[site].lower() != 'none':
                            newlines.append('%s\t%s' % (renumber_d[site], remainder))
                        foundsites.add(site)
                    else:
                        raise ValueError("Cannot renumber site %s in %s, as no new number was specified for this site. Offending line in %s:\n%s" % (site, oldfile, oldfile, line))
        assert foundsites == set(renumber_d.keys()), "Failed to find all sites specified for renumbering in %s" % oldfile
        with open(newfile, 'w') as f:
            f.write(''.join(newlines))

    print('\nSuccessful completion of %s' % prog)


if __name__ == '__main__':
    main() # run the script
