## A428
## Joan Ginés i Ametllé

import sympy as sy
from math import *
from numpy import *
from cdi_haar import *
from matplotlib.pyplot import *

# Daubechies transform (D4)
r2 = sy.sqrt(2); r3 = sy.sqrt(3)
d = 4*r2

a1 = (1+r3)/d; a2 = (3+r3)/d
a3 = (3-r3)/d; a4 = (1-r3)/d

def D4trend(f):
    N = len(f)
    if N%2: 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)]
    
def D4fluct(f):
    N = len(f)
    if N%2: 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)]
    
def D4(f):
    N = len(f)
    if N%2: return 'D4 transform: Error'
    return D4trend(f) + D4fluct(f)
    
# signal to plot
F = lambda x: cos(32*atan(1)*x)*sin(8*atan(1)*x**2)

X = arange(0.0, 1.0, 0.001)
Y = [F(x) for x in X]

close('all')
fig, axes = subplots(nrows=3, ncols=1)
fig.tight_layout()

# plot original signal
subplot(3,1,1, title = "Original signal")
plot(X, Y, color='r', lw = 1)
grid('on')

# plot haar transform
D2Y = haar(Y)
subplot(3,1,2, title = "D2 transform")
plot(X, D2Y, color='b', lw = 1)
grid('on')
 
# plot daubechies transform
D4Y = D4(Y)
subplot(3,1,3, title = "D4 transform")
plot(X, D4Y, color = 'g', lw = 1)
grid('on')

# both D2 and D4 produce similar results for the given input