#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# plotsto.py -- Plot all muscle activation signals in a .sto file
#
# Copyright (C) 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 getopt
import os, sys
import numpy as np
import matplotlib.pyplot as plt

def usage():
    print("""usage: {} [OPTION...] STO-FILE...

Plot all muscle activation signals in a .sto file.

options:

    -h  show this help and exit""".format(os.path.basename(sys.argv[0])))

def read_sto(sto):
    f = open(sto, 'r')

    # .sto header format:
    # control
    # version=N
    # nRows=N
    # nColumns=N
    # endheader
    name = f.readline().strip()
    line = f.readline().strip().split('=')
    if line[0] != 'version' or not line[1].isdigit():
        f.close()
        print("Error: invalid version header line in {}: {}".format(sto, "=".join(line)))
        return None
    version = int(line[1])

    line = f.readline().strip().split('=')
    if line [0] != 'nRows' or not line[1].isdigit():
        f.close()
        print("Error: invalid nRows header line in {}: {}".format(sto, "=".join(line)))
        return None
    nRows = int(line[1])

    line = f.readline().strip().split('=')
    if line[0] != 'nColumns' or not line[1].isdigit():
        f.close()
        print("Error: invalid nColumns header line in {}: {}".format(sto, "=".join(line)))
        return None
    nColumns = int(line[1])

    endheader = f.readline().strip()
    if endheader != 'endheader':
        f.close()
        print("Error: header not properly terminated by 'endheader' in {}".format(sto))
        return None

    print("name: " + name)
    print("version: " + str(version))
    print("nRows: " + str(nRows))
    print("nColumns: " + str(nColumns))

    cols = f.readline().strip().split()
    if len(cols) != nColumns:
        f.close()
        print("Error: nColumns does not match number of columns in file")
        return None

    i = 0
    for line in f:
        i += 1

    f.close()

    if i != nRows:
        print("Error: nRows does not match number of rows in file")
        return None

    act = np.zeros((nRows, nColumns))

    f = open(sto, 'r')
    # skip first 6 lines (header)
    for i in range(6):
        _ = f.readline()
    i = 0
    for line in f:
        data = line.split()
        if len(data) != nColumns:
            f.close()
            print("Error: row {} has less than {} columns".format(i+7, nColumns))
            return None
        for j, val in enumerate(data):
            act[i,j] = float(val)
        i += 1
    f.close()

    return act, cols

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

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

    for o, a in opts:
        if o == '-h':
            usage()
            sys.exit(0)
        else:
            assert False, "unhandled option"

    for sto in args:
        if not os.path.exists(sto):
            print "Error: File %s not found, skipping" % sto
            continue

        act, cols = read_sto(sto)
        # transpose so we can access it more conveniently
        act = act.transpose()
        t, ys = act[0], act[1:]
        l = []
        # see matplotlib.org/examples/pylab_examples/subplots.demo.html
        f, ax = plt.subplots(nrows=(len(cols) - 1) / 2, ncols=2, sharex='col', sharey='row')
        for i, y in enumerate(ys):
            ax[i/2,i%2].plot(t, y)
            l.append(cols[i+1])

        plt.axis([0, t.max(), 0, 1.0], 'equal')
    #    plt.legend(l, bbox_to_anchor=(1.05, 1.), loc=2, borderaxespad=0.)
        plt.show()

if __name__ == '__main__':
    main()