from brian import *
from scipy import signal, weave, random
from filterbank import Filterbank, RestructureFilterbank
from ..bufferable import Bufferable
from itertools import izip

__all__ = ['LinearFilterbank']

def _scipy_apply_linear_filterbank(b, a, x, zi):
    '''
    Parallel version of scipy lfilter command for a bank of n sequences of length 1
    
    In scipy.lfilter, you can apply a filter to multiple sounds at the same time,
    but you can't apply a bank of filters at the same time. This command does
    that. The coeffs b, a must be of shape (n,m,p), x must be of shape (s, n),
    and zi must be of shape (n,m-1,p). Here n is the number of channels in the
    filterbank, m is the order of the filter, p is the number of filters in
    a chain (cascade) to apply (you do first with (:,:,0) then (:,:,1), etc.),
    and s is the size of the buffer segment.
    '''
    alf_cache_b00 = [0]*zi.shape[2]
    alf_cache_a1 = [0]*zi.shape[2]
    alf_cache_b1 = [0]*zi.shape[2]
    alf_cache_zi00 = [0]*zi.shape[2]
    alf_cache_zi0 = [0]*zi.shape[2]
    alf_cache_zi1 = [0]*zi.shape[2]
    for curf in xrange(zi.shape[2]):
        alf_cache_b00[curf] = b[:, 0, curf]
        alf_cache_zi00[curf] = zi[:, 0, curf]
        alf_cache_b1[curf] = b[:, 1:b.shape[1], curf]
        alf_cache_a1[curf] = a[:, 1:b.shape[1], curf]
        alf_cache_zi0[curf] = zi[:, 0:b.shape[1]-1, curf]
        alf_cache_zi1[curf] = zi[:, 1:b.shape[1], curf]
    X = x.copy()
    output = empty_like(X)
    num_cascade = zi.shape[2]
    b_loop_size = b.shape[1]-2
    y = zeros(zi.shape[0])
    yr = reshape(y, (1, len(y))).T
    t = zeros(alf_cache_b1[0].shape, order='F')
    t2 = zeros(alf_cache_b1[0].shape, order='F')
    for sample, (x, o) in enumerate(izip(X, output)):
        xr = reshape(x, (1, len(x))).T
        for curf in xrange(num_cascade):
            #y = b[:, 0, curf]*x+zi[:, 0, curf]
            multiply(alf_cache_b00[curf], x, y)
            add(y, alf_cache_zi00[curf], y)
            #zi[:, :i-1, curf] = b[:, 1:i, curf]*xr+zi[:, 1:i, curf]-a[:, 1:i, curf]*yr
            multiply(alf_cache_b1[curf], xr, t)
            add(t, alf_cache_zi1[curf], t)
            multiply(alf_cache_a1[curf], yr, t2)
            subtract(t, t2, alf_cache_zi0[curf])
            u = x
            ur = xr
            x = y
            xr = yr
            y = u
            yr = ur
        #output[sample] = y
        o[:] = x
    return output


def _weave_apply_linear_filterbank(b, a, x, zi,
                                   cpp_compiler, extra_compile_args):
    if zi.shape[2]>1:
        # we need to do this so as not to alter the values in x in the C code below
        # but if zi.shape[2] is 1 there is only one filter in the chain and the
        # copy operation at the end of the C code will never happen.
        x = array(x, copy=True, order='C')
    else:
        # make sure that the array is in C-order
        x = asarray(x, order='C')
    y = empty_like(x)
    n, m, p = b.shape
    n1, m1, p1 = a.shape
    numsamples = x.shape[0]
    if n1!=n or m1!=m or p1!=p or x.shape!=(numsamples, n) or zi.shape!=(n, m, p):
        raise ValueError('Data has wrong shape.')
    if numsamples>1 and not x.flags['C_CONTIGUOUS']:
        raise ValueError('Input data must be C_CONTIGUOUS')
    if not b.flags['F_CONTIGUOUS'] or not a.flags['F_CONTIGUOUS'] or not zi.flags['F_CONTIGUOUS']:
        raise ValueError('Filter parameters must be F_CONTIGUOUS')
    code = '''
    #define X(s,i) x[(s)*n+(i)]
    #define Y(s,i) y[(s)*n+(i)]
    #define A(i,j,k) a[(i)+(j)*n+(k)*n*m]
    #define B(i,j,k) b[(i)+(j)*n+(k)*n*m]
    #define Zi(i,j,k) zi[(i)+(j)*n+(k)*n*(m-1)]
    for(int s=0; s<numsamples; s++)
    {
        for(int k=0; k<p; k++)
        {
            for(int j=0; j<n; j++)
                         Y(s,j) =   B(j,0,k)*X(s,j) + Zi(j,0,k);
            for(int i=0; i<m-2; i++)
                for(int j=0;j<n;j++)
                    Zi(j,i,k) = B(j,i+1,k)*X(s,j) + Zi(j,i+1,k) - A(j,i+1,k)*Y(s,j);
            for(int j=0; j<n; j++)
                  Zi(j,m-2,k) = B(j,m-1,k)*X(s,j)               - A(j,m-1,k)*Y(s,j);
            if(k<p-1)
                for(int j=0; j<n; j++)
                    X(s,j) = Y(s,j);
        }
    }
    '''
    weave.inline(code, ['b', 'a', 'x', 'zi', 'y', 'n', 'm', 'p', 'numsamples'],
                 compiler=cpp_compiler,
                 extra_compile_args=extra_compile_args)
    return y 


