#! /usr/bin/env python
###############################################################################
#                                                                             #
#   Copyright 2005 University of Cambridge Computer Laboratory.               #
#                                                                             #
#   This file is part of Nprobe.                                              #
#                                                                             #
#   Nprobe is free software; you can redistribute it and/or modify            #
#   it under the terms of the GNU General Public License as published by      #
#   the Free Software Foundation; either version 2 of the License, or         #
#   (at your option) any later version.                                       #
#                                                                             #
#   Nprobe is distributed in the hope that it will be useful,                 #
#   but WITHOUT ANY WARRANTY; without even the implied warranty of            #
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the             #
#   GNU General Public License for more details.                              #
#                                                                             #
#   You should have received a copy of the GNU General Public License         #
#   along with Nprobe; if not, write to the Free Software                     #
#   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA #
#                                                                             #
###############################################################################


##############################################################################
## 
## Smooth time series data by averaging over specified period (in s)
##  - various specified smooths can be applied, successively if required
## 
##
## 

############################################################################ 

import string
import glob
import os
from os.path import basename
import sys
from copy import deepcopy

from math import floor, sin, pi

#import nprobe
from sys import argv
import getopt
from signal import *
from np_plot import *
from Tkinter import *
import tkSimpleDialog, tkFileDialog
from np_printplot import OverPrint
	    
	     
##############################################################################
##############################################################################

#
# Can't get data exception
#

class NoDataError:

    def __init__(self, val):
        self.val = val

    def __str__(self):
	return self.val
	    
	     
##############################################################################
##############################################################################

#
# Smooth error exception
#

class SmoothError:

    def __init__(self, val):
        self.val = val

    def __str__(self):
	return self.val

############################################################################## 
#############################################################################

class TsSmooth:

    def __init__(self, data=None, infile='', spec=None, draw=0, show_stages=0,
                 savefile='', savesde=0):

        self.data = data
        self.infile=infile
        self.spec = spec
        self.do_draw = draw
        self.show_stages = show_stages
        self.savesde = savesde
        if savefile:
            self.savefile = savefile
        elif infile:
            self.savefile = infile
        else:
            self.savefile = 'saved_smooth'

        self.smooths = []
        self.opath = ''

        if (not data) and infile:
            self.load_data(infile)

        if spec:
            print savefile
            self.build_smooths(spec, self.savefile)

        if self.data and self.smooths:
            self.smooth()
            
            if draw:
                self.draw()
            
	     
##############################################################################

    def smooth_data(self, data):

        self.data = data
        self.smooth()
        return self.smoothed_data
            
	     
##############################################################################

    def draw(self):

        if self.savesde:
            fields = ['X', 'Y', 'SD', 'SE']
        else:
            fields = ['X', 'Y']

        self.smoothed[-1].fields = fields

        if not self.show_stages:
            smoothed = [self.smoothed[-1]]
        else:
            smoothed = self.smoothed

        try:
            np_Plot(smoothed, standalone='yes', path=self.savefile,  
                title='%s (smoothed %s)' % (basename(self.infile), self.opath),
                    xlab='time s', ylab='ms')
        except EmptyDataSetError:
            print '%s - empty data set' % (ds.path)
            sys.exit(1)
            
	     
##############################################################################

    def smooth(self):

        if not self.data:
            raise NoDataError('No data to smooth')

        if not self.smooths:
            raise SmoothError('No smooth specification')

        data = self.data
        data.sort()
        

        i = 0
        if self.show_stages:
            self.smoothed = [DataSet(data, DATA_TS, 'Raw data', i, None)]
            i += 1
        else:
            self.smoothed = []
            
        regd = 0

        #if self.savesde:
            #fields = ['X', 'Y', ]
        for smooth in self.smooths:
            data = smooth[0](data, smooth[1])#[:]
            self.smoothed.append(DataSet(data, DATA_TS, smooth[2], i, None))
            i += 1

        self.smoothed_data = data
            
	     
##############################################################################

    def get_moothed_data(self):
            
	   return self.smoothed_data
       
