def test_per_axis_wavelets_and_modes():
    # tests seperate wavelet and edge mode for each axis.
    rstate = np.random.RandomState(1234)
    data = rstate.randn(16, 16, 16)

    # wavelet can be a string or wavelet object
    wavelets = (pywt.Wavelet('haar'), 'sym2', 'db4')

    # mode can be a string or a Modes enum
    modes = ('symmetric', 'periodization',
             pywt._extensions._pywt.Modes.reflect)

    coefs = pywt.dwtn(data, wavelets, modes)
    assert_allclose(pywt.idwtn(coefs, wavelets, modes), data, atol=1e-14)

    coefs = pywt.dwtn(data, wavelets[:1], modes)
    assert_allclose(pywt.idwtn(coefs, wavelets[:1], modes), data, atol=1e-14)

    coefs = pywt.dwtn(data, wavelets, modes[:1])
    assert_allclose(pywt.idwtn(coefs, wavelets, modes[:1]), data, atol=1e-14)

    # length of wavelets or modes doesn't match the length of axes
    assert_raises(ValueError, pywt.dwtn, data, wavelets[:2])
    assert_raises(ValueError, pywt.dwtn, data, wavelets, mode=modes[:2])
    assert_raises(ValueError, pywt.idwtn, coefs, wavelets[:2])
    assert_raises(ValueError, pywt.idwtn, coefs, wavelets, mode=modes[:2])

    # dwt2/idwt2 also support per-axis wavelets/modes
    data2 = data[..., 0]
    coefs2 = pywt.dwt2(data2, wavelets[:2], modes[:2])
    assert_allclose(pywt.idwt2(coefs2, wavelets[:2], modes[:2]), data2,
                    atol=1e-14)

def test_negative_axes():
    data = np.array([[0, 1, 2, 3],
                     [1, 1, 1, 1],
                     [1, 4, 2, 8]])
    coefs1 = pywt.dwtn(data, 'haar', axes=(1, 1))
    coefs2 = pywt.dwtn(data, 'haar', axes=(-1, -1))
    assert_equal(coefs1, coefs2)

    rec1 = pywt.idwtn(coefs1, 'haar', axes=(1, 1))
    rec2 = pywt.idwtn(coefs1, 'haar', axes=(-1, -1))
    assert_equal(rec1, rec2)

def test_idwtn_mixed_complex_dtype():
    rstate = np.random.RandomState(0)
    x = rstate.randn(8, 8, 8)
    x = x + 1j*x
    coeffs = pywt.dwtn(x, 'db2')

    x_roundtrip = pywt.idwtn(coeffs, 'db2')
    assert_allclose(x_roundtrip, x, rtol=1e-10)

    # mismatched dtypes OK
    coeffs['a' * x.ndim] = coeffs['a' * x.ndim].astype(np.complex64)
    x_roundtrip2 = pywt.idwtn(coeffs, 'db2')
    assert_allclose(x_roundtrip2, x, rtol=1e-7, atol=1e-7)
    assert_(x_roundtrip2.dtype == np.complex128)

def test_idwtn_none_coeffs():
    data = np.array([[0, 1, 2, 3],
                     [1, 1, 1, 1],
                     [1, 4, 2, 8]])
    data = data + 1j*data  # test with complex data
    coefs = pywt.dwtn(data, 'haar', axes=(1, 1))

    # verify setting coefficients to None is the same as zeroing them
    coefs['dd'] = np.zeros_like(coefs['dd'])
    result_zeros = pywt.idwtn(coefs, 'haar', axes=(1, 1))

    coefs['dd'] = None
    result_none = pywt.idwtn(coefs, 'haar', axes=(1, 1))

    assert_equal(result_zeros, result_none)

