## L414
## AMA


from cdi import *
from math import *
import numpy as np
import matplotlib.pyplot as plt


def trend(f):
    r2 = 1.4142135623730951 # sqrt(2)
    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 # sqrt(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)  # // es la divisio amb part entera nomes
    h = haar(f,r-1)
    a = h[:m]  # aixo agafa la part desde l'inici de la llista fins a m
    return trend(a) + fluct(a) + h[m:]
    
    #a = trend(f)
    #d = fluct(f)
    #return a + d

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


## Tests

def fun(x): return (np.sin(np.pi * x) * 3*(x**(2)))

def test():
    plt.close("all")
    x = np.arange(0,1,0.0001)
    plt.figure("f(x) = sin(x)*(x**3) + 6*(x**2)", figsize=(8,10))
    plt.plot(x, fun(x),'r-', lw=2)
    plt.plot(x, haar(fun(x)), 'g-', lw=2)
    plt.plot(x, haar(fun(x),2), 'b-', lw=2)
    plt.plot(x, haar(fun(x),3), 'y-', lw=2)
    plt.show()
    
    print('Energy of f: ')
    e = energy(fun(x))
    print(e)
    print('Energy of H1: ')
    e1 = energy(haar(fun(x)))
    print(e1)
    print('Energy of H2: ')
    e2 = energy(haar(fun(x),2))
    print(e2)
    print('Energy of H3: ')
    e3 = energy(haar(fun(x),3))
    print(e3)