##############################################################################

    def load_data(self, file):

        self.infile = file
        data = []
        try:
            f = open(file, 'r')
        except IOError, s:
            raise NoDataError(s)
        line = 0
        while 1:
            s = f.readline()
            line = line +1
            if not len(s):
                #print 'empty'
                break
            if s[0] == '#' or s[0] == ' ' or s[0] =='\n':
                #print 'comment'
                continue
            vs = string.split(s)
            try:
                data.append([string.atof(vs[0]), string.atof(vs[1])])
            except IndexError:
                raise NoDataError('Malformed data line %d file %s' % (line, file))

        if not data:
            raise NoDataError('No data in file %s' % (file))


        self.data = data
       
##############################################################################

    def save_data(self):
        
	fnm = self.filepath
        fnm = fnm.replace(':', '-')
	root = Tk()
	root.geometry("0x0")
	results = []
	#fnm = '%s.p%s.smoothed.e' % (args[0], pstr)
	OverPrint(root, title='Save smoothed data', 
		  msg='save data as:', 
		  entries = [('FILE:', fnm)],results=results)
	if len(results):
	    fnm = results[0]
        else:
            fnm = ''
        #print results
        #print fnm
	while len(fnm):
	    results = []
	    try:
		f = open(fnm, 'r')
	    except:
		# doesn't already exist - carry on
		break
	    else:
		# check
		OverPrint(root, title='Overprint', 
			     msg='file already exists - overwrite?', 
			     entries = [('FILE:', fnm)],results=results)
                print results
		if len(results):
		    if fnm == results[0]:
			break
		    fnm = results[0]
                else:
                    fnm = ''
                    break
	
	root.destroy()

	if len(fnm):
            if self.savesde:
                fieldstr = '#fields=X:Y:SD:SE'
            else:
                fieldstr = '#fields=X:Y'
            
            f = open(fnm, 'w')
            f.write('%s\n' % (fieldstr))
            
	    for e in self.smoothed_data:
                if self.savesde:
                    f.write('%f\t%f\t%f\t%f\n' % (e[0], e[1], e[2], e[3]))
                else:
                    f.write('%f\t%f\n' % (e[0], e[1]))
	    f.close()

            print 'File %s written' % (fnm)
	    
	     
##############################################################################

    def smooth_mean(self, d, args):

        #
        # Replace data occuring in each period P with mean for period at period
        # mid point. Generate SD and SE for each period
        #


        p = args[0]
        if p == None:
            print 'WARNING not performing initial mean binning'
            return d
        ss = self.savesde

        print 'performing mean smooth, period %d' % (p)

        start = d[0][0]
        tm = 0
        while tm < start:
            print tm, p
            tm += p
        #tm = start + p
        n = 0

        smoothed = []
        agg = 0.0
        aggsq = 0.0

        for e in d:
            t = e[0] - start
            v = e[1]

            if t > tm:
                #print '(%.3f) t=%.3f tm=%.3f n=%d' % (e[0], t, tm, n)
                if n:
                    #print 'ap'
                    M = agg/n
                    #print 'n=%d M=%f agg=%f aggsq=%f ' % (n,M,agg,aggsq),
                    V= (aggsq - (pow(agg, 2)/n))/n  #(n-1)?
                    sd = sqrt(V)
                    if ss:
                        smoothed.append([tm, M, sd, sd/sqrt(n)])
                    else:
                        smoothed.append([tm, M])
                    #print 'sd=%f stderr=%f' % (sd, sd/sqrt(n))
                    n = 0
                    agg = 0.0
                    aggsq = 0.0
                while tm < t:
                    tm += p  
                    #print'tm now %.3f' % (tm)
            n += 1
            agg += v
            aggsq += pow(v, 2)

        if n:
            M = agg/n
            #print 'n=%d M=%f agg=%f aggsq=%f ' % (n,M,agg,aggsq),
            V= (aggsq - (pow(agg, 2)/n))/n #(n-1)?
            sd = sqrt(V)
            if ss:
                smoothed.append([tm, M, sd, sd/sqrt(n)])
            else:
                smoothed.append([tm, M])
            #print 'sd=%f stderr=%f' % (sd, sd/sqrt(n))

        #print 'mean smoothed', smoothed

        return smoothed
	    
	     