class LinearFilterbank(Filterbank):
    '''
    Generalised linear filterbank
    
    Initialisation arguments:

    ``source``
        The input to the filterbank, must have the same number of channels or
        just a single channel. In the latter case, the channels will be
        replicated.
    ``b``, ``a``
        The coeffs b, a must be of shape ``(nchannels, m)`` or
        ``(nchannels, m, p)``. Here ``m`` is
        the order of the filters, and ``p`` is the number of filters in a
        chain (first you apply ``[:, :, 0]``, then ``[:, :, 1]``, etc.).
    
    The filter parameters are stored in the modifiable attributes ``filt_b``,
    ``filt_a`` and ``filt_state`` (the variable ``z`` in the section below).
    
    Has one method:
    
    .. automethod:: decascade
    
    **Notes**
    
    These notes adapted from scipy's :func:`~scipy.signal.lfilter` function.
    
    The filterbank is implemented as a direct II transposed structure.
    This means that for a single channel and element of the filter cascade,
    the output y for an input x is defined by::

        a[0]*y[m] = b[0]*x[m] + b[1]*x[m-1] + ... + b[m]*x[0]
                              - a[1]*y[m-1] - ... - a[m]*y[0]

    using the following difference equations::

        y[i] = b[0]*x[i] + z[0,i-1]
        z[0,i] = b[1]*x[i] + z[1,i-1] - a[1]*y[i]
        ...
        z[m-3,i] = b[m-2]*x[i] + z[m-2,i-1] - a[m-2]*y[i]
        z[m-2,i] = b[m-1]*x[i] - a[m-1]*y[i]

    where i is the output sample number.

    The rational transfer function describing this filter in the
    z-transform domain is::
    
                                -1              -nb
                    b[0] + b[1]z  + ... + b[m] z
            Y(z) = --------------------------------- X(z)
                                -1              -na
                    a[0] + a[1]z  + ... + a[m] z
        
    '''
    def __init__(self, source, b, a):
        # Automatically duplicate mono input to fit the desired output shape
        if b.shape[0]!=source.nchannels:
            if source.nchannels!=1:
                raise ValueError('Can only automatically duplicate source channels for mono sources, use RestructureFilterbank.')
            source = RestructureFilterbank(source, b.shape[0])
        Filterbank.__init__(self, source)
        # Weave version of filtering requires Fortran ordering of filter params
        if len(b.shape)==2 and len(a.shape)==2:
            b = reshape(b, b.shape+(1,))
            a = reshape(a, a.shape+(1,))
        self.filt_b = array(b, order='F')
        self.filt_a = array(a, order='F')
        self.filt_state = zeros((b.shape[0], b.shape[1], b.shape[2]), order='F')
        self.use_weave = get_global_preference('useweave')
        if self.use_weave:
            log_info('brian.hears.filtering.linearfilterbank', 'Using weave')
            self.cpp_compiler = get_global_preference('weavecompiler')
            self.extra_compile_args = ['-O3']
            if self.cpp_compiler=='gcc':
                self.extra_compile_args += get_global_preference('gcc_options')


    def reset(self):
        self.buffer_init()
        
    def buffer_init(self):
        Filterbank.buffer_init(self)
        self.filt_state[:] = 0
    
    def buffer_apply(self, input):
        
        if self.use_weave:
            return _weave_apply_linear_filterbank(self.filt_b, self.filt_a, input,
                                                  self.filt_state, self.cpp_compiler,
                                                  self.extra_compile_args)
        else:
            return _scipy_apply_linear_filterbank(self.filt_b, self.filt_a, input,
                                                  self.filt_state)
        
    def decascade(self, ncascade=1):
        '''
        Reduces cascades of low order filters into smaller cascades of high order filters.
        
        ``ncascade`` is the number of cascaded filters to use, which should be
        a divisor of the original number.
        
        Note that higher order filters are often numerically unstable.
        '''
        n, m, p = self.filt_b.shape
        if p%ncascade!=0:
            raise ValueError('Number of cascades must be a divisor of original number of cascaded filters.')
        b = zeros((n, (m-1)*(p/ncascade)+1, ncascade))
        a = zeros((n, (m-1)*(p/ncascade)+1, ncascade))
        for i in xrange(n):
            for k in xrange(ncascade):
                bp = ones(1)
                ap = ones(1)
                for j in xrange(k*(p/ncascade), (k+1)*(p/ncascade)):
                    bp = polymul(bp, self.filt_b[i, ::-1, j])
                    ap = polymul(ap, self.filt_a[i, ::-1, j])
                bp = bp[::-1]
                ap = ap[::-1]
                a0 = ap[0]
                ap /= a0
                bp /= a0
                b[i, :len(bp), k] = bp
                a[i, :len(ap), k] = ap
        self.filt_b = array(b, order='F')
        self.filt_a = array(a, order='F')
        self.filt_state = zeros((b.shape[0], b.shape[1], b.shape[2]), order='F')
                
# Use the GPU version if available
try:
    if get_global_preference('brianhears_usegpu'):
        import pycuda
        from gpulinearfilterbank import LinearFilterbank
        use_gpu = True
    else:
        use_gpu = False
except ImportError:
    use_gpu = False
