#!/usr/bin/env python

# file: $(ISIP)/exp/tuh_eeg/exp_0022/scripts/gen_feats_energy.py
#
# revision history:
#  20150616 (JP): completed cleaning up the code
#  20150616 (AM): cleaned up and documented the code
#  20140601 (AH): initial version
#
# usage:
#  gen_feats_energy.py -i input.htk -o output.htk
#
# arguments:
#  -i (--input):  .htk formatted feature file
#  -o (--output): .htk formatted feature file with energy
#  -w (--window): window length in frames (default = 4)
#
# This script computes a differential energy feature and adds it to
# a feature vector. This differential energy feature is defined as the 
# difference between the maximum and minimum energy in a window.
#------------------------------------------------------------------------------

# import required modules:
#  note that the path to the module htkmfc must be included in the
#  PYTHONPATH environment variable.
#
import os
import sys
import getopt
import htkmfc
import numpy as np

# main: this is the main function of this Python
#
def main(argv):

    # declare parameters for the input and output filenames
    #
    ifile_a = ""
    ofile_a = ""

    # declare a variable for the window length:
    #  (1) note that the window extends from [-wlen,+wlen].
    #  (2) the default value is 4.
    #
    wlen_a = 4

    # define the command line options
    #
    try:
        opts, args = getopt.getopt(
            argv, "hi:o:w:", ["input=", "output=", "window="])

    # error handling for command line options
    #
    except getopt.GetoptError:
        print("*> %s: the option does not exist", sys.argv[0]);
        sys.exit(-1)

    # parse the command line arguments
    #
    for opt, arg in opts:

        # option: help
        #
        if opt == '-h':
            print 'Usage: gen_feats_energy.py -i input.htk -o output.htk'
            sys.exit()

        # the input argument which is the full path of an htk file
        #
        elif opt in ("-i", "--input"):
            ifile_a = arg

        # the output argument which is the full path of an htk file
        #
        elif opt in ("-o", "--output"):
            ofile_a = arg

        # window length argument. It is optional. The default value is 4.
        #
        elif opt in ("-w", "--window"):
            wlen_a = int(arg)

    # open the input file and read the data
    #
    inp_mfc = htkmfc.open(ifile_a)
    data = inp_mfc.getall()

    # extract N, the number of frames, and D, the number of features
    #
    (N, D) = data.shape

    # create a matrix of zeros with N rows and 1 columns
    #
    z = np.zeros(shape=(N, 1))

    # create a new matrix with the data concatenated with z:
    #  note that this generates data_out with Nx(D+1) with the last 
    #  column set to 0.
    #
    data_out = np.append(data, z, 1)

    # loop over all frames to calculate the differential energy
    #
    for n in range(N):

        # create a matrix with the length of 2*wlen_a+1, which
        # is the total size of the history used to compute the min and max
        #
        w = np.zeros(shape=(2 * wlen_a + 1, 1))

        # set the elements of matrix of w to be equal to the energy (first
        # feature) at time n.
        #
        k = 0
        w[k] = data[n, 0]
        k = k + 1

        # loop over all elements in the submatrix
        #
        for t0 in range(wlen_a):

            # check boundary conditions so we don't exceed
            # the edges of the signal
            #
            t = t0 + 1
            if n + t > N - 1:
                ind2 = N - 1
            else:
                ind2 = n + t
            if n - t < 0:
                ind1 = 0
            else:
                ind1 = n - t

            # fill the w vector with information that occurs inside window
            #
            w[k] = data[ind1, 0]
            k = k + 1
            w[k] = data[ind2, 0]
            k = k + 1

            # end of loop
            #
            pass

        # compute the difference of the max and min (an estimate of the
        # derivative, and set the last feature coefficient for data_out
        #
        data_out[n, D:] = np.max(w) - np.min(w)

        # end of loop
        #
        pass

    # obtain the shape of data_out
    #
    (N, Dout) = data_out.shape

    # create a new .htk file that contains the new information:
    #  ofile_a will be generated in the same place as the input file
    #
    out_mfc = htkmfc.HTKFeat_write(
        ofile_a, inp_mfc.nSamples, Dout,
        inp_mfc.sampPeriod, inp_mfc.paramKind)
    out_mfc.writeall(data_out)

    # end of loop
    #
    pass

# begin gracefully
#
if __name__ == "__main__":
    main(sys.argv[1:])

#
# end of file