#!/usr/bin/env python
import os
import pickle
import numpy as np
from numpy.testing import assert_allclose, assert_

import pywt


def test_wavelet_properties():
    w = pywt.Wavelet('db3')

    # Name
    assert_(w.name == 'db3')
    assert_(w.short_family_name == 'db')
    assert_(w.family_name, 'Daubechies')

    # String representation
    fields = ('Family name', 'Short name', 'Filters length', 'Orthogonal',
              'Biorthogonal', 'Symmetry')
    for field in fields:
        assert_(field in str(w))

    # Filter coefficients
    dec_lo = [0.03522629188210, -0.08544127388224, -0.13501102001039,
              0.45987750211933, 0.80689150931334, 0.33267055295096]
    dec_hi = [-0.33267055295096, 0.80689150931334, -0.45987750211933,
              -0.13501102001039, 0.08544127388224, 0.03522629188210]
    rec_lo = [0.33267055295096, 0.80689150931334, 0.45987750211933,
              -0.13501102001039, -0.08544127388224, 0.03522629188210]
    rec_hi = [0.03522629188210, 0.08544127388224, -0.13501102001039,
              -0.45987750211933, 0.80689150931334, -0.33267055295096]
    assert_allclose(w.dec_lo, dec_lo)
    assert_allclose(w.dec_hi, dec_hi)
    assert_allclose(w.rec_lo, rec_lo)
    assert_allclose(w.rec_hi, rec_hi)

    assert_(len(w.filter_bank) == 4)

    # Orthogonality
    assert_(w.orthogonal)
    assert_(w.biorthogonal)

    # Symmetry
    assert_(w.symmetry)

    # Vanishing moments
    assert_(w.vanishing_moments_phi == 0)
    assert_(w.vanishing_moments_psi == 3)


def test_wavelet_coefficients():
    families = ('db', 'sym', 'coif', 'bior', 'rbio')
    wavelets = sum([pywt.wavelist(name) for name in families], [])
    for wavelet in wavelets:
        if (pywt.Wavelet(wavelet).orthogonal):
            check_coefficients_orthogonal(wavelet)
        elif(pywt.Wavelet(wavelet).biorthogonal):
            check_coefficients_biorthogonal(wavelet)
        else:
            check_coefficients(wavelet)


def check_coefficients_orthogonal(wavelet):

    epsilon = 5e-11
    level = 5
    w = pywt.Wavelet(wavelet)
    phi, psi, x = w.wavefun(level=level)

    # Lowpass filter coefficients sum to sqrt2
    res = np.sum(w.dec_lo)-np.sqrt(2)