<commit_before># -*- coding: utf-8 -*-

# Copyright (c) 2006-2012 <NAME> <http://en.ig.ma/>
# See COPYING for license details.

"""
2D Discrete Wavelet Transform and Inverse Discrete Wavelet Transform.
"""

from __future__ import division, print_function, absolute_import

__all__ = ['dwt2', 'idwt2', 'swt2', 'dwtn']

from itertools import cycle

import numpy as np

from ._pywt import Wavelet, MODES
from ._pywt import dwt, idwt, swt, downcoef


def dwt2(data, wavelet, mode='sym'):
    """
    2D Discrete Wavelet Transform.

    Parameters
    ----------
    data : ndarray
        2D array with input data
    wavelet : Wavelet object or name string
        Wavelet to use
    mode : str, optional
        Signal extension mode, see MODES (default: 'sym')

    Returns
    -------
    (cA, (cH, cV, cD)) : tuple
        Approximation, horizontal detail, vertical detail and diagonal
        detail coefficients respectively.

    Examples
    --------
    >>> from pywt import multidim
    >>> data = np.ones((4,4), dtype=np.float64)
    >>> coeffs = multidim.dwt2(data, 'haar')
    >>> cA, (cH, cV, cD) = coeffs
    >>> cA
    array([[ 2.,  2.],
           [ 2.,  2.]])
    >>> cV
    array([[ 0.,  0.],
           [ 0.,  0.]])

    """
    data = np.asarray(data)
    if len(data.shape) != 2:
        raise ValueError("Expected 2D data array")

    if not isinstance(wavelet, Wavelet):
        wavelet = Wavelet(wavelet)

    mode = MODES.from_object(mode)

    # filter rows
    H, L = [], []
    for row in data:
        cA, cD = dwt(row, wavelet, mode)
        L.append(cA)
        H.append(cD)

    # filter columns
    H = np.transpose(H)
    L = np.transpose(L)

    LL, LH = [], []
    for row in L:
        cA, cD = dwt(np.array(row, np.float64), wavelet, mode)
        LL.append(cA)
        LH.append(cD)

    HL, HH = [], []
    for row in H:
        cA, cD = dwt(np.array(row, np.float64), wavelet, mode)
        HL.append(cA)
        HH.append(cD)

    # build result structure: (approx, (horizontal, vertical, diagonal))
    ret = (np.transpose(LL), (np.transpose(LH), np.transpose(HL), np.transpose(HH)))

    return ret


def idwt2(coeffs, wavelet, mode='sym'):
    """
    2D Inverse Discrete Wavelet Transform. Reconstruct data from coefficients
    arrays.

    Parameters
    ----------
    coeffs : tuple
        (cA, (cH, cV, cD)) A tuple with approximation coefficients and three
        details coefficients 2D arrays like from `dwt2()`
    wavelet : Wavelet object or name string
        Wavelet to use
    mode : str, optional
        Signal extension mode, see MODES (default: 'sym')

    Examples
    --------
    >>> from pywt import multidim
    >>> data = np.array([[1,2], [3,4]], dtype=np.float64)
    >>> coeffs = multidim.dwt2(data, 'haar')
    >>> multidim.idwt2(coeffs, 'haar')
    array([[ 1.,  2.],
           [ 3.,  4.]])

    """
    if len(coeffs) != 2 or len(coeffs[1]) != 3:
        raise ValueError("Invalid coeffs param")

    # L -low-pass data, H - high-pass data
    LL, (LH, HL, HH) = coeffs

    if LL is not None:
        LL = np.transpose(LL)
    if LH is not None:
        LH = np.transpose(LH)
    if HL is not None:
        HL = np.transpose(HL)
    if HH is not None:
        HH = np.transpose(HH)

    all_none = True
    for arr in (LL, LH, HL, HH):
        if arr is not None:
            all_none = False
            if len(arr.shape) != 2:
                raise TypeError("All input coefficients arrays must be 2D.")

    if all_none:
        raise ValueError(
            "At least one input coefficients array must not be None.")

    if not isinstance(wavelet, Wavelet):
        wavelet = Wavelet(wavelet)

    mode = MODES.from_object(mode)

    # idwt columns
    L = []
    if LL is None and LH is None:
        L = None
    else:
        if LL is None:
            # IDWT can handle None input values - equals to zero-array
            LL = cycle([None])
        if LH is None:
            # IDWT can handle None input values - equals to zero-array
            LH = cycle([None])
        for rowL, rowH in zip(LL, LH):
            L.append(idwt(rowL, rowH, wavelet, mode, 1))

    H = []
    if HL is None and HH is None:
        H = None
    else:
        if HL is None:
            # IDWT can handle None input values - equals to zero-array
            HL = cycle([None])
        if HH is None:
            # IDWT can handle None input values - equals to zero-array
            HH = cycle([None])
        for rowL, rowH in zip(HL, HH):
            H.append(idwt(rowL, rowH, wavelet, mode, 1))

    if L is not None:
        L = np.transpose(L)
    if H is not None:
        H = np.transpose(H)

    # idwt rows
    data = []
    if L is None:
        # IDWT can handle None input values - equals to zero-array
        L = cycle([None])
    if H is None:
        # IDWT can handle None input values - equals to zero-array
        H = cycle([None])
    for rowL, rowH in zip(L, H):
        data.append(idwt(rowL, rowH, wavelet, mode, 1))

    return np.array(data, np.float64)


