#!/usr/bin/python -tt
#
# Werner Lustermann
# ETH Zurich
#
import numpy as np

import fir_filter as fir

class DRSSpikes(object):
    """ remove spikes (single or double false readings) from DRS4 data
    Strategy:
    * filter the data, removing the signal, thus spike(s) are clearly visible
    * search single and double spikes
    * replace the spike by a value derived from the neighbors
    
    """
    
    def __init__(self, threshold=7., 
                 single_pattern=np.array( [-0.5, 1.0, -0.5]) ,
                 double_pattern=np.array([-1., 1., 1., -1.]), 
                 user_action=lambda DRSSpikes_instance: None,
                 debug = False):
        """ initialize spike filter 
        template_single: template of a single slice spike
        template_double: template of a two slice spike

        """

        self.threshold = threshold
        self.single_pattern = single_pattern * threshold
        self.double_pattern = double_pattern * threshold
        
        self.remove_signal = fir.RemoveSignal()

        self.user_action = user_action
        self.debug = debug

    def __call__(self, data):

        self.row, self.col = data.shape
        indicator = self.remove_signal(data)
        a = indicator.flatten()
        singles = []
        doubles = []

        # a spike in the first or last channel is considered as a filter artefact
        candidates = np.where(a[1:-2] > self.threshold)
        # candidates = np.where(a[1:1022] > self.threshold)
        cc = candidates[0]
        #print 'cc: ', cc
        #: find single spikes
        p = self.single_pattern * np.sign( self.single_pattern )
        for i, can in enumerate( zip(a[cc], a[cc+1], a[cc+2]) ):
            #print 'can : p', can, p
            can = can * np.sign(self.single_pattern)
            if np.all(can > p):
                singles.append(cc[i])

        #: find double spikes
        p = self.double_pattern * np.sign( self.double_pattern )
        for i, can in enumerate( zip(a[cc], a[cc+1], a[cc+2], a[cc+3]) ):
            #print 'data:  ', [data[0,cc[i]+k] for k in range(3)]
            #print 'can : p', can, p
            can = can * np.sign(self.double_pattern)
            if np.all(can > p):
                doubles.append(cc[i])

        if self.debug:
            print 'singles: ', singles
            print 'doubles: ', doubles
        
        # make relevant information a member,
        self.singles = singles
        self.doubles = doubles 
        self.indicator = a
        self.candidates = cc

        self.user_action(self)

        data = self.remove_single_spikes(singles, data)
        data = self.remove_double_spikes(doubles, data)
        return data

    def remove_single_spikes(self, singles, data):
        data = data.flatten() 
        for spike in singles:
            data[spike] = (data[spike-1] + data[spike+1]) / 2.
        return data.reshape(self.row, self.col)
    
    def remove_double_spikes(self, doubles, data):
        data = data.flatten() 
        for spike in doubles:
            data[spike:spike+2] = (data[spike-1] + data[spike+2]) / 2.
        return data.reshape(self.row, self.col)


def _test():
  
    a = np.ones((3,12)) * 3.
    a[0,3] = 7.
    a[1,7] = 14.
    a[1,8] = 14.
    a[2,4] = 50.
    a[2,5] = 45.
    a[2,8] = 20.
    
    print a 

    SpikeRemover = DRSSpikes(3., debug=True)
    print 'single spike pattern ', SpikeRemover.single_pattern
    print 'double spike pattern ', SpikeRemover.double_pattern
    afilt = SpikeRemover(a)
    print afilt
    
if __name__ == '__main__':
    """ test the class """
    _test()