#!/usr/bin/env python
from __future__ import division, print_function, absolute_import

import numpy as np
from numpy.testing import (run_module_suite, assert_allclose, assert_,
                           assert_raises)

import pywt

# Check that float32 and complex64 are preserved.  Other real types get
# converted to float64.
dtypes_in = [np.int8, np.float32, np.float64, np.complex64, np.complex128]
dtypes_out = [np.float64, np.float32, np.float64, np.complex64, np.complex128]


def test_dwt_idwt_basic():
    x = [3, 7, 1, 1, -2, 5, 4, 6]
    cA, cD = pywt.dwt(x, 'db2')
    cA_expect = [5.65685425, 7.39923721, 0.22414387, 3.33677403, 7.77817459]
    cD_expect = [-2.44948974, -1.60368225, -4.44140056, -0.41361256,
                 1.22474487]
    assert_allclose(cA, cA_expect)
    assert_allclose(cD, cD_expect)

    x_roundtrip = pywt.idwt(cA, cD, 'db2')
    assert_allclose(x_roundtrip, x, rtol=1e-10)


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

        cA, cD = pywt.dwt(x, wavelet)
        assert_(cA.dtype == cD.dtype == dt_out, "dwt: " + errmsg)

        x_roundtrip = pywt.idwt(cA, cD, wavelet)
        assert_(x_roundtrip.dtype == dt_out, "idwt: " + errmsg)


def test_dwt_idwt_basic_complex():
    x = np.asarray([3, 7, 1, 1, -2, 5, 4, 6])
    x = x + 0.5j*x
    cA, cD = pywt.dwt(x, 'db2')
    cA_expect = np.asarray([5.65685425, 7.39923721, 0.22414387, 3.33677403,
                            7.77817459])
    cA_expect = cA_expect + 0.5j*cA_expect
    cD_expect = np.asarray([-2.44948974, -1.60368225, -4.44140056, -0.41361256,
                            1.22474487])
    cD_expect = cD_expect + 0.5j*cD_expect
    assert_allclose(cA, cA_expect)
    assert_allclose(cD, cD_expect)

    x_roundtrip = pywt.idwt(cA, cD, 'db2')
    assert_allclose(x_roundtrip, x, rtol=1e-10)


def test_dwt_wavelet_kwd():
    x = np.array([3, 7, 1, 1, -2, 5, 4, 6])
    w = pywt.Wavelet('sym3')
    cA, cD = pywt.dwt(x, wavelet=w, mode='constant')
    cA_expect = [4.38354585, 3.80302657, 7.31813271, -0.58565539, 4.09727044,
                 7.81994027]
    cD_expect = [-1.33068221, -2.78795192, -3.16825651, -0.67715519,
                 -0.09722957, -0.07045258]
    assert_allclose(cA, cA_expect)
    assert_allclose(cD, cD_expect)


def test_dwt_coeff_len():
    x = np.array([3, 7, 1, 1, -2, 5, 4, 6])
    w = pywt.Wavelet('sym3')
    ln = pywt.dwt_coeff_len(data_len=len(x), filter_len=w.dec_len, mode='symmetric')
    assert_(ln == 6)
    ln_modes = [pywt.dwt_coeff_len(len(x), w.dec_len, mode) for mode in
                pywt.Modes.modes]
    assert_allclose(ln_modes, [6, 6, 6, 6, 6, 4])


def test_idwt_none_input():
    # None input equals arrays of zeros of the right length
    res1 = pywt.idwt([1, 2, 0, 1], None, 'db2', 'symmetric')
    res2 = pywt.idwt([1, 2, 0, 1], [0, 0, 0, 0], 'db2', 'symmetric')
    assert_allclose(res1, res2, rtol=1e-15, atol=1e-15)

    res1 = pywt.idwt(None, [1, 2, 0, 1], 'db2', 'symmetric')
    res2 = pywt.idwt([0, 0, 0, 0], [1, 2, 0, 1], 'db2', 'symmetric')
    assert_allclose(res1, res2, rtol=1e-15, atol=1e-15)

    # Only one argument at a time can be None
    assert_raises(ValueError, pywt.idwt, None, None, 'db2', 'symmetric')


def test_idwt_invalid_input():
    # Too short, min length is 4 for 'db4':
    assert_raises(ValueError, pywt.idwt, [1, 2, 4], [4, 1, 3], 'db4', 'symmetric')


def test_dwt_single_axis():
    x = [[3, 7, 1, 1],
         [-2, 5, 4, 6]]

    cA, cD = pywt.dwt(x, 'db2', axis=-1)

    cA0, cD0 = pywt.dwt(x[0], 'db2')
    cA1, cD1 = pywt.dwt(x[1], 'db2')

    assert_allclose(cA[0], cA0)
    assert_allclose(cA[1], cA1)

    assert_allclose(cD[0], cD0)
    assert_allclose(cD[1], cD1)


def test_idwt_single_axis():
    x = [[3, 7, 1, 1],
         [-2, 5, 4, 6]]

    cA, cD = pywt.dwt(x, 'db2', axis=-1)

    x0 = pywt.idwt(cA[0], cD[0], 'db2', axis=-1)
    x1 = pywt.idwt(cA[1], cD[1], 'db2', axis=-1)

    assert_allclose(x[0], x0)
    assert_allclose(x[1], x1)


def test_dwt_axis_arg():
    x = [[3, 7, 1, 1],
         [-2, 5, 4, 6]]

    cA_, cD_ = pywt.dwt(x, 'db2', axis=-1)
    cA, cD = pywt.dwt(x, 'db2', axis=1)

    assert_allclose(cA_, cA)
    assert_allclose(cD_, cD)


def test_idwt_axis_arg():
    x = [[3, 7, 1, 1],
         [-2, 5, 4, 6]]

    cA, cD = pywt.dwt(x, 'db2', axis=1)

    x_ = pywt.idwt(cA, cD, 'db2', axis=-1)
    x = pywt.idwt(cA, cD, 'db2', axis=1)

    assert_allclose(x_, x)


def test_dwt_idwt_axis_excess():
    x = [[3, 7, 1, 1],
         [-2, 5, 4, 6]]
    # can't transform over axes that aren't there
    assert_raises(ValueError,
                  pywt.dwt, x, 'db2', 'symmetric', axis=2)

    assert_raises(ValueError,
                  pywt.idwt, [1, 2, 4], [4, 1, 3], 'db2', 'symmetric', axis=1)


if __name__ == '__main__':
    run_module_suite()
