#!/usr/bin/python

""" meanstep.py
 ...
"""

import sys
import os

results = {}

if len(sys.argv) != 4 and len(sys.argv) != 5:
    print "usage: ./meanstep.py <datadir> <matching col 1>,...,<matching col n> <mean col 1>,...,<mean col n> [<outdir>]"
    exit(0)

matchcols = sys.argv[2].split(",")
meancols = sys.argv[3].split(",")

if len(sys.argv) == 5:
    outdir = sys.argv[4]
    writetofile = True
else:
    writetofile = False

for root, dirs, filenames in os.walk(sys.argv[1]):
    for f in filenames:
        if f[-4:] != ".dat":
            continue

        if not f[:-4] in results:
            results[f[:-4]] = {}
            print "Found new fileid: " + f[:-4]

        with open(root + "/" + f) as inf:
            for line in inf:
                if line[0:1] == "#":
                    continue

                linecols = line.split()
                matchkey = ""
                for imatchcol in range(0,len(matchcols)):
                    matchkey = matchkey + linecols[int(matchcols[imatchcol])-1]
                
                if not matchkey in results[f[:-4]]:
                    results[f[:-4]][matchkey] = linecols + [1]
                else:
                    results[f[:-4]][matchkey][ len(results[f[:-4]][matchkey]) - 1 ] = \
                        results[f[:-4]][matchkey][ len(results[f[:-4]][matchkey]) - 1 ] + 1 
         
                    for imeancol in range(0,len(meancols)):
                        results[f[:-4]][matchkey][int(meancols[imeancol])-1] = \
                            float(results[f[:-4]][matchkey][int(meancols[imeancol])-1]) + \
                            float(linecols[int(meancols[imeancol])-1])

for fileid in results:
    print "\n[" + fileid + "]"

    if writetofile:
        outfile = open(outdir + "/" + fileid + "_meanstep.dat", "w+")

    for matchkey in results[fileid]:
        for imeancol in range(0,len(meancols)):
            results[fileid][matchkey][int(meancols[imeancol])-1] = \
                float(results[fileid][matchkey][int(meancols[imeancol])-1]) / \
                results[fileid][matchkey][ len(results[f[:-4]][matchkey]) - 1 ]
        
        for icol in range(0,len(results[fileid][matchkey])):
            sys.stdout.write(str(results[fileid][matchkey][icol]))
            if writetofile:
                outfile.write(str(results[fileid][matchkey][icol]));

            if icol != len(results[fileid][matchkey])-1:
                sys.stdout.write("\t")
                if writetofile:
                    outfile.write("\t")
                                  
        print ""
        if writetofile:
            outfile.write("\n")

    if writetofile:
        outfile.close()