
import numpy as np

import inspect

class scattering_tree:
    
    def __init__(self, propagation_filters, output_filters, nonlinearities, poolings):
        self._prop_filters = propagation_filters
        self._out_filters = output_filters
        self._nonlin = nonlinearities
        self._pools = poolings

    def scatter_linearized(self, signal):
        return self.scatter(signal)
        
        
    def scatter(self, signal):
        """
        Applies the scattering transform onto the signal and returns the result as a flat numpy array of floats.
        Complex values get splittet into their real and imaginary part.
        """
        return _scatter_flat(signal, self._prop_filters, self._out_filters, self._nonlin, self._pools)


def _scatter_flat(signal, propagation_filters, output_filters, nonlinearities, poolings):
    """
    Applies the scattering transform onto the signal and returns the result as a flat numpy array of floats.
    Complex values get splittet into their real and imaginary part.
    """
    
    sc = _scatter(signal, propagation_filters, output_filters, nonlinearities, poolings)
    
    flat = [s.flatten() for s in sc]
    flat = _split_complex(flat)

    return np.concatenate(flat)


def _scatter(signal, propagation_filters, output_filters, nonlinearities, poolings):
    """
    @param signal  Either a numpy array or a list of numpy arrays.
    """
    output = []
    nodes = []
    
    """
    This should allow for multiple input signals (as a list) such as R,G,B channels of an image.
    """
    if type(signal) is list:
        nodes.extend(signal)
    else:
        nodes.append(signal)
        
    meta = [None]*len(nodes)
    
    for prop_filters, out_filters, nl, pooling in zip(propagation_filters, output_filters, nonlinearities, poolings):
        next_nodes = []
        next_meta = []
        # scatter each node on that level
        for n, m in zip(nodes, meta):
            sc = scatter_single_node(n, prop_filters, out_filters, nl, pooling, meta=m)
            output.extend(sc['out'])
            next_nodes.extend(sc['prop'])
            next_meta.extend(sc['meta'])
            
        nodes = next_nodes
        meta = next_meta

    output = _split_complex(output)

    return output

def _split_complex(arr):
    """Split complex arrays into their real and imaginary part.
    
    If an output image is complex, it is replaced by its real part, and its imaginary part gets
    added as seperate output image at the end of list of output images.
    """
    
    real = []
    
    for x in arr:
        if np.iscomplexobj(x):
            """
            Split complex arrays into their real and imaginary part.
            """
            real.append(np.real(x))
            real.append(np.imag(x))
        else:
            real.append(x)
    
    return real

def scatter_single_node(signal, propagation_filters, output_filters, nonlinearity, pooling, meta=None):
    """
    Do a single level scattering of the input image.
    
    signal	=	Input image.
    filters	=	List of linear filter functions performing their task in Fourier-space. Can be None.
    nonlinearity	=	Function mapping R onto R. Can be None.
    pooling	=	For dimension reduction. Can be None.
    
    meta = Optional meta information.
    
    Returns a dict containing values to be propagated and values for output.
        {'prop': [...], 'out': [...]}
    """
    in_shape = signal.shape
    
    
    # Space for filter outputs.
    filtered = []
    output = []
    new_meta = []
    
    """
    Apply each filter and nonlinearity
    """
    for f in propagation_filters:
        if f is None:
            filtered.append(signal)
        else:
            
            # get arguments of f.apply_filter
            f_args = inspect.getargspec(f.apply_filter).args
            
            if len(f_args) > 2:
                # Pass meta data only to filters that want it.
                res = f.apply_filter(signal, meta=meta)
            else:
                # Be nice to filters that can't handle meta information.
                res = f.apply_filter(signal)

            if type(res) is dict:
                filtered.extend(res['prop'])
                
                if 'out' in res:
                    output.extend(res['out'])
                
                """
                Store meta data if filter provides it.
                """
                if 'meta' in res:
                    new_meta.extend(res['meta']);
                else:
                    new_meta.extend([None]*len(res['prop']))
            else:
                """
                This filter returned an ordinary array.
                """
                filtered.extend(res)
    
    # Something has gone wrong if this does not hold:
    assert(len(filtered) == len(new_meta));
        
    # create output signals
    for f in output_filters:
        if f is None:
            output.append(signal)	# output unfiltered signal
        else:
            res = f.apply_filter(signal)
                
            if type(res) is dict:
                output.extend(res['prop'])	# TODO: Not nice enough.
                output.extend(res['out'])
            else:
                output.extend(res)

    
    if nonlinearity is not None:
        filtered = [nonlinearity(x) for x in filtered]
    
    if pooling is None:
        pooled = filtered
    else:
        pooled = [pooling.apply_pooling(x) for x in filtered]
        
    return {'prop': pooled, 'out': output, 'meta': new_meta}