def test_wavelet_decomposition3d_and_reconstruction3d():
    # Test 3D wavelet decomposition and reconstruction and verify that
    # they perform as expected
    x = np.random.rand(16, 16, 16)
    mode = 'sym'
    wbasis = pywt.Wavelet('db5')
    nscales = 1
    wavelet_coeffs = wavelet_decomposition3d(x, wbasis, mode, nscales)
    aaa = wavelet_coeffs[0]
    reference = pywt.dwtn(x, wbasis, mode)
    aaa_reference = reference['aaa']
    assert all_almost_equal(aaa, aaa_reference)

    reconstruction = wavelet_reconstruction3d(wavelet_coeffs, wbasis, mode,
                                              nscales)
    reconstruction_reference = pywt.idwtn(reference, wbasis, mode)
    assert all_almost_equal(reconstruction, reconstruction_reference)
    assert all_almost_equal(reconstruction, x)
    assert all_almost_equal(reconstruction_reference, x)

    wbasis = pywt.Wavelet('db1')
    nscales = 3
    wavelet_coeffs = wavelet_decomposition3d(x, wbasis, mode, nscales)
    shape_true = (nscales + 1, )
    assert all_equal(np.shape(wavelet_coeffs), shape_true)

    reconstruction = wavelet_reconstruction3d(wavelet_coeffs, wbasis, mode,
                                              nscales)
    assert all_almost_equal(reconstruction, x)

def test_idwtn_axes():
    data = np.array([[0, 1, 2, 3],
                     [1, 1, 1, 1],
                     [1, 4, 2, 8]])
    data = data + 1j*data  # test with complex data
    coefs = pywt.dwtn(data, 'haar', axes=(1, 1))
    assert_allclose(pywt.idwtn(coefs, 'haar', axes=(1, 1)), data, atol=1e-14)

def test_ignore_invalid_keys():
    data = np.array([[0, 4, 1, 5, 1, 4], [0, 5, 6, 3, 2, 1], [2, 5, 19, 4, 19, 1]])

    wavelet = pywt.Wavelet("haar")

    LL, (HL, LH, HH) = pywt.dwt2(data, wavelet)
    d = {"aa": LL, "da": HL, "ad": LH, "dd": HH, "foo": LH, "a": HH}

    assert_allclose(pywt.idwt2((LL, (HL, LH, HH)), wavelet), pywt.idwtn(d, wavelet), atol=1e-15)

def test_idwtn_take():
    data = np.array([
        [[1, 4, 1, 5, 1, 4],
         [0, 5, 6, 3, 2, 1],
         [2, 5, 19, 4, 19, 1]],
        [[1, 5, 1, 2, 3, 4],
         [7, 12, 6, 52, 7, 8],
         [5, 2, 6, 78, 12, 2]]])
    wavelet = pywt.Wavelet('haar')
    d = pywt.dwtn(data, wavelet)

    assert_(data.shape != pywt.idwtn(d, wavelet).shape)
    assert_allclose(data, pywt.idwtn(d, wavelet, take=data.shape), atol=1e-15)

    # Check shape for take not equal to data.shape
    data = np.random.randn(51, 17, 68)
    d = pywt.dwtn(data, wavelet)
    assert_equal((2, 2, 2), pywt.idwtn(d, wavelet, take=2).shape)
    assert_equal((52, 18, 68), pywt.idwtn(d, wavelet, take=0).shape)

def test_dwdtn_idwtn_allwavelets():
    rstate = np.random.RandomState(1234)
    r = rstate.randn(16, 16)
    # test 2D case only for all wavelet types
    wavelist = pywt.wavelist()
    if 'dmey' in wavelist:
        wavelist.remove('dmey')
    for wavelet in wavelist:
        for mode in pywt.Modes.modes:
            coeffs = pywt.dwtn(r, wavelet, mode=mode)
            assert_allclose(pywt.idwtn(coeffs, wavelet, mode=mode),
                            r, rtol=1e-7, atol=1e-7)

