#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# random-muscle-activation.py -- Generate random muscle activation data in CSV
#
# 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 os
import sys
import getopt
import numpy as np
import matplotlib.pyplot as plt
import plotcsv

DEFAULT_N = 1000
DEFAULT_XMIN = 0.0
DEFAULT_XMAX = 100.0
DEFAULT_YMIN = 0.0
DEFAULT_YMAX = 50.0
DEFAULT_YDMIN = -0.5
DEFAULT_YDMAX = 0.5

WINDOWS = [ 'flat', 'hanning', 'hamming', 'bartlett', 'blackman' ]
WINDOW_LEN = 11
DEFAULT_WINDOW = 'blackman'

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

Generate random muscle activation data in CSV.

options:
    -f    force overwrite of existing file
    -n N  set number of data points to N (default: {})
    -p    plot data after generating it
    -s    smoothen the data using blackman window function
    -i N  set minimum initial y value to N (default: {})
    -I N  set maximum initial y value to N (default: {})
    -x N  set minimum x value to N (default: {})
    -X N  set maximum x value to N (default: {})
    -y N  set minimum y value to N (default: {})
    -Y N  set maximum y value to N (default: {})
    -d N  set minimum y change value to N (default: {})
    -D N  set maximum y change value to N (default: {})""".format(
        os.path.basename(sys.argv[0]), DEFAULT_N,
        DEFAULT_YMIN, DEFAULT_YMAX,
        DEFAULT_XMIN, DEFAULT_XMAX,
        DEFAULT_YMIN, DEFAULT_YMAX,
        DEFAULT_YDMIN, DEFAULT_YDMAX))

def generate_activation_data(N, xmin, xmax, ymin, ymax, ydmin, ydmax, imin, imax):
    X = np.array([ ((x / float(N)) * (xmax - xmin)) + xmin for x in range(0, N) ])
    Y = np.zeros(N)
    # Y = np.array([ ((y / float(N)) * (ymax - ymin)) + ymin for y in range(0, N) ])

    Y[0] = (imax - imin) * np.random.random_sample() + imin
    for y in range(1, N):
        Y[y] = Y[y-1] + ((ydmax - ydmin) * np.random.random_sample() + ydmin)
        if Y[y] < ymin:
            Y[y] = ymin
        if Y[y] > ymax:
            Y[y] = ymax

    act = np.array((X, Y)).transpose()

    return act

# based on scipy.org/Cookbook/SignalSmooth
def smoothen(x, window=DEFAULT_WINDOW, window_len=WINDOW_LEN):
    """
    smoothen signal using a window with requested size

    input:
      x:          the input signal
      window:     type of window from 'flat', 'hanning', 'hamming', 'bartlett',
                  'blackman'; the 'flat' method will apply a moving average
      window_len: dimension of smoothing window, must be an odd integer

    output:
      smoothened signal
    """

    if x.ndim != 1:
        raise ValueError, "only 1-dimensional arrays can be smoothened"
    if x.size < window_len:
        raise ValueError, "input vector needs to be bigger than window size"
    if type(window_len) != int or window_len % 2 != 1:
        raise ValueError, "window length must be an odd integer"

    if window_len < 3:
        return x
    if not window in WINDOWS:
        raise ValueError, "window type must be one of " + ", ".join([ "'" + w + "'" for w in WINDOWS ])

    s = np.r_[x[window_len-1:0:-1],x,x[-1:-window_len:-1]]
    if window == 'flat':
        w = np.ones(window_len, 'd')
    else:
        w = eval('np.' + window + '(window_len)')

    return np.convolve(w / w.sum(), s, mode='valid')

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "fn:psi:I:x:X:y:Y:d:D:h")
    except getopt.GetoptError as err:
        print(str(err))
        usage()
        sys.exit(-1)

    do_overwrite = False
    do_plot = False
    do_smooth = False
    N = DEFAULT_N
    imin = DEFAULT_YMIN
    imax = DEFAULT_YMAX
    xmin = DEFAULT_XMIN
    xmax = DEFAULT_XMAX
    ymin = DEFAULT_YMIN
    ymax = DEFAULT_YMAX
    ydmin = DEFAULT_YDMIN
    ydmax = DEFAULT_YDMAX
    for o, a in opts:
        if o == '-f':
            do_overwrite = True
        elif o == '-n':
            try:
                N = int(a)
            except ValueError:
                print("Error: number of data points must be an integer > 0")
                usage()
                sys.exit(-1)
        elif o == '-p':
            do_plot = True
        elif o == '-s':
            do_smooth = True
        elif o == '-i':
            try:
                imin = float(a)
            except ValueError:
                print("Error: imin must be a float")
                usage()
                sys.exit(-1)
        elif o == '-I':
            try:
                imax = float(a)
            except ValueError:
                print("Error: imax must be a float")
                usage()
                sys.exit(-1)
        elif o == '-x':
            try:
                xmin = float(a)
            except ValueError:
                print("Error: xmin must be a float")
                usage()
                sys.exit(-1)
        elif o == '-X':
            try:
                xmax = float(a)
            except ValueError:
                print("Error: xmax must be a float")
                usage()
                sys.exit(-1)
        elif o == '-y':
            try:
                ymin = float(a)
            except ValueError:
                print("Error: ymin must be a float")
                usage()
                sys.exit(-1)
            if imin == DEFAULT_YMIN:
                imin = ymin
        elif o == '-Y':
            try:
                ymax = float(a)
            except ValueError:
                print("Error: ymax must be a float")
                usage()
                sys.exit(-1)
            if imax == DEFAULT_YMAX:
                imax = ymax
        elif o == '-d':
            try:
                ydmin = float(a)
            except ValueError:
                print("Error: ydmin must be a float")
                usage()
                sys.exit(-1)
        elif o == '-D':
            try:
                ydmax = float(a)
            except ValueError:
                print("Error: ydmax must be a float")
                usage()
                sys.exit(-1)
        elif o == '-h':
            usage()
            sys.exit(0)
        else:
            assert False, "unhandled option"

    if N <= 0:
        print("Error: invalid number of data points ({}), must be > 0".format(N))
        sys.exit(-1)

    if not xmin < xmax:
        print("Error: xmin must be smaller than xmax")
        sys.exit(-1)

    if not ymin < ymax:
        print("Error: ymin must be smaller than ymax")
        sys.exit(-1)

    if not ydmin <= ydmax:
        print("Error: ydmin must be smaller than or equal to ydmax")
        sys.exit(-1)

    if not imin < imax:
        print("Error: imin must be smaller than imax")
        sys.exit(-1)

    if imin < ymin or imin > ymax:
        print("Error: imin must be between ymin and ymax ({} <= {} <= {})".format(ymin, imin, ymax))
        sys.exit(-1)

    if imax < ymin or imax > ymax:
        print("Error: imax must be between ymin and ymax ({} <= {} <= {})".format(ymin, imax, ymax))
        sys.exit(-1)

    if len(args) < 1:
        print("Error: no output file(s) specified")
        usage()
        sys.exit(-1)

    for a in args:
        if not do_overwrite and os.path.exists(a):
            print("Error: output file {} already exists, use -f to overwrite".format(a))
            sys.exit(-1)

    print("Generating {} files with {} data points each...".format(len(args), N))
    for a in args:
        with open(a, 'w') as f:
            act = generate_activation_data(N, xmin, xmax, ymin, ymax, ydmin, ydmax, imin, imax)
            if do_smooth:
                all_in_range = False
                n = 0
                while not all_in_range:
                    print("Smooth pass " + str(n))
                    sact = smoothen(act[:,1], DEFAULT_WINDOW, WINDOW_LEN)
                    sact = sact[(WINDOW_LEN - 1) / 2:-(WINDOW_LEN - 1) / 2]
                    if sact.min() >= ymin and sact.max() <= ymax:
                        all_in_range = True
                    n += 1
                plt.plot(act[:,1])
                plt.plot(sact)
                plt.legend(['original', 'smoothened'])
                plt.show()
                act[:,1] = sact
            # XXX: temp
            #    plt.plot(act[:,1])
            #    l = ['original']
            #    for w in WINDOWS:
            #        sact = smoothen(act[:,1], w, WINDOW_LEN)
            #        plt.plot(sact[(WINDOW_LEN-1)/2:-(WINDOW_LEN-1)/2])
            #        l.append(w)
            #    plt.legend(l)
            #    plt.show()

            f.write('x,y\n')
            for x, y in act:
                f.write('{},{}\n'.format(x, y))
            if do_plot:
                print("Plotting data for {}, please close window to continue...".format(a))
                plotcsv.plot(act[:,0], act[:,1], xmax, ymax, 'x', 'y', a, False, False)

    sys.exit(0)

if __name__ == '__main__':
    main()