
from scipy import signal
import math, numpy as np
from matplotlib import pyplot

import roslib; roslib.load_manifest('sandbox_tapo_darpa_m3')
import rospy
import hrl_lib.matplotlib_util as mpu
import hrl_lib.util as ut

if __name__ == '__main__':

    ### Frequency Analysis for Experimental Data

    ta_fs = ut.load_pickle('time_varying_data_fixed_styrofoam.pkl')
    ta_fp = ut.load_pickle('time_varying_data_fixed_pillow.pkl')
    ta_ms = ut.load_pickle('time_varying_data_movable_styrofoam.pkl')
    ta_mp = ut.load_pickle('time_varying_data_movable_pillow.pkl')

    ### FOR Fixed_Styrofoam (Total_Forces)
    
    # some constants
    samp_rate = 100
    sim_time = ta_mp[-1,0]
    nsamps = int(samp_rate*sim_time)
    #cutoff_freq = 5 #for HPF ??
    cutoff_freq = 0.5 #for LPF ??
    
    fig = pyplot.figure()

    # generate input signal
    t = ta_mp[:,0]
    x = ta_mp[:,1]
    time_dom = fig.add_subplot(232)
    pyplot.plot(t, x)
    pyplot.title('Filter Input - Time Domain')
    pyplot.grid(True)

    # input signal spectrum
    xfreq = np.fft.fft(x)
    fft_freqs = np.fft.fftfreq(nsamps, d=1./samp_rate)
    fig.add_subplot(233)
    pyplot.loglog(fft_freqs[0:nsamps/2], np.abs(xfreq)[0:nsamps/2])
    pyplot.title('Filter Input - Frequency Domain')
    pyplot.grid(True)

    # design filter (Low-Pass Butterworth)
    norm_pass = 2*math.pi*cutoff_freq/samp_rate
    norm_stop = 1.5*norm_pass
    (N, Wn) = signal.buttord(wp=norm_pass, ws=norm_stop, gpass=2, gstop=30, analog=0)
    (b, a) = signal.butter(N, Wn, btype='low', analog=0, output='ba')
    b *= 1e3
    print("b="+str(b)+", a="+str(a))

    # design filter (High-Pass Butterworth)
    #norm_stop = 2*math.pi*cutoff_freq/samp_rate
    #norm_pass = 1.5*norm_stop
    #(N, Wn) = signal.buttord(wp=norm_pass, ws=norm_stop, gpass=2, gstop=30, analog=0)
    #(b, a) = signal.butter(N, Wn, btype='high', analog=0, output='ba')
    #b *= 1e3
    #print("b="+str(b)+", a="+str(a))

    # design filter (Low-Pass FIR)
    #b = signal.firwin(5,cutoff_freq)
    #a = 1

    # design filter (High-Pass FIR) (This feature is enabled only in the latest version of SciPy(v0.10.dev)
    #b = signal.firwin(5,cutoff_freq,pass_zero=False)
    #a = 1
 
    # filter frequency response
    (w, h) = signal.freqz(b, a)
    fig.add_subplot(131)
    pyplot.loglog(w, np.abs(h))
    pyplot.title('Filter Frequency Response')
    pyplot.text(2e-3, 1e-5, str(N)+"-th order Butterworth filter")
    pyplot.grid(True)

    # filtered output
    y = signal.lfilter(b, a, x)
    fig.add_subplot(235)
    pyplot.plot(t, y)
    pyplot.title('Filter output - Time Domain')
    pyplot.grid(True)

    # output spectrum
    yfreq = np.fft.fft(y)
    fig.add_subplot(236)
    pyplot.loglog(fft_freqs[0:nsamps/2], np.abs(yfreq)[0:nsamps/2])
    pyplot.title('Filter Output - Frequency Domain')
    pyplot.grid(True)

    pyplot.show()

    





