## A414
## Joan Ginés i Ametllé

'''
Compute the Haar transforms of level 2 and 3 of your favorite signal 
and draw their graphics as in the example on T6.7
'''

from math import *
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from fractions import Fraction

# Haar transform: trend(), fluct() and haar()
def trend(f):
    r2 = 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 = 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)
    h = haar(f,r-1); a = h[:m]
    return trend(a) + fluct(a) + h[m:]
    

# Utilities hline(), vline(), htick(), vtick() and write() from 
# cdi15-graphics-305.py. Some minor modifications are introduced

# hline draws a horizontal line fron (a,h) to (b,h)
def hline(a, b, h, lw=1, dash='k-'):
    x = np.arange(a, b+0.0001, b-a)
    return plt.plot(x, h+0*x, dash, lw=lw)
    
# vline draws a vertical line fron (d,a) to (d,b)
def vline(a, b, d, lw=1, dash='k-'):
    y = np.arange(a, b+0.0001, b-a)
    return plt.plot(d+0*y, y, dash, lw=lw)
    
# htick ticks and labels the vertical axis
def htick(X, T, run=0, lw='1', **kwargs):
    plt.plot([X[0]-run, X[0]+run], [X[1],X[1]], lw=lw, color='black')
    write((X[0]-2*run,X[1]), T, ha='left', fs='13', **kwargs)
    return None
    
# vtick ticks and labels the horizontal axis
def vtick(X, T, run=0, lw='1', **kwargs):
    plt.plot([X[0], X[0]], [X[1]-run,X[1]+run], lw=lw, color='black')
    write((X[0],X[1]-3*run), T, fs='13', **kwargs)
    return None
    
# utility to write text T at position X
def write(X, T, ha='center', va='center', fs='16', rot=0.0, **kwargs):
    return plt.text(X[0], X[1], T, horizontalalignment=ha, verticalalignment=va, fontsize=fs, rotation=rot, **kwargs)
    

# draw plot common features
def setplot():
    plt.axis('off')
    # title
    myargs1 = {'color':'b', 'size':'xx-large'}
    write((0.5, 0.8), r'$f(x) = 0.25\cdot x\cdot\sin{20x^2}$', **myargs1)
    # draw vertical markers
    for x in np.arange(0, 1.1, 0.1):
        vline(-0.7, 0.7, x, lw=0.15)
    # draw horizontal markers
    for y in np.arange(-0.7, 0.8, 0.1):
        hline(0, 1, y, lw=0.15)
    # draw points (1/2 and 1/4)
    plt.plot(0.5, 0, 'ko')
    plt.plot(0.25, 0, 'ko')
    # draw arrows (d^1 and d^2)
    arrow(0.25, -0.55, 0.25, 0, r'$d^2$')
    arrow(0.5, -0.55, 0.5, 0, r'$d^1$')
    
# write numbers in x-axis (r numbers in log2 scale)
def tick_ys(r, pos):
    if not r:
        # write '0'
        vtick((0, -0.6), '0')
    else: 
        # write current pos
        vtick((pos,-0.6), str(Fraction(pos)))
        # write next pos
        tick_ys(r-1, pos/2)
        
# write numbers in y-axis (2*pos/step equispaced numbers)
def tick_xs(step, pos):
    if not pos: 
        # write 0
        htick((1.02, pos), str(pos))
    else:
        # write current and symmetric pos
        htick((1.02, pos), str(Fraction(pos)))
        htick((1.02, -pos), str(Fraction(-pos)))
        # write next pos
        tick_xs(step, pos-step)
        
# plot an arrow at position (x,y) with length (dx,dy)
def arrow(x, y, dx, dy, text=''):
    plt.annotate('', xy=(x+dx,y+dy), xytext=(x,y), arrowprops=dict(arrowstyle="<->"))
    write((x+dx/2, y+0.05), text)
    
        
# define signal
def f(x): return sin(20*x**2)*(x/4)

X = np.arange(0, 1, 0.001)
Y = [f(x) for x in X]

plt.close('all')
myargs2 = {'color':'g', 'size':'xx-large'}

# plot haar transform of level 2
plt.figure('Haar transform level 2')
plt.plot(X, Y)
plt.plot(X, haar(Y,2))
setplot()

write((0.5, -0.85), r'$H_2(f)$: 2-level haar transform', **myargs2)
tick_xs(0.5, 0.5)
tick_ys(3, 1)
arrow(0, -0.55, 0.25, 0, '$a^2$')


# plot haar transform of level 3
plt.figure('Haar transform level 3')
plt.plot(X, Y)
plt.plot(X, haar(Y, 3))
setplot()

write((0.5, -0.85), r'$H_3(f)$: 3-level haar transform', **myargs2)
tick_xs(0.5, 0.5)
tick_ys(4, 1)
arrow(0, -0.55, 0.125, 0, '$a^3$')
arrow(0.125, -0.55, 0.125, 0, '$d^3$')