#!/bin/python3

"""dbfullpivot takes an input file with time/key/value pairs, and
   pivots the table into a wide table with one new column per key value.

   For example, if the input was this:

       #fsdb -F s time key value
       1 foo 10
       2 bar 20
       2 foo 30
       3 bar 40

   It would convert this to:

       #fsdb -F s time foo bar
       1 10 0
       2 30 20
       3 0 40

   """

import sys
import argparse
import pyfsdb


def parse_args():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter, description=__doc__
    )

    parser.add_argument(
        "input_file", type=argparse.FileType("r"), nargs="?", default=sys.stdin, help=""
    )

    parser.add_argument(
        "output_file",
        type=argparse.FileType("w"),
        nargs="?",
        default=sys.stdout,
        help="",
    )

    parser.add_argument(
        "-k", "--key-column", default="key", type=str, help="The key to pivot around"
    )

    parser.add_argument(
        "-t",
        "--time-column",
        default="time",
        type=str,
        help="The binned time column to use",
    )

    args = parser.parse_args()
    return args


def main():
    args = parse_args()

    # set up storage structures
    storage = {}
    columns = {}

    # open the input file stream
    input = pyfsdb.Fsdb(file_handle=args.input_file)

    # from the input, get extract column numbers/names
    time_column = input.get_column_number(args.time_column)
    key_column = input.get_column_number(args.key_column)
    column_names = input.column_names

    # for each row, remember each value based on time and key
    for row in input:
        # if the time hasn't been seen before, allocate the sub-structure
        if row[time_column] not in storage:
            storage[row[time_column]] = {}

        for column_num in range(0, len(row)):
            # remember all values of non-time and non-key columns
            if column_num != time_column and column_num != key_column:
                storage[row[time_column]][row[key_column]] = row[column_num]
                # record that we've seen this column before
                columns[row[key_column]] = 1

    # open the output stream, and set it's properties
    out = pyfsdb.Fsdb(out_file_handle=args.output_file)

    # the output columns will be a merge of the time column, and
    # previously seen key-index values.
    output_columns = [args.time_column]
    output_columns.extend(columns.keys())
    out.out_column_names = output_columns

    # Output all data, grouped by time_key
    for time_key in storage:

        # create a row containing a column for every seen key
        row = [time_key]
        for column in columns:
            if column not in storage[time_key] or storage[time_key][column] == "":
                row.append("0")
            else:
                row.append(storage[time_key][column])

        # write it out
        out.append(row)


if __name__ == "__main__":
    main()
