#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2012-2013 Tobias Klauser <tklauser@distanz.ch>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.

import os, sys
import getopt
import csv
import numpy as np

DEFAULT_NAME = 'control'
DEFAULT_MAXTIME = 100.0

def usage():
    print """usage: %s [OPTION...] CSV-FILE STO-FILE

  -f    force overwrite of existing files
  -n    name of the sto file (default: %s)
  -T N  scale time such that it goes from 0.0 to N
  -h    show this help and exit""" % (DEFAULT_NAME, os.path.basename(sys.argv[0]))

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "fn:hT:")
    except getopt.GetoptError, err:
        print(str(err))
        usage()
        sys.exit(-1)

    if len(args) < 2:
        usage()
        sys.exit(-1)

    overwrite = False
    name = DEFAULT_NAME
    maxtime = DEFAULT_MAXTIME
    for o, a in opts:
        if o == '-f':
            overwrite = True
        elif o == '-n':
            name = a
        elif o == '-T':
            maxtime = float(a)
        elif o == '-h':
            usage()
            sys.exit(0)
        else:
            assert False, "unhandled option"

    fs_csv = args[0:-1]
    for f_csv in fs_csv:
        if not os.path.exists(f_csv):
            print "Error: CSV file %s does not exist" % f_csv
            sys.exit(-1)
    f_sto = args[-1]
    if not overwrite and os.path.exists(f_sto):
        print "Error: STO file %s already exists" % f_sto
        sys.exit(-1)

    # determine number of columns and rows
    tot_nRows = tot_nCols = 0
    for f_csv in fs_csv:
        fd = open(f_csv, 'r')
        nRows = nColumns = 0
        for line in fd:
            if nRows == 0:
                nCols = len(line.split(','))
            nRows += 1
        if tot_nRows == 0:
            tot_nRows = nRows
        elif tot_nRows != nRows:
            print "Error: Number of rows in CSV files do not match"
            sys.exit(-1)
        if tot_nCols == 0:
            tot_nCols = nCols
        elif tot_nCols != nCols:
            print "Error: Number of columns in CSV files do not match"
            sys.exit(-1)

        fd.close()

    print "Writing sto file %s with name %s" % (f_sto, name)
    if maxtime != DEFAULT_MAXTIME:
        print "Scaling to maxtime %f" % maxtime

    fd_sto = open(f_sto, 'w')
    fd_sto.write(name + "\n")
    fd_sto.write("version=1\n")
    fd_sto.write("nRows=" + str(tot_nRows) + "\n")
    fd_sto.write("nColumns=" + str(tot_nCols) + "\n")
    fd_sto.write("endheader\n")

    fd_sto.write("time")
    for f_csv in fs_csv:
        colname = os.path.splitext(os.path.basename(f_csv))[0]
        print "[+] Adding muscle %s..." % colname
        fd_sto.write("\t" + colname)
    fd_sto.write("\n")

    fds_csv = [ open(f_csv, 'r') for f_csv in fs_csv ]
    crs = [ csv.reader(fd_csv, delimiter=',') for fd_csv in fds_csv ]

    times = np.zeros((tot_nRows - 1, len(fds_csv)))
    vals = np.zeros((tot_nRows - 1, len(fds_csv)))
    for i, cr in enumerate(crs):
        cr.next()   # skip header line
        for j, row in enumerate(cr):
            t,y = map(float, row)
            if maxtime != DEFAULT_MAXTIME:
                # Assume default 0-100 time in all input files for now
                t = (t / 100.0) * maxtime
            times[j,i] = t
            vals[j,i] = y

    # XXX: For now assume same time spacing in all CSV files
    for i, row in enumerate(times):
        fd_sto.write(str(row[0]))
        for val in vals[i]:
            fd_sto.write("\t" + str(val))
        fd_sto.write("\n");

    for fd_csv in fds_csv:
        fd_csv.close()
    fd_sto.close()

if __name__ == '__main__':
    main()