##############################################################################
	    
	     
##############################################################################

    def smooth_meanmin(self, d, args):

        #
        # Replace data occuring in each period P with mean for period at period
        # mid point. Generate SD and SE for each period
        #


        p = args[0]
        if p == None:
            print 'WARNING not performing initial mean binning'
            return d
        ss = self.savesde

        print 'performing mean smooth, period %d' % (p)

        start = d[0][0]
        tm = 0
        while tm < start:
            print tm, p
            tm += p
        #tm = start + p
        n = 0

        smoothed = []
        Min = 1000000

        for e in d:
            t = e[0] - start
            v = e[1]

            if t > tm:
                #print '(%.3f) t=%.3f tm=%.3f n=%d' % (e[0], t, tm, n)
                if n:
                    smoothed.append([tm, Min])
                    Min = 1000000
                    n = 0
                while tm < t:
                    tm += p  
                    #print'tm now %.3f' % (tm)
            Min = min(Min, v)
            n += 1

        if n:
            smoothed.append([tm, Min])

        return smoothed
	    
	     
##############################################################################


    def smooth_sma(self, d, args):

        #
        # Simple moving average
        #


        span = args[0]

        l = len(d)   
        sm = deepcopy(d)
        d = [e[1] for e in d]

        print 'performing sma smooth, span %d' % (span)

        if span == 0:
            raise SmoothError('sma smooth: span = 0')
        if span >= l:
            raise SmoothError('sma smooth: span %d >= data length %d' % (span, l))

        for x in range(1, l):
            p = min(x, span)
            if p <= span:
                sum = 0.0
                for y in range(p):
                    sum += d[x-y]
            else:
                sum = sum - d[x-p] + d[x]

            sm[x][1] = sum/p

        return sm   
	    
	     
##############################################################################

    def smooth_ema(self, d, args):

        #
        # Single exponential moving average
        #


        a = args[0]

        l = len(d)   
        sm = deepcopy(d)
        d = [e[1] for e in d]

        if not 0 < a <= 1:
            raise SmoothError('ema smooth: invalid alpha %f (0<alpha<=1)' % (a))

        print 'performing ema smooth, alpha %f' % (a)  

        se = a
        le = 1.0 - a

        # get something sensible for first value
        n = min(4, l)
        acc = 0.0
        for x in range(n):
            acc += d[x]
        sm[0][1] = ((acc/n)*le) + (sm[0][1]*se)

        for x in range(1, l):
            sm[x][1] = d[x]*se + sm[x-1][1]*le

        return sm    
	     
##############################################################################

    def han(self, d, args):

        #
        # Hanning smooth
        #


        sm = deepcopy(d)
        l = len(d)

        if l < 3:
            raise SmoothError('han: need > 3 data points')

        print 'hanning'

        d = [e[1] for e in d]

        sm[0][1] = 0.75*d[0] + 0.25*d[1]
        for x in range(1, l-1):
            sm[x][1] = 0.25*d[x-1] + 0.5*d[x] + 0.25*d[x+1]
        sm[l-1][1] = 0.25*d[l-2] + 0.75*d[l-1]

        return sm
	     