def _downcoef(data, wavelet, mode, type):
    """Adapts pywt.downcoef call for apply_along_axis"""
    return downcoef(type, data, wavelet, mode, level=1)


def dwtn(data, wavelet, mode='sym'):
    """
    Single-level n-dimensional Discrete Wavelet Transform.

    Parameters
    ----------
    data : ndarray
        nD array with input data
    wavelet : Wavelet object or name string
        Wavelet to use
    mode : str, optional
        Signal extension mode, see MODES (default: 'sym')

    Returns
    -------
    coeffs : dict
        Results are arranged in a dictionary, where key specifies
        the transform type on each dimension and value is a n-dimensional
        coefficients array.

        For example, for a 2D case the result will look something like this::

            {
                'aa': <coeffs>  # A(LL) - approx. on 1st dim, approx. on 2nd dim
                'ad': <coeffs>  # H(LH) - approx. on 1st dim, det. on 2nd dim
                'da': <coeffs>  # V(HL) - det. on 1st dim, approx. on 2nd dim
                'dd': <coeffs>  # D(HH) - det. on 1st dim, det. on 2nd dim
            }

    """
    data = np.asarray(data)
    dim = len(data.shape)
    coeffs = [('', data)]
    for axis in range(dim):
        new_coeffs = []
        for subband, x in coeffs:
            new_coeffs.extend([
                (subband + 'a', np.apply_along_axis(_downcoef, axis,
                 x, wavelet, mode, 'a')),
                (subband + 'd', np.apply_along_axis(_downcoef, axis,
                 x, wavelet, mode, 'd'))])

        coeffs = new_coeffs
    return dict(coeffs)


