
import numpy as np
import pywt	# python wavelets

class swt2:
    
    def __init__(self, wavelet='haar', level=1, start_level=0, frequency_decreasing_path=False):
        self._wavelet = wavelet
        self._level = level
        self._start_level = start_level
        self._frequency_decreasing_path = frequency_decreasing_path
        
    def apply_filter(self, img, meta=None):
        return _swt2(img, meta, self._wavelet, self._level, self._start_level, self._frequency_decreasing_path)


"""
pywt.swt2 returns

            [
                (cA_n,
                    (cH_n, cV_n, cD_n)
                ),
                (cA_n+1,
                    (cH_n+1, cV_n+1, cD_n+1)
                ),
                ...,
                (cA_n+level,
                    (cH_n+level, cV_n+level, cD_n+level)
                )
            ]

"""

def _swt2(signal, meta, wavelet, level=1, start_level=0, frequency_decreasing_path=False):
    """
    @param signal	The input signal.
    @param wavelet	The wavelet type.
    @param meta		Optional meta data.
    """

    tf = pywt.swt2(signal, pywt.Wavelet(wavelet), level=level, start_level=start_level)

    """
    tf =
                [
                    (cA_n,
                        (cH_n, cV_n, cD_n)
                    ),
                    (cA_n+1,
                        (cH_n+1, cV_n+1, cD_n+1)
                    ),
                    ...,
                    (cA_n+level,
                        (cH_n+level, cV_n+level, cD_n+level)
                    )
                ]
    """

    """
    Get the maximum scale to propagate.
    """
    if meta is not None and 'scale' in meta:
        # We get the scale from the previous filter.
        max_scale = meta['scale']
    else:
        # The previous filter was not polite enouth to provide us with that information.
        max_scale = len(tf)

    prop = []
    out = []
    new_meta = []

    for (cA, cHVD), scale in zip(tf, range(len(tf), 0, -1)):
        # scale takes values in: len(tf), ..., 3,2,1
        
        if not frequency_decreasing_path or scale <= max_scale:
            prop.extend(cHVD)
            m = {'scale': scale}
            new_meta.extend([m]*len(cHVD))
        
    out = [cA] # output approximate of last level

    return {'prop': prop, 'out': out, 'meta': new_meta}

