#!/usr/bin/env python
from scipy.signal import butter, lfilter
from nptdms import TdmsFile
import sys, string, pickle
import matplotlib.pyplot as plt
import numpy as np
from scipy import interpolate
def blackman_smooth(x, window_len):
    s=np.r_[2*x[0]-x[window_len:1:-1], x, 2*x[-1]-x[-1:-window_len:-1]]
    w = np.blackman(window_len)
    y = np.convolve(w/w.sum(),s,mode='same')
    return y[window_len-1:-window_len+1] 

def moving_average(interval, window_size):
    window = np.ones(int(window_size),float)/float(window_size)
    return np.convolve(interval,window,'same')

def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a

def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y
def main(filename,min_time,max_time):
    lowcut =1.0 
    highcut = 200.0
    tdms_file = TdmsFile(filename)
    min_time = float(min_time)
    max_time = float(max_time)
    groups = tdms_file.groups()
    dat = {}
    time = {}
    for g in groups:
        dat[g] = {}
        time[g] = {}
        channels = ['LoadCell', 'LVDT']
        for c in channels:
            channel = tdms_file.object(g,c)
            dat[g][c] = channel.data
            time[g][c] = channel.time_track()
            fs = 1.0/(time[g][c][1]-time[g][c][0])
            min_ind = next(x[0] for x in enumerate(time[g][c]) if x[1] > min_time)
            max_ind =  next(x[0] for x in enumerate(time[g][c]) if x[1] > max_time)
            dmy = dat[g][c][min_ind:max_ind]
            '''
            ps = np.abs(np.fft.fft(dmy))**2
            time_step = time[g][c][1] - time[g][c][0]
            freqs = np.fft.fftfreq(dmy.size,time_step)
            idx = np.argsort(freqs)
            plt.plot(freqs[idx],ps[idx])
            plt.show()
            plt.clf()
            '''
            dmy = butter_bandpass_filter(np.array(dmy),lowcut,highcut,fs,order=6)
            time[g][c] = time[g][c][min_ind:max_ind]
            plt.plot(time[g][c],dmy)
            plt.show()
            dat[g][c] = blackman_smooth(dmy,10)
            # interpolate the data
            ti = np.arange(min(time[g][c]),max(time[g][c]),0.1)
            f = interpolate.interp1d(time[g][c],dat[g][c])
            dat[g][c] = f(ti)
            time[g][c] = ti
            #plt.plot(time[g][c],dat[g][c],'b')
            '''
            plt.plot(time[g][c],dat[g][c],'r')
            plt.show()
            '''
    output = open(string.replace(filename,'.tdms','.pkl'),'wb')
    dat = {'Time': time, 'Load': dat}
    pickle.dump(dat,output)
            
if __name__ == "__main__":
    main(sys.argv[-3],sys.argv[-2],sys.argv[-1])