# -*- coding: utf-8 -*-

"""Waveform filtering routines."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import numpy as np
from scipy import signal


#------------------------------------------------------------------------------
# Waveform filtering routines
#------------------------------------------------------------------------------

def bandpass_filter(rate=None, low=None, high=None, order=None):
    """Butterworth bandpass filter."""
    assert low < high
    assert order >= 1
    return signal.butter(order,
                         (low / (rate / 2.), high / (rate / 2.)),
                         'pass')


def apply_filter(x, filter=None, axis=0):
    """Apply a filter to an array."""
    if isinstance(x, list):
        x = np.asarray(x)
    if x.shape[axis] == 0:
        return x
    b, a = filter
    return signal.filtfilt(b, a, x[:], axis=axis)


class Filter(object):
    """Multichannel bandpass filter.

    The filter is applied on every column of a 2D array.

    Example
    -------

    ```python
    fil = Filter(rate=20000., low=5000., high=15000., order=4)
    traces_f = fil(traces)
    ```

    """
    def __init__(self, rate=None, low=None, high=None, order=None):
        self._filter = bandpass_filter(rate=rate,
                                       low=low,
                                       high=high,
                                       order=order,
                                       )

    def __call__(self, data):
        return apply_filter(data, filter=self._filter)


#------------------------------------------------------------------------------
# Whitening
#------------------------------------------------------------------------------

class Whitening(object):
    """Compute a whitening matrix and apply it to data.

    Contributed by Pierre Yger.

    """
    def fit(self, x, fudge=1e-18):
        """Compute the whitening matrix.

        Parameters
        ----------

        x : array
            An `(n_samples, n_channels)` array.

        """
        assert x.ndim == 2
        ns, nc = x.shape
        x_cov = np.cov(x, rowvar=0)
        assert x_cov.shape == (nc, nc)
        d, v = np.linalg.eigh(x_cov)
        d = np.diag(1. / np.sqrt(d + fudge))
        # This is equivalent, but seems much slower...
        # w = np.einsum('il,lk,jk->ij', v, d, v)
        w = np.dot(np.dot(v, d), v.T)
        self._matrix = w
        return w

    def transform(self, x):
        """Whiten some data.

        Parameters
        ----------

        x : array
            An `(n_samples, n_channels)` array.

        """
        return np.dot(x, self._matrix)