def test_3D_reconstruct():
    data = np.array(
        [
            [[0, 4, 1, 5, 1, 4], [0, 5, 26, 3, 2, 1], [5, 8, 2, 33, 4, 9], [2, 5, 19, 4, 19, 1]],
            [[1, 5, 1, 2, 3, 4], [7, 12, 6, 52, 7, 8], [2, 12, 3, 52, 6, 8], [5, 2, 6, 78, 12, 2]],
        ]
    )

    wavelet = pywt.Wavelet("haar")
    for mode in pywt.Modes.modes:
        d = pywt.dwtn(data, wavelet, mode=mode)
        assert_allclose(data, pywt.idwtn(d, wavelet, mode=mode), rtol=1e-13, atol=1e-13)

def test_dwtn_idwtn_dtypes():
    wavelet = pywt.Wavelet('haar')
    for dt_in, dt_out in zip(dtypes_in, dtypes_out):
        x = np.ones((4, 4), dtype=dt_in)
        errmsg = "wrong dtype returned for {0} input".format(dt_in)

        coeffs = pywt.dwtn(x, wavelet)
        for k, v in coeffs.items():
            assert_(v.dtype == dt_out, "dwtn: " + errmsg)

        x_roundtrip = pywt.idwtn(coeffs, wavelet)
        assert_(x_roundtrip.dtype == dt_out, "idwtn: " + errmsg)

def test_ignore_invalid_keys():
    data = np.array([
        [0, 4, 1, 5, 1, 4],
        [0, 5, 6, 3, 2, 1],
        [2, 5, 19, 4, 19, 1]])

    wavelet = pywt.Wavelet('haar')

    LL, (HL, LH, HH) = pywt.dwt2(data, wavelet)
    d = {'aa': LL, 'da': HL, 'ad': LH, 'dd': HH,
         'foo': LH, 'a': HH}

    assert_allclose(pywt.idwt2((LL, (HL, LH, HH)), wavelet),
                    pywt.idwtn(d, wavelet), atol=1e-15)

def test_idwtn_missing():
    # Test to confirm missing data behave as zeroes
    data = np.array([
        [0, 4, 1, 5, 1, 4],
        [0, 5, 6, 3, 2, 1],
        [2, 5, 19, 4, 19, 1]])

    wavelet = pywt.Wavelet('haar')

    LL, (HL, _, HH) = pywt.dwt2(data, wavelet)
    d = {'aa': LL, 'da': HL, 'dd': HH}

    assert_allclose(pywt.idwt2((LL, (HL, None, HH)), wavelet),
                    pywt.idwtn(d, 'haar'), atol=1e-15)

def test_3D_reconstruct_complex():
    # All dimensions even length so `take` does not need to be specified
    data = np.array(
        [
            [[0, 4, 1, 5, 1, 4], [0, 5, 26, 3, 2, 1], [5, 8, 2, 33, 4, 9], [2, 5, 19, 4, 19, 1]],
            [[1, 5, 1, 2, 3, 4], [7, 12, 6, 52, 7, 8], [2, 12, 3, 52, 6, 8], [5, 2, 6, 78, 12, 2]],
        ]
    )
    data = data + 1j

    wavelet = pywt.Wavelet("haar")
    d = pywt.dwtn(data, wavelet)
    # idwtn creates even-length shapes (2x dwtn size)
    original_shape = [slice(None, s) for s in data.shape]
    assert_allclose(data, pywt.idwtn(d, wavelet)[original_shape], rtol=1e-13, atol=1e-13)

def test_idwtn_idwt2_complex():
    data = np.array([
        [0, 4, 1, 5, 1, 4],
        [0, 5, 6, 3, 2, 1],
        [2, 5, 19, 4, 19, 1]])
    data = data + 1j
    wavelet = pywt.Wavelet('haar')

    LL, (HL, LH, HH) = pywt.dwt2(data, wavelet)
    d = {'aa': LL, 'da': HL, 'ad': LH, 'dd': HH}

    for mode in pywt.Modes.modes:
        assert_allclose(pywt.idwt2((LL, (HL, LH, HH)), wavelet, mode=mode),
                        pywt.idwtn(d, wavelet, mode=mode),
                        rtol=1e-14, atol=1e-14)