##############################################################################

    def smooth_sft(self, d, args):

        #
        # FFT - based smooth (low-pass filter)
        #

        def four1(d, nn, isign):

            n = nn << 1
            j = 1
            for i in range(1, n, 2):
                if j > i:
                    tmp = d[j]
                    d[j] = d[i]
                    d[i] = tmp
                    tmp = d[j+1]
                    d[j+1] = d[i+1]
                    d[i+1] = tmp
                m = n/2
                while  j > m >= 2:
                    j -= m;
                    m  >> 1
                j += m

            mmax = 2
            while n > mmax:
                istep = mmax << 1
                theta = (pi*2)/(isign*mmax)
                wtmp = sin(0.5*theta)
                wpr = -2.0*wtmp*wtmp
                wpi = sin(theta)
                wr = 1.0
                wi = 0.0
                for m in range(1, mmax, 2):
                    for i in range(m, n+1, istep):
                        j = i + mmax
                        print j
                        tmpr = wr*d[j] - wi*d[j+1]
                        tmpi = wr*d[j+1] + wi*d[j]
                        d[j] = d[i] - tmpr
                        d[j+1] = d[i+1] - tmpi
                        d[i] += tmpr
                        d[i+1] += tmpi
                    wtmp = wr
                    wr = wr*wpr - wi*wpi + wr
                    wi = wi*wpr + wtmp*wpi + wi
                mmax = istep
                    
                    

        def realfft(d, n, isign):

            theta = pi/n
            c1 = 0.5
            if isign == 1:
                c2 = -0.5
                four1(d, n, 1)
            else:
                c2 = 0.5
                theta = -theta

            wpr = -2.0*pow(sin(0.5*theta), 2)
            wpi = sin(theta)
            wr = 1.0 + wpr
            wi = wpi
            np3 = 2*n+3

            for i in range(2, n/2+1):
                i1 = 2*i-1
                i2 = i1 + 1
                i3 = np3 - i2
                #i3 = n-2*i
                i4 = i3 + 1
                #print i1, i2, i3, i4
                h1r = c1*(d[i1] + d[i3])
                h1i = c1*(d[i2] - d[i4])
                h2r = -c2*(d[i2] + d[i4])
                h2i = c2*(d[i1] - d[i3])
                d[i1] = h1r + wr*h2r - wi*h2i
                d[i2] = h1i + wr*h2i + wi*h2r
                d[i3] = h1r - wr*h2r + wi*h2i
                d[i4] = -h1i + wr*h2i + wi*h2r
                wtmp = wr
                wr = wr*wpr - wi*wpi + wr
                wi = wi*wpr + wtmp*wpi + wi

            if isign == 1:
                h1r = d[1]
                d[1] = h1r + d[2]
                d[2] = h1r - d[2]

            else:
                h1r = d[1]
                d[1] = c1*(h1r + d[2])
                d[2] = c1*(h1r - d[2])
                four1(d, n, -1)

            

        span = args[0]

        sm = deepcopy(d)
        d = [e[1] for e in d]
        n = len(d)
        d.insert(0, 0.0)

        # calculate o/a size of working array (s >= 2**x >= n + 2*span) 
        l = n + 1 + int(floor((2*span)+0.5))
        m = 2
        while m < l:
            m *= 2

        for i in range(m-n):
            d.append(0.0)

        print len(d)

        #
        # here's the smooth
        #
        cnst = pow(span/m, 2)
        d1 = d[1]
        dn = d[n]
        rn1 = 1.0/(n-1.0)
        for i in range(1, n+1):
            d[i] -= rn1*(d1*(n-i)+dn*(i-1))
        mo2 = m/2

        realfft(d, mo2, 1)

        d[1] /= mo2
        fac = 1.0
        #print cnst
        for i in range(1, mo2):
            k = (2*i+1)
            print k, fac
            if fac != 0.0:
                fac = (1.0 - cnst*i*i)/mo2
                if fac < 0.0:
                    fac = 0.0
                d[k] *= fac
                d[k+1] *= fac
            else:
                d[k] = 0.0
                d[k+1] = 0.0
        fac = (1.0 - (0.25*span*span))/mo2
        if fac < 0.0:
            fac = 0.0
        d[2] *= fac

        realfft(d, mo2, -1)

        for i in range(1, n+1):
            d[i] += rn1*(d1*(n-i) + dn*(i-1))
            #print d[i-1]

                    
        for i in range(n):
            sm[i][1] = d[i-1]

        return sm

        
	     
