## A414 - 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)

# A414
prec = 0.0001
X = arange(0.0, 1.0, prec)

def f(x): 
    return 4*power(x, 2)*exp(-256*power(x, 8)*power(sin(40*x), 2))

# Auxiliar arrow function
def doubleArrow(x, x2, y, y2, text, ax, offset):
    annotate('', xy=(x, y), xycoords='data', xytext=(x2, y2), textcoords='data',
        arrowprops={'arrowstyle': '<->'})
    ax.text((x+x2)/2, (y+y2)/2 + offset, text, color='k', fontsize=15)

    
close('all')

# Initial options HAAR Level 2
figure("Haar transform of level 2", facecolor = 'white')

ax = axes(frameon=True)
ax.get_xaxis().set_ticks([0, 1/4, 1/2, 1])
ax.get_xaxis().set_ticklabels([r'$0$', r'$\frac{1}{4}$', 
    r'$\frac{1}{2}$', r'$1$'], fontsize = 18)
grid(True)

# Functions 
plot(X, f(X), color='#005200', lw = 1)
plot(X, haar(f(X), 2), color='#104E8B',lw = 1)

# Points
x = [0, 1/4, 1/2, 1]
plot(x, [0, 0, 0, 0], 'ko')

# Arrows    
doubleArrow(x[0], x[1], -1/2, -1/2, r'$a^2$', ax, -0.4)
doubleArrow(x[1], x[2], -1/2, -1/2, r'$d^2$', ax, -0.4)
doubleArrow(x[2], x[3], -1/2, -1/2, r'$d^1$', ax, -0.4)

# Text
text(0.4, 6, r'$f(x) =4x^2 e^{-256x^8 \sin^2(40x)}$', color='#005200', fontsize=22)
text(0.4, 5, r'$Haar(f(x), 2)$', color='#104E8B', fontsize=22)

# Initial options HAAR Level 3
figure("Haar transform of level 3", facecolor = 'white')

ax = axes(frameon=True)
ax.get_xaxis().set_ticks([0, 1/8, 1/4, 1/2, 1])
ax.get_xaxis().set_ticklabels([r'$0$', r'$\frac{1}{8}$', r'$\frac{1}{4}$', 
    r'$\frac{1}{2}$', r'$1$'], fontsize = 18)
grid(True)

# Functions 
plot(X, f(X), color='#005200', lw = 1)
plot(X, haar(f(X), 3), color='#104E8B',lw = 1)

# Points
x = [0, 1/8, 1/4, 1/2, 1]
plot(x, [0, 0, 0, 0, 0], 'ko')

# Arrows    
doubleArrow(x[0], x[1], -1, -1, r'$a^3$', ax, -0.7)
doubleArrow(x[1], x[2], -1, -1, r'$d^3$', ax, -0.7)
doubleArrow(x[2], x[3], -1, -1, r'$d^2$', ax, -0.7)
doubleArrow(x[3], x[4], -1, -1, r'$d^1$', ax, -0.7)

# Text
text(0.4, 8, r'$f(x) =4x^2 e^{-256x^8 \sin^2(40x)}$', color='#005200', fontsize=22)
text(0.4, 7, r'$Haar(f(x), 3)$', color='#104E8B', fontsize=22)