##L414
'''
Haar wavelets alorithm
'''

def trend(f) : #f is a numerical vector
    r2 = 1.414213562370951 #square root of 2
    N = len(f)
    J = range(N//2)
    return[(f[2*j] + f[2*j+1]) / r2 for j in J]
    
def fluct(f) : #f is a numerical vector
    r2 = 1.414213562370951 #square root of 2
    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) #Divided as integer
    h =  haar(f,r-1); a = h[:m]
    return trend(a) + fluct(a) + h[m:] #concatenation
    
def energy(f) :
    return sum (t**2 for t in f)

##Tests
import numpy as np
import matplotlib.pyplot as plt
from math import *
pi = 3.141592653589793
x = np.arange(0,1,1/(2**9))
def func(x) : return (30*(x**2))*((1-x)**6)*np.sin(pi*5*x)

print('Energy of original function: ', energy(func(x)))
h1 = haar(func(x),1)
print('Energy of first level Haar transformation: ',energy(h1))
h2 = haar(func(x),2)
print('Energy of second level Haar transformation: ', energy(h2))
h3 = haar(func(x),3)
print('Energy of third level Haar transformation: ',energy(h3))

plt.figure("30*(x^2)*((1-x)^6)*sin(5*pi*x)")
plt.plot(x,func(x),'b-', lw=2)
plt.plot(x,haar(func(x),1),'g-', lw=2)
plt.plot(x,haar(func(x),2),'r-', lw=2)
#plt.plot(x,haar(func(x),3),'y-', lw=2)