## L414_sanchez
'''
Haar wavelets algorithms: continued
'''

from cdi import *

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

def fluct(f):
    r2 = 1.4142135623730951
    N = len(f)
    J = range(N//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)

 
##Testing
import math
import numpy as np
import matplotlib.pyplot as plt

def f(p): 
    l = []
    for x in p:
        l.append(5*x**2 * math.cos(x*math.pi))
    return np.array(l)

plt.close("all")
x = np.arange(0,10.0,0.01)
plt.figure("Haar transformation of a function")
plt.plot(x, haar(f(x),0),'k-', lw=2)
plt.plot(x, haar(f(x),1),'r-', lw=2)
plt.plot(x, haar(f(x),2),'g-', lw=2)
plt.plot(x, haar(f(x),3),'b-', lw=2)
plt.show()
print('--- Testing of energy ---')
print("Energy of Haar 0-level:",energy(haar(f(x),0)), ", color: Black")
print("Energy of Haar 1-level:",energy(haar(f(x),1)), ", color: Red")
print("Energy of Haar 2-level:",energy(haar(f(x),2)), ", color: Green")
print("Energy of Haar 3-level:",energy(haar(f(x),3)), ", color: Blue")