#!/usr/bin/python -tt
#
# Dominik Neise, Werner Lustermann
# TU Dortmund, ETH Zurich
#
import numpy as np

class SignalGenerator(object):
    """ Signal Generator
        generates signals for testing several helper classes like:
        * fir filters
        * signal extractors
    """
    
    def __init__(self, option_str = 'len 100 noise 3', name = 'SignalGenerator'):
        """ initialize the generator
            sets default signal to generate
        """
        self.__module__ = 'generator'
        self.option_str = option_str.lower()
        self.options = make_options_from_str(option_str)
        self.parse_options()
        self.name = name
        
    def parse_options(self):
        o = self.options #shortcut
        if 'len' in o:
            self.npoints = int(o['len'][0])
        else: 
            self.npoints = 100
        if 'noise' in o:
            self.sigma = float(o['noise'][0])
        else:
            self.sigma = 1
        if 'bsl' in o:
            self.bsl = float(o['bsl'][0])
        else:
            self.bsl = -0.5
        
        if 'step' in o:
            self.step_height = float(o['step'][0])
            self.step_start = int(o['step'][1])
            self.step_stop = int(o['step'][2])

        if 'triangle' in o:
            self.pulses = []
            # append 1st pulse to list of pulses
            self.pulses.append( ( float(o['triangle'][0]) , float(o['triangle'][1]), int(o['triangle'][2]), int(o['triangle'][3]) ) )
            number_of_pulses_after_1st = (len(o['triangle'])-4)/2
            for i in range(number_of_pulses_after_1st):
                self.pulses.append( ( float(o['triangle'][2*i+4]) , float(o['triangle'][2*i+5]), int(o['triangle'][2]), int(o['triangle'][3]) ) )

        if 'spike' in o:
            self.spikes = []
            for i in range(len(o['spike'])/2):
                self.spikes.append( ( int(o['spike'][2*i]), float(o['spike'][2*i+1]) ) )

    def __call__(self, option_str = ''):
        if option_str:
            self.option_str = option_str.lower()
            self.options = make_options_from_str(self.option_str)
            self.parse_options()

        signal = np.zeros(self.npoints)
        signal += self.bsl
        signal += np.random.randn(self.npoints) * self.sigma
        if 'step' in self.options:
            signal[self.step_start:self.step_stop] += self.step_height
        if 'triangle' in self.options:
            for pulse in self.pulses:
                pos = pulse[0]
                height = pulse[1]
                rise = pulse[2]
                fall = pulse[3]
                start = pos - rise
                stop = pos + fall
                signal[start:pos] += np.linspace(0., height, rise)
                signal[pos:stop] += np.linspace(height, 0. , fall)
        if 'spike' in self.options:
            for spike in self.spikes:
                signal[spike[0]] += spike[1]
        return signal 

    def __str__(self):
        s = self.name + '\n'
        s += 'possible options and parameters\n'
        s += ' * len:      number of samples (100)\n'
        s += ' * noise:    sigma (1)\n'
        s += ' * bsl:      level (-0.5)\n'
        s += ' * step:     height, start, end\n'
        s += ' * triangle: pos height risingedge, fallingedge [pos height ...]\n'
        s += ' * spike:    pos height [pos height ...]\n'
        
        s += 'current options are:\n'
        for key in self.options.keys():
            s += key + ':' + str(self.options[key]) + '\n'
        return s


