#!/usr/bin/env python

import defapp
from optparse import make_option as o
from os.path  import splitext, basename

import sys

from util import load_csv_file

def split_rows(rows, key_col):
    by_key_col = {}
    key_order  = []

    # hash each row according by its key
    for r in rows:
        key = r[key_col]
        if not key in by_key_col:
            key_order.append(key)
            by_key_col[key] = []
        by_key_col[key].append(r)

    return (by_key_col, key_order)


def group(rows, key_col):
    (by, order) = split_rows(rows, key_col)
    for k in order:
        yield (k, by[k])

def aggregate_column(rows, key_col, aggregate_fun):
    by_key, order = split_rows(rows, key_col)
    for key in order:
        yield aggregate_fun(by_key[key])


def weighted_schedulability(data, key_col, ucap_col):
    def wsched(rows):
        # Every column except for the key_col and the ucap_col
        # is presumed to be schedulability values.
        key   = rows[0][key_col]
        cols  = len(rows[0])
        accum = [0.0 for _ in xrange(cols)]
        norm  = 0.0
        umax  = 0.0
        for r in rows:
            ucap = r[ucap_col]
            umax = max(ucap, umax)
            norm += ucap
            for i, val in enumerate(r):
                if i != key_col and i != ucap_col:
                    accum[i] += ucap * val
        return [(accum[i] / norm) if i != key_col and i != ucap_col
                else (key if i == key_col else umax) for i in xrange(cols)]

    return aggregate_column(data, key_col, wsched)


def print_row(row, out, fmt):
    out.write(", ".join(fmt % d for d in row))
    out.write("\n")

def filter_wsched(fname, key, ucap, out=sys.stdout, fmt="%10.3f"):
    data = load_csv_file(fname, dtype=float)
    for row in weighted_schedulability(data, key, ucap):
        print_row(row, out, fmt)
    out.flush()

def filter_comments(fname, is_comment=lambda line: line and line[0] == '#',
                    out=sys.stdout):
    f = open(fname, 'r')
    for line in f:
        if is_comment(line):
            out.write(line)
    out.flush()
    f.close()

options = [
    # output options
    o('-o', '--output-prefix', action='store', dest='prefix'),
    o('-k', '--key-column', action='store', type='int', dest='keycol'),
    o('-u', '--ucap-column', action='store', type='int', dest='ucol'),
    o(None, '--no-comments', action='store_true', dest='no_comments'),
    o('-w', '--col-width', action='store', type='int', dest='colwidth'),
    o(None, '--split', action='store_true', dest='want_split'),
    o('-f', '--key-format', action='store', dest='key_fmt'),
    ]

defaults = {
    'prefix'      : None,
    'keycol'      : 0,
    'ucol'        : 1,
    'no_comments' : False,
    'colwidth'    : 10,
    'precision'   :  3,
    'indent'      :  2,
    'key_fmt'     : "_key=%d",
    'want_split'  : False,
    }


class WschedFilter(defapp.App):
    def __init__(self):
        defapp.App.__init__(self, options, defaults, no_std_opts=True)
        self.out = sys.stdout


    def filter_wsched(self, fname, out=sys.stdout):
        data = load_csv_file(fname, dtype=float)
        for row in weighted_schedulability(
            data, self.options.keycol, self.options.ucol):
            out.write(" " * self.options.indent)
            out.write(", ".join(self.fmt % d for d in row))
            out.write("\n")
        out.flush()

    def target_name(self, fname, extra=None):
        bname = basename(fname)
        if extra:
            bname, ext = splitext(bname)
            bname = "%s%s%s" % (bname, extra, ext)
        return self.options.prefix + bname

    def filter_file(self, fin, fout=None):
        out = open(fout, 'w') if fout else sys.stdout
        if not self.options.no_comments:
            filter_comments(fin, out=out)
        self.filter_wsched(fin, out=out)
        if fout:
            out.close()

    def write_split(self, fin, fout, rows):
        out = open(fout, 'w')
        if not self.options.no_comments:
            filter_comments(fin, out=out)
        for r in rows:
            print_row(r, out, self.fmt)
        out.close()

    def split_file(self, fname):
        data = load_csv_file(fname, dtype=float)
        for (key, rows) in group(data, self.options.keycol):
            fout = self.target_name(fname, self.options.key_fmt % key)
            print '\t->', fout
            self.write_split(fname, fout, rows)

    def wsched(self, _):
        if self.options.prefix:
            for i, datafile in enumerate(self.args):
                fout = self.target_name(datafile)
                print "[%d/%d] %s -> %s" % \
                    (i + 1, len(self.args), datafile, fout)
                try:
                    self.filter_file(datafile, fout)
                except IOError as ioe:
                    sys.stderr.write("[!!] ")
                    sys.stderr.write(str(ioe))
                    sys.stderr.write("\n")
        else:
            for fname in self.args:
                self.filter_file(fname)

    def split(self, _):
        if self.options.prefix:
            for i, datafile in enumerate(self.args):
                print "[%d/%d] Splitting %s..." % \
                    (i + 1, len(self.args), datafile)
                try:
                    self.split_file(datafile)
                except IOError as ioe:
                    sys.stderr.write("[!!] ")
                    sys.stderr.write(str(ioe))
                    sys.stderr.write("\n")
        else:
            print "split: requires output prefix (-o)."

    def default(self, _):
        self.fmt = "%%%d.%df" % (self.options.colwidth, self.options.precision)
        if self.options.want_split:
            self.split(_)
        else:
            self.wsched(_)

if __name__ == '__main__':
    WschedFilter().launch()