#!/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:

  -m NAMES  plot only columns with names from comma-separated list NAMES
  -h        show this help and exit""".format(os.path.basename(sys.argv[0])))

# .sto header format:
# control
# version=N
# nRows=N
# nColumns=N
# ... (optional header lines)
# endheader
def parse_sto_header(f):
    endheader = False
    name = ''
    nRows = nColumns = 0

    name = f.readline()

    while True:
        line = f.readline().strip()
        if line is None:
            break

        if line == 'endheader':
            endheader = True
            break

        k, v = line.split('=')
        if k == 'nRows':
            nRows = int(v)
        elif k == 'nColumns':
            nColumns = int(v)

    if not endheader:
        print("Error: header not properly terminated by 'endheader' in {}".format(f.name))
        return -1, -1

    return nRows, nColumns

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

    nRows, nColumns = parse_sto_header(f)
    if nRows <= 0 or nColumns <= 0:
        return

    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 = len(f.readlines())
    f.close()

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

    # by default use all columns
    idx = list(range(nColumns))
    _cols = []
    # filter out selected muscles
    if not muscles is None:
        for i, m in enumerate(cols):
            if m in muscles or m == 'time':
                _cols.append(m)
            else:
                idx[i] = None
        cols = _cols

    act = np.zeros((nRows, len(cols)))
    i = 0
    endheader = 0
    f = open(sto, 'r')
    for line in f:
        # skip header and column title lines (we know there's an endheader line)
        if line.strip() == 'endheader':
            endheader = 1
            continue
        if endheader == 0:
            continue
        elif endheader == 1:
            endheader += 1
            continue

        data = line.split()
        if len(data) != nColumns:
            f.close()
            print("Error: row {} has less than {} columns".format(i+7, nColumns))
            return None
        k = 0
        for j, val in enumerate(data):
            if j in idx:
                act[i,k] = float(val.strip())
                k += 1
        i += 1
    f.close()

    return act, cols

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

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

    muscles = None
    for o, a in opts:
        if o == '-m':
            muscles = [ x.strip() for x in a.split(',') ]
        elif 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".format(sto))
            continue

        act, cols = read_sto(sto, muscles)
        # 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)
            ax[i/2,i%2].set_title(cols[i+1], fontsize=9)
            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()