## L414 - cdi_wavelets.py
## Albert Puente Encinas

'''
Haar wavelets algorithms
'''

from math import *
from numpy import *
from matplotlib.pyplot import *

def trend(f):
    r2 = sqrt(2)
    J = range(len(f)//2)
    return [(f[2*j]+f[2*j+1])/r2 for j in J]
    
def fluct(f):
    r2 = sqrt(2)
    J = range(len(f)//2)
    return [(f[2*j]-f[2*j+1])/r2 for j in J]

def haar(f, r = 1):
    if r == 0: return f
    if r == 1: return (trend(f) + fluct(f))
    N = len(f)
    m = N // 2**(r-1)
    h = haar(f, r-1)
    a = h[:m]
    return trend(a) + fluct(a) + h[m:]
    
def energy(f):
    return sum(t**2 for t in f)

## Test

X = arange(0, 1, 0.001)

def f(x): return 20*x*x*((1-x)**5)*cos(11*pi*x)

close('all')
plot(f(X), color='k')
text(400, 0.6, 'Energies:', color='k')
text(400, 0.5, 'f(x): ' + str(energy(f(X))), color='k')

plot(haar(f(X)), color='r')
text(400, 0.4, 'haar(f(x), 1): ' + str(energy(haar(f(X), 1))), color='r')

plot(haar(f(X), 2), color='b')
text(400, 0.3, 'haar(f(x), 2): ' + str(energy(haar(f(X), 2))), color='b')

plot(haar(f(X), 3), color='g')
text(400, 0.2, 'haar(f(x), 3): ' + str(energy(haar(f(X), 3))), color='g')