## A428_martin
## David Martín Alaminos

from cdi_haar import *
from numpy import *
import matplotlib.pyplot as plt


'''
Daubechies functions
'''

r2 = sqrt(2)
r3 = sqrt(3)

d = 4*r2
a1 = (1+r3)/d
a2 = (3+r3)/d
a3 = (3-r3)/d
a4 = (1-r3)/d


#Daubechies trend
def D4trend(f):
    N = len(f)
    if N%2 != 0: return 'D4trend: %d is not divisible by 2' %N
    return [a1*f[2*j] + a2*f[2*j+1] + a3*f[(2*j+2)%N] + a4*f[(2*j+3)%N] \
    for j in range(N//2)]

#Daubechies fluctuation
def D4fluct(f):
    N = len(f)
    if N%2 != 0: return 'D4fluct: %d is not divisible by 2' %N
    return [a4*f[2*j] - a3*f[2*j+1] + a2*f[(2*j+2)%N] - a1*f[(2*j+3)%N] \
    for j in range(N//2)]

#D4 transform (first level)
def D4(f):
    N = len(f)
    if N%2 != 0: return 'D4: %d is not divisible by 2' %N
    return D4trend(f) + D4fluct(f)



'''
Draw the graph of D4(f) for your favorite signal f.
Compare the result with D2(f).
'''

def write(X,T,ha='center', va = 'center', fs='16', rot=0.0, color='k'):
    return plt.text(X[0],X[1], T, horizontalalignment=ha,verticalalignment=va,fontsize=fs,rotation=rot,color=color)


n = 10
N = 2**n

F = lambda x: 17*x**2 * ((1-x)**5)*cos(20*pi*x)
f = sample(F,N-1)

plt.close('all')
plt.figure("Comparison of D2 and D4 transforms")
plt.title("D2 and D4 transforms of function $f(x) = 17x^2 (1-x^{5}) cos(20 \pi x)$, $x \in [0..2^{%d}) $ \n" % n, color='black')
plt.xlim(0,N)

offset = 1
plt.plot(D4(f))
plt.plot([offset+x for x in HaarT(f)])
plt.plot([2*offset + x for x in f])

write((950,2*offset + 0.2),"$f(x)$", color='r')
write((950,offset+0.2),"$D2(x)$", color='g')
write((950,0.2),"$D4(x)$", color='b')

plt.show()