#! /usr/bin/env python
"""Try smoothing data with numpy.

Use cookbook example to start with.


Date: November 28, 2008

Author: Bob Wimmer

"""
from math import pi
from numpy import *
from numpy.random import *
import Gnuplot, Gnuplot.funcutils
import numpy
from numpy.fft import *


def plot(x,y,desc='plot'):
    data = Gnuplot.Data(x,y,using=(1,2)) #this ensures that t is used as x axis
    gg = Gnuplot.Gnuplot()
    string = 'set title "'+str(desc)+'"'
    gg(string)
    gg('set ylabel "y-axis [arb. units]"')
    gg('set xlabel "x-axis [arb. units]"')
    gg('set style data lines')
    gg.plot(data)
    


class timeseries(object):
    """Time series given by x vlaues (e.g. time) and y values (e.g. temperature).
    Perform various operations on time series wuc as:
    - treat/replace missing values
    - smoothing
    - FFT
    - and more to come (maybe)

    Note: requires numpy

    """
    def __init__(self,x,y):
        if len(x) < 2:
            print "this is not a time series, but a single point!"
        if len(y) != len(x):
            print "len(x) != len(y) -- this needs to be fixed!"

        self.x = x
        self.y = y
        self.beginning = x[0]
        self.end = x[-1]
        
    
    def max(self):
        """returns the maximum of x, y"""
        return self.x.max(), self.y.max()

    def min(self):
        """returns the minimum of x, y"""
        return self.x.min(), self.y.min()

    def cut_ind(self, lo,hi):
        """return subset of time series between lo and hi indices."""
        self.xci = self.x[lo:hi]
        self.yci = self.y[lo:hi]
        return self.xci, self.yci

    def cut_x(self, x_lo,x_hi):
        """return subset of time series between x_lo and x_hi values."""
        """find indices of x_lo and x_hi, then call cut_ind"""
        lo = self.x.searchsorted(x_lo)
        hi = self.x.searchsorted(x_hi)
        print lo, self.x[lo], self.x[lo-1]
        self.xcx = self.x[lo:hi]
        self.ycx = self.y[lo:hi]
        return self.xcx, self.ycx

    def clean(self):
        """remove and/or linearly interpolate appropriately missing data values"""
        """This is not yet implemented"""

    def fft(self):
        self.foury = fft(self.y)
        N = len(self.y)
        timestep = self.x[1] - self.x[0]        # if unit=day -> freq unit=cycles/day
        self.fourx = fftfreq(N, d=timestep)
        self.fourx = fftshift(self.fourx)
        self.foury = fftshift(self.foury)
        return self.fourx, self.foury
        
        
    def smooth(self, x,window_len=10,window='hanning'):
        self.s=numpy.r_[2*x[0]-x[window_len:1:-1],x,2*x[-1]-x[-1:-window_len:-1]]
        print(len(self.s))
        if window == 'flat': #moving average
            w=ones(window_len,'d')
        else:
            w=eval('numpy.'+window+'(window_len)')
        y=numpy.convolve(w/w.sum(),self.s,mode='same')
        dy = self.s.copy()
        dy[window_len-1:-window_len+1] = self.s[window_len-1:-window_len+1] - y[window_len-1:-window_len+1]
        return y[window_len-1:-window_len+1], dy[window_len-1:-window_len+1]


    def plot(self):
        """Plot the results with Gnuplot."""
        data = Gnuplot.Data(self.x,self.y,using=(1,2)) #this ensures that t is used as x axis
        g = Gnuplot.Gnuplot()
        g('set ylabel "y-axis [arb. units]"')
        g('set xlabel "x-axis [arb. units]"')
        g('set style data lines')
        g.plot(data)


 
def demo():
    tmax = 5.*pi
    steps = 256
    t=linspace(0,tmax,steps)
    xp=sin(5*t)
    xn=xp+randn(len(t))*0.3
    ts = timeseries(t,xn)
    xp2 = sin(t) + 1.3*sin(2*t) + 1.2*sin(4*t) + 1.5*sin(6*t) + 1.9*sin(10.*t)
    xn2=xp2+randn(len(t))*0.2
    ts2 = timeseries(t,xn2)
    data1 = Gnuplot.Data(t,xn,using=(1,2))
    data2 = Gnuplot.Data(t,xn2,using=(1,2))
    g = Gnuplot.Gnuplot()
    g('set style data lines')
    g('set title "series 1 and series 2"')
    g.plot(data1,data2)
    ts_fft_k, ts_fft_y = ts.fft()
    norm = real_if_close(ts_fft_y*conjugate(ts_fft_y))
    nn = len(norm)
    xx = ts_fft_k[nn/2:nn]*2.*pi
    yy = 2.*norm[nn/2:nn]/steps**2
    plot(xx,yy,'Power spectral density series 1')
    ts2_fft_k, ts2_fft_y = ts2.fft()
    norm2 = real_if_close(ts2_fft_y*conjugate(ts2_fft_y))
    nn2 = len(norm2)
    xx2 = ts2_fft_k[nn2/2:nn2]*2.*pi
    yy2 = 2.*norm2[nn2/2:nn2]/steps**2
    plot(xx2,yy2,'Power spectral density series 2')

if __name__=='__main__':
    demo()
    raw_input('Please press the return key to continue...\n')