#! /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'):
    #x, y = tsc.fft()
    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'):
        #print "in smooth method"
        #print x, window_len, window
        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 = 6*pi
    steps = 128
    empty = zeros(steps)
    t=linspace(0,tmax,steps)
    xp1=sin(t)
    xp2 = sin(t-1) 
    xn2=xp2+randn(len(t))*0.15
    xn1=xp1+randn(len(t))*0.08
    xn1=append(xn1,empty) #Add zero padding
    xn2=append(xn2,empty) #Add zero padding
    t = linspace(0,2*tmax,2*steps) #add because of xn1, xn2 zero padding
    ts1 = timeseries(t,xn1)
    ts2 = timeseries(t,xn2)
    data1 = Gnuplot.Data(t,xn1,using=(1,2))
    data2 = Gnuplot.Data(t,xn2,using=(1,2))
    g = Gnuplot.Gnuplot()
    g('set style data lines')
    g.plot(data1,data2)

    fft_xn1 = fft(xn1)
    fft_xn2 = fft(xn2)
    norm1 = real(ifft(fft_xn2*conjugate(fft_xn1)))/steps*2.
    print norm1
    nn1 = len(norm1)
    print nn1
    xx = t[0:nn1/2]
    yy = norm1[0:nn1/2]
    #plot(xx,yy,'xx-yy')
    print len(t), len(xx), len(yy)
    plot(xx,yy,'correlation function between xn1 and xn2')
    

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