class SignalGeneratorCSV(object):
    
    def __init__(self, file_name, option_str = 'len 100 noise 3', name = 'SignalGenerator'):
        time, maxprob, mean, median = np.loadtxt( file_name, delimiter=',', unpack=True)
        csv_data = maxprob
        
        # csv data was downshifted, I shift it up here
        self.csv_data = csv_data - csv_data.min()
        

        self.__module__ = 'CSV generator'
        self.option_str = option_str.lower()
        self.options = make_options_from_str(option_str)
        self.parse_options()
        self.name = name
        
    def parse_options(self):
        o = self.options #shortcut
        if 'len' in o:
            self.npoints = int(o['len'][0])
        else: 
            self.npoints = 100
        if 'noise' in o:
            self.sigma = float(o['noise'][0])
        else:
            self.sigma = 1
        if 'bsl' in o:
            self.bsl = float(o['bsl'][0])
        else:
            self.bsl = -0.5
        
        if 'step' in o:
            self.step_height = float(o['step'][0])
            self.step_start = int(o['step'][1])
            self.step_stop = int(o['step'][2])

        if 'triangle' in o:
            self.pulses = []
            # append 1st pulse to list of pulses
            self.pulses.append( ( float(o['triangle'][0]) , float(o['triangle'][1]), int(o['triangle'][2]), int(o['triangle'][3]) ) )
            number_of_pulses_after_1st = (len(o['triangle'])-4)/2
            for i in range(number_of_pulses_after_1st):
                self.pulses.append( ( float(o['triangle'][2*i+4]) , float(o['triangle'][2*i+5]), int(o['triangle'][2]), int(o['triangle'][3]) ) )

        if 'spike' in o:
            self.spikes = []
            for i in range(len(o['spike'])/2):
                self.spikes.append( ( int(o['spike'][2*i]), float(o['spike'][2*i+1]) ) )
        
        if 'csv' in o:
            self.csvs = []
            for i in range(len(o['csv'])/2):
                time = int( o['csv'][2*i] )
                amplitude = float( o['csv'][2*i+1] )
                self.csvs.append( (time, amplitude) )

    def __call__(self, option_str = ''):
        if option_str:
            self.option_str = option_str.lower()
            self.options = make_options_from_str(self.option_str)
            self.parse_options()

        signal = np.zeros(self.npoints)
        signal += self.bsl
        
        if 'step' in self.options:
            signal[self.step_start:self.step_stop] += self.step_height
        if 'triangle' in self.options:
            for pulse in self.pulses:
                pos = pulse[0]
                height = pulse[1]
                rise = pulse[2]
                fall = pulse[3]
                start = pos - rise
                stop = pos + fall
                signal[start:pos] += np.linspace(0., height, rise)
                signal[pos:stop] += np.linspace(height, 0. , fall)
        if 'spike' in self.options:
            for spike in self.spikes:
                signal[spike[0]] += spike[1]
        if 'csv' in self.options:
            for csv in self.csvs:
                amplitude = csv[1]
                time = csv[0]
                csv_data = self.csv_data.copy()
                #scale
                csv_data *= amplitude
                # add shifted
                print 'bumm', len(csv_data)
                print csv_data.shape
                print signal[time:time+len(csv_data)].shape
                signal[time:time+len(csv_data)] += csv_data
                
        # add noise
        signal += + np.random.normal(0.0,self.sigma, signal.shape)
        return signal 

    def __str__(self):
        s = self.name + '\n'
        s += 'possible options and parameters\n'
        s += ' * len:      number of samples (100)\n'
        s += ' * noise:    sigma (1)\n'
        s += ' * bsl:      level (-0.5)\n'
        s += ' * step:     height, start, end\n'
        s += ' * triangle: pos height risingedge, fallingedge [pos height ...]\n'
        s += ' * spike:    pos height [pos height ...]\n'
        s += ' * csv:      pos height [pos height ...]\n'
        
        s += 'current options are:\n'
        for key in self.options.keys():
            s += key + ':' + str(self.options[key]) + '\n'
        return s


# Helper function to parse signalname and create a dictionary
# dictionary layout :
# key : string
# value : [list of parameters]
def make_options_from_str(signalname):
    options = {}
    for word in (signalname.lower()).split():
        if word.isalpha():
            current_key = word
            options[current_key] = []
#        if word.isdigit():
        else:
            options[current_key].append(word)
#        else:
#            print '-nothing'
    return options
    
if __name__ == '__main__':
    from plotters import Plotter
    myGenerator = SignalGenerator('len 400 noise 0.3 bsl -2.5 triangle 50 10.2 10 100 65 10 150 20 spike 100 50. 20 50 21 49')
    sig = myGenerator()
    print myGenerator
    
    p = Plotter('generator test')
    p(sig)
    
    anothergen = SignalGeneratorCSV('PulseTemplate_PointSet_0.csv', 
                'len 1000 noise 0.4 bsl -2.0 csv 300 1 60 2 650 1 spike 110 50')
    sig2 = anothergen()
    print anothergen
    
    pp = Plotter('CSV Gen Test')
    pp(sig2)
    
    raw_input('any key to quit')