##############################################################################


    def build_smooths(self, smoothopts, path):

        need_reg = 1
        first = 1
        smooths = []
        opath = ''

        for s in smoothopts:
            #opath += '.%s' % (s)
            sp = s.split(':')
            if sp[0] == 'mean':
                if not len(sp) == 2:
                    raise SmoothError('mean - no averaging period given')
                try:
                    p = float(sp[1])
                    opath += '.%s' % (s)
                except:
                    if sp[1] == '-':
                        # inhibit initial binning
                        p = None
                    else:
                        raise SmoothError('mean - invalid period')
                smooths.append((self.smooth_mean, (p,), opath))

            elif sp[0] == 'sma':
                if not len(sp) >= 2:
                    raise SmoothError('sma - no span(s) given')
                opath += sp[0]
                for span in sp[1:]:
                    try:
                        p = int(span)
                    except:
                        raise SmoothError('sma - invalid span %d' % (span))
                    opath += ':%s' % (span)
                    smooths.append((self.smooth_sma, (p,), opath))
            elif sp[0] == 'ema':
                if not len(sp) == 2:
                    raise SmoothError('ema - no alpha given')
                try:
                    p = float(sp[1])
                except:
                    raise SmoothError('ema - invalid alpha')
                opath += '.%s' % (s)
                smooths.append((self.smooth_ema, (p,), opath))
            elif sp[0] == 'sft':
                if not len(sp) == 2:
                    raise SmoothError('sft - no alpha given')
                try:
                    p = float(sp[1])
                except:
                    raise SmoothError('sft - invalid alpha')
                opath += '.%s' % (s)
                smooths.append((self.smooth_sft, (p,), opath))
            elif sp[0] == 'han':
                opath += '.%s' % (s)
                smooths.append((self.han, (0,), opath))
            else:
                raise SmoothError('smooth \'%s\' not recognised' % (sp[0]))

            # if not starting with mean smooth normalise time periods
            if first and sp[0] != 'mean':
                smooths.insert(0, (self.smooth_mean, (1,), 'mean:1'))
            first = 0

        if path[-1] == '.':
            opath = opath[1:]

        self.smooths = smooths
        self.opath = opath
        self.filepath = path + opath

    

	     
##############################################################################
##############################################################################

def usage(s):

    print 'usage %s: <opts> <data file>' % (s)
    print 'opts:'
    print '\t-d draw results (default if not saving)'
    print '\t-s:spec1 -s:spec2 ...  smooth specifiers in order applied'
    print '\t  available smooths:'
    print '\t    -s<mean:p> - average over periods p'
    print '\t    -s<sma:r1:r2:r3... - simple moving average successively\n\t      over distances r1, r2, r3 ...'
    print '\t    -s<ema:a> - exponential moving average over alpha a'
    print '\t  if no initial mean smooth is given prior to sma or ema a\n          period 1 mean will be applied'
    print '\t-p:x shorthand for -smean:x'
    print '\t-S save smoothed data'
    print '\t-i show original data, intermediate results and final data'
    print '\t-e save s.d. and s.e.m. for averaging periods (output cols. 3/4)'
    print '\t-h(elp)'
    
    sys.exit(1)
	    
	     
##############################################################################

def get_file(dir):

    root = Tk()
    root.winfo_toplevel().title('Get file:')
    fnm =  tkFileDialog.askopenfilename( \
        title='File', initialdir=dir,     
        filetypes=[('All files', '.*')])

    root.destroy()
    del root
    return fnm
       
##############################################################################




def main():

    scriptname = sys.argv[0]

    p = None
    pstr = ''
    savedata = 0
    smoothopts = []
    smoothed = []
    show_stages = 0
    savesde = 0
    draw = 0
    
    try:
        optlist, args = getopt.getopt(sys.argv[1:], 'hSs:ip:ed')
	
    except getopt.GetoptError, s:
        print 'ts_smooth: %s' % (s)
        usage()
        sys.exit(1)
	
    for opt in optlist:
	if opt[0] == "-h":
	    usage(scriptname)
	if opt[0] == "-S":
	    savedata = 1
	if opt[0] == "-s":
	    smoothopts.append(opt[1])
	if opt[0] == "-p":
	    smoothopts.append('mean:' + opt[1])
	if opt[0] == "-i":
	    show_stages = 1
	if opt[0] == "-e":
	    savesde = 1
	if opt[0] == "-d":
	    draw = 1

    if not savedata:
        draw = 1

    print args
    
    if len(args) == 0:
        fnm = get_file(os.getcwd())
    elif len(args) == 1 and os.path.isdir(args[0]):
        fnm = get_file(os.path.dirname(args[0]))
    elif len(args) == 1 and os.path.isfile(args[0]):
        fnm = args[0]

    if not fnm:
        sys.exit(0)
        
    basepath = fnm

    if not smoothopts:
        print 'No smooth specified'
        usage()
    else:
        try:
            smoother = TsSmooth(infile=fnm, spec=smoothopts, draw=draw,
                           show_stages=show_stages, savesde=savesde)
        except NoDataError, s:
            print s.val
            sys.exit(1)
        except SmoothError, s:
            print s
            sys.exit(1)

    if savedata:
        smoother.save_data()
	     
##############################################################################


# Call main when run as script
if __name__ == '__main__':
        main()
