#!/usr/local/bin/python
import sys
import shlex #for smart splitting to preserve enquoted text
from pylab import *
import numpy as np

if len(sys.argv) <= 1:
    print 'USAGE: ', sys.argv[0],'data_file [columnNumber]'
    print '    where data_file has data in columns with an optional one-line header.'
    print '    Data are delimited by white-space. The first column is assumed to be'
    print '    the independent variable'
    print '    The middle-click button is then used to select the FFT range in the'
    print '    plot window.'
    sys.exit()

# stop reading file if change in any of the data values is > THRESHOLD
#THRESHOLD = 9e9
#THRESHOLD = 1e3
#THRESHOLD = 10

data = []
with open(sys.argv[1]) as f:

    #line = f.readline().split()
    line = shlex.split(f.readline())
    print line
    Nvar = len(line) - 1
    try:
        float(line[0]) # first line is data
        varNames = ['x'] + ['y'+str(i) for i in range(1,Nvar+1)]
    except ValueError: # first line is header
        varNames = line
        line = f.readline().split()
    print Nvar,'dependent variables in file :',varNames

    # save first line of data
    for val in line:
        data.append([float(val)])

    # read rest of file
    prevValues = []
    for line in f:
        values = [float(s) for s in line.split()]
#        if len(prevValues) > 0:
#            delta = [abs(v2/v1) for v1,v2 in zip(prevValues,values)]
#            #if any(delta > THRESHOLD):
#            if any([r > THRESHOLD for r in delta]):
#                print 'WARNING: solution blow up detected! (TRESHOLD='+str(THRESHOLD)+')'
#                print '  previous values',prevValues
#                print '  current values ',values
#                break
        # save data
        for col,val in zip(data,values):
            col.append(float(val))
        prevValues = values

if Nvar==0: #single column of data
    Nvar = 1
    N = len(data[0])
    data = [range(1,N+1), data[0]]
    varNames += ['y']

# allow specification of specific columns to plot
#if len(sys.argv[2:]) > 0: 
#    plotlist = [int(var) for var in sys.argv[2:]]
#else: plotlist = range(1,Nvar+1)
#print 'Plotting variables:',[varNames[v] for v in plotlist]
selectedVar = int(sys.argv[2])
print 'Plotting variable:',varNames[selectedVar]

fig = figure()
dataAxes = fig.add_subplot(211)
#for var in plotlist:
#    dataAxes.plot(data[0],data[var],'s-')
dataAxes.plot(data[0],data[selectedVar],'s-')
xlabel('t [s]')
ylabel('signal')

#=========================================================================================
# begin FFT code
#
selectedPoints = []
selectedLines = []
dt = data[0][1] - data[0][0]
print 'dt =',dt
fftAxes = fig.add_subplot(212)
def onClick(event):
    #print 'button=%d, x=%d, y=%d, xdata=%f, ydata=%f' % \
    #    (event.button, event.x, event.y, event.xdata, event.ydata)
    if event.button==2: # middle click
        try:
            print 'Searching for point near %f,%f' % (event.xdata,event.ydata)
        except: return
        idx = np.argmin(np.abs( data[0] - event.xdata ))
        selectedPoints.append(idx)
        print '  - found point at %f,%f' % (data[0][idx],data[selectedVar][idx])

        if len(selectedPoints) > 1: #perform fft
            idx0 = min(selectedPoints[-2:])
            idx1 = max(selectedPoints[-2:])
            print '  - updating FFT for signal between t=',data[0][idx0],data[0][idx1]
            print '    dy =',data[selectedVar][idx1]-data[selectedVar][idx0]
            Nfft = idx1 - idx0
            print '    Nfft =',Nfft
            if Nfft==0: return

            signal = data[selectedVar][idx0:idx1]
            f = np.fft.fftfreq(Nfft,dt)
            F = np.fft.fft(signal)/Nfft
            P = np.abs(F)**2
            imax = np.argmax(P[1:Nfft/2]) + 1
            print '    fmax =',f[imax],'Hz'
            print '    signal mean =',F[0]
            fftAxes.cla()
            #fftAxes.plot(f[:Nfft/2],np.abs(F[:Nfft/2])**2)
            fftAxes.plot(f[1:Nfft/2],P[1:Nfft/2])
            xlabel('f [Hz]')
            ylabel('power')
            draw()

            #for line in selectedLines: line.remove()
            while len(selectedLines) > 0:
                dataAxes.lines.remove( selectedLines.pop(0) )
            selectedLines.append( \
                dataAxes.plot(data[0][selectedPoints[-2]],data[selectedVar][selectedPoints[-2]],\
                    'go', markersize=14)[0] )
            selectedLines.append( \
                dataAxes.plot(data[0][selectedPoints[-1]],data[selectedVar][selectedPoints[-1]],\
                    'r^', markersize=14)[0] )
            selectedLines.append( \
                dataAxes.plot(data[0][idx0:idx1],data[selectedVar][idx0:idx1],\
                    'k-', linewidth=2)[0] )
            draw()

cid = fig.canvas.mpl_connect('button_press_event', onClick)
# end FFT code
#=========================================================================================

show()

