# L414 - Andrés Mingorance
from math import *
from numpy import *
from matplotlib.pyplot import *
from cdi import round

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 < 1 :
        if r < 0 : print("haar: invalid r. returning f")
        return f
    if r == 1: return trend(f) + fluct(f)
    a = haar(f, r-1)
    N = len(f)
    m = N // 2**(r-1)   # len(a^(r-1))
    h = a[:m]
    return haar(h) + a[m:]

def energy(f):
    return sum(t**2 for t in f)


############################################################

X = arange(0, 1, 0.0001)
def f(x): return 17*x**2 * ((1-x)**5)*cos(20*pi*x)

close('all')
for r in range(5) :
    h_r = haar(f(X),r)
    offset = -sqrt(2) ** (1+r)
    plot(list(map(lambda x:x+offset, h_r))) # offset them vertically a bit so they don't overlap
    text(6000, offset, "energy (lvl " + str(r) + ") = " + str(round(energy(h_r), 8)))
show()