def swt2(data, wavelet, level, start_level=0):
    """
    2D Stationary Wavelet Transform.

    Parameters
    ----------
    data : ndarray<commit_msg>Bring dwt2 into agreement with dwtn docstring<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2006-2012 <NAME> <http://en.ig.ma/>
# See COPYING for license details.

"""
2D Discrete Wavelet Transform and Inverse Discrete Wavelet Transform.
"""

from __future__ import division, print_function, absolute_import

__all__ = ['dwt2', 'idwt2', 'swt2', 'dwtn']

from itertools import cycle

import numpy as np

from ._pywt import Wavelet, MODES
from ._pywt import dwt, idwt, swt, downcoef


def dwt2(data, wavelet, mode='sym'):
    """
    2D Discrete Wavelet Transform.

    Parameters
    ----------
    data : ndarray
        2D array with input data
    wavelet : Wavelet object or name string
        Wavelet to use
    mode : str, optional
        Signal extension mode, see MODES (default: 'sym')

    Returns
    -------
    (cA, (cH, cV, cD)) : tuple
        Approximation, horizontal detail, vertical detail and diagonal
        detail coefficients respectively.

    Examples
    --------
    >>> from pywt import multidim
    >>> data = np.ones((4,4), dtype=np.float64)
    >>> coeffs = multidim.dwt2(data, 'haar')
    >>> cA, (cH, cV, cD) = coeffs
    >>> cA
    array([[ 2.,  2.],
           [ 2.,  2.]])
    >>> cV
    array([[ 0.,  0.],
           [ 0.,  0.]])

    """
    data = np.asarray(data)
    if len(data.shape) != 2:
        raise ValueError("Expected 2D data array")

    if not isinstance(wavelet, Wavelet):
        wavelet = Wavelet(wavelet)

    mode = MODES.from_object(mode)

    # filter rows
    H, L = [], []
    for row in data:
        cA, cD = dwt(row, wavelet, mode)
        L.append(cA)
        H.append(cD)

    # filter columns
    H = np.transpose(H)
    L = np.transpose(L)

    LL, HL = [], []
    for row in L:
        cA, cD = dwt(np.array(row, np.float64), wavelet, mode)
        LL.append(cA)
        HL.append(cD)

    LH, HH = [], []
    for row in H:
        cA, cD = dwt(np.array(row, np.float64), wavelet, mode)
        LH.append(cA)
        HH.append(cD)

    # build result structure: (approx, (horizontal, vertical, diagonal))
    ret = (np.transpose(LL), (np.transpose(HL), np.transpose(LH), np.transpose(HH)))

    return ret


def idwt2(coeffs, wavelet, mode='sym'):
    """
    2D Inverse Discrete Wavelet Transform. Reconstruct data from coefficients
    arrays.

    Parameters
    ----------
    coeffs : tuple
        (cA, (cH, cV, cD)) A tuple with approximation coefficients and three
        details coefficients 2D arrays like from `dwt2()`
    wavelet : Wavelet object or name string
        Wavelet to use
    mode : str, optional
        Signal extension mode, see MODES (default: 'sym')

    Examples
    --------
    >>> from pywt import multidim
    >>> data = np.array([[1,2], [3,4]], dtype=np.float64)
    >>> coeffs = multidim.dwt2(data, 'haar')
    >>> multidim.idwt2(coeffs, 'haar')
    array([[ 1.,  2.],
           [ 3.,  4.]])

    """
    if len(coeffs) != 2 or len(coeffs[1]) != 3:
        raise ValueError("Invalid coeffs param")

    # L -low-pass data, H - high-pass data
    LL, (LH, HL, HH) = coeffs

    if LL is not None:
        LL = np.transpose(LL)
    if LH is not None:
        LH = np.transpose(LH)
    if HL is not None:
        HL = np.transpose(HL)
    if HH is not None:
        HH = np.transpose(HH)

    all_none = True
    for arr in (LL, LH, HL, HH):
        if arr is not None:
            all_none = False
            if len(arr.shape) != 2:
                raise TypeError("All input coefficients arrays must be 2D.")

    if all_none:
        raise ValueError(
            "At least one input coefficients array must not be None.")

    if not isinstance(wavelet, Wavelet):
        wavelet = Wavelet(wavelet)

    mode = MODES.from_object(mode)

    # idwt columns
    L = []
    if LL is None and LH is None:
        L = None
    else:
        if LL is None:
            # IDWT can handle None input values - equals to zero-array
            LL = cycle([None])
        if LH is None:
            # IDWT can handle None input values - equals to zero-array
            LH = cycle([None])
        for rowL, rowH in zip(LL, LH):
            L.append(idwt(rowL, rowH, wavelet, mode, 1))

    H = []
    if HL is None and HH is None:
        H = None
    else:
        if HL is None:
            # IDWT can handle None input values - equals to zero-array
            HL = cycle([None])
        if HH is None:
            # IDWT can handle None input values - equals to zero-array
            HH = cycle([None])
        for rowL, rowH in zip(HL, HH):
            H.append(idwt(rowL, rowH, wavelet, mode, 1))

    if L is not None:
        L = np.transpose(L)
    if H is not None:
        H = np.transpose(H)

    # idwt rows
    data = []
    if L is None:
        # IDWT can handle None input values - equals to zero-array
        L = cycle([None])
    if H is None:
        # IDWT can handle None input values - equals to zero-array
        H = cycle([None])
    for rowL, rowH in zip(L, H):
        data.append(idwt(rowL, rowH, wavelet, mode, 1))

    return np.array(data, np.float64)


def _downcoef(data, wavelet, mode, type):
    """Adapts pywt.downcoef call for apply_along_axis"""
    return downcoef(type, data, wavelet, mode, level=1)


def dwtn(data, wavelet, mode='sym'):
    """
    Single-level n-dimensional Discrete Wavelet Transform.

    Parameters
    ----------
    data : ndarray
        nD array with input data
    wavelet : Wavelet object or name string
        Wavelet to use
    mode : str, optional
        Signal extension mode, see MODES (default: 'sym')

    Returns
    -------
    coeffs : dict
        Results are arranged in a dictionary, where key specifies
        the transform type on each dimension and value is a n-dimensional
        coefficients array.

        For example, for a 2D case the result will look something like this::

            {
                'aa': <coeffs>  # A(LL) - approx. on 1st dim, approx. on 2nd dim
                'ad': <coeffs>  # V(LH) - approx. on 1st dim, det. on 2nd dim
                'da': <coeffs>  # H(HL) - det. on 1st dim, approx. on 2nd dim
                'dd': <coeffs>  # D(HH) - det. on 1st dim, det. on 2nd dim
            }

    """
    data = np.asarray(data)
    dim = len(data.shape)
    coeffs = [('', data)]
    for axis in range(dim):
        new_coeffs = []
        for subband, x in coeffs:
            new_coeffs.extend([
                (subband + 'a', np.apply_along_axis(_downcoef, axis,
                 x, wavelet, mode, 'a')),
                (subband + 'd', np.apply_along_axis(_downcoef, axis,
                 x, wavelet, mode, 'd'))])

        coeffs = new_coeffs
    return dict(coeffs)


def swt2(data, wavelet, level, start_level=0):
    """
    2D Stationary Wavelet Transform.

    Parameters
    ----------
    data : ndarray
{'fraction_non_alphanumeric': 0.12501829893134242, 'fraction_numerical': 0.010906163080076123, 'mean_word_length': 1.9338630019325747, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 75, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5239409', 'n_tokens_mistral': 4716, 'n_tokens_neox': 4420, 'n_words': 1477}