def test_idwtn_idwt2():
    data = np.array([[0, 4, 1, 5, 1, 4], [0, 5, 6, 3, 2, 1], [2, 5, 19, 4, 19, 1]])

    wavelet = pywt.Wavelet("haar")

    LL, (HL, LH, HH) = pywt.dwt2(data, wavelet)
    d = {"aa": LL, "da": HL, "ad": LH, "dd": HH}

    for mode in pywt.Modes.modes:
        assert_allclose(
            pywt.idwt2((LL, (HL, LH, HH)), wavelet, mode=mode),
            pywt.idwtn(d, wavelet, mode=mode),
            rtol=1e-14,
            atol=1e-14,
        )

def test_dwdtn_idwtn_allwavelets():
    rstate = np.random.RandomState(1234)
    r = rstate.randn(16, 16)
    # test 2D case only for all wavelet types
    wavelist = pywt.wavelist()
    if 'dmey' in wavelist:
        wavelist.remove('dmey')
    for wavelet in wavelist:
        if wavelet in ['cmor', 'shan', 'fbsp']:
            # skip these CWT families to avoid warnings
            continue
        if isinstance(pywt.DiscreteContinuousWavelet(wavelet), pywt.Wavelet):
            for mode in pywt.Modes.modes:
                coeffs = pywt.dwtn(r, wavelet, mode=mode)
                assert_allclose(pywt.idwtn(coeffs, wavelet, mode=mode),
                                r, rtol=1e-7, atol=1e-7)

def test_idwtn_missing():
    # Test to confirm missing data behave as zeroes
    data = np.array([[0, 4, 1, 5, 1, 4], [0, 5, 6, 3, 2, 1], [2, 5, 19, 4, 19, 1]])

    wavelet = pywt.Wavelet("haar")

    coefs = pywt.dwtn(data, wavelet)

    # No point removing zero, or all
    for num_missing in range(1, len(coefs)):
        for missing in combinations(coefs.keys(), num_missing):
            missing_coefs = coefs.copy()
            for key in missing:
                del missing_coefs[key]
            LL = missing_coefs.get("aa", None)
            HL = missing_coefs.get("da", None)
            LH = missing_coefs.get("ad", None)
            HH = missing_coefs.get("dd", None)

            assert_allclose(pywt.idwt2((LL, (HL, LH, HH)), wavelet), pywt.idwtn(missing_coefs, "haar"), atol=1e-15)

reconstructed = pywt.idwt2(coeffs2, 'bior1.3')
fig = plt.figure()
plt.imshow(reconstructed, interpolation="nearest", cmap=plt.cm.gray)

# Check that reconstructed image is close to the original
np.testing.assert_allclose(original, reconstructed, atol=1e-13, rtol=1e-13)


# Now do the same with dwtn/idwtn, to show the difference in their signatures

coeffsn = pywt.dwtn(original, 'bior1.3')
fig = plt.figure()
for i, key in enumerate(['aa', 'ad', 'da', 'dd']):
    ax = fig.add_subplot(2, 2, i + 1)
    ax.imshow(coeffsn[key], origin='image', interpolation="nearest",
              cmap=plt.cm.gray)
    ax.set_title(titles[i], fontsize=12)

fig.suptitle("dwtn coefficients", fontsize=14)

# Now reconstruct and plot the original image
reconstructed = pywt.idwtn(coeffsn, 'bior1.3')
fig = plt.figure()
plt.imshow(reconstructed, interpolation="nearest", cmap=plt.cm.gray)

# Check that reconstructed image is close to the original
np.testing.assert_allclose(original, reconstructed, atol=1e-13, rtol=1e-13)


plt.show()

def test_idwtn_axes_subsets():
    data = np.array(np.random.standard_normal((4, 4, 4, 4)))
    # test all combinations of 3 out of 4 axes transformed
    for axes in combinations((0, 1, 2, 3), 3):
        coefs = pywt.dwtn(data, 'haar', axes=axes)
        assert_allclose(pywt.idwtn(coefs, 'haar', axes=axes), data, atol=1e-14)