## A428 - Andrés Mingorance

import sympy as s

r2 = s.sqrt(2)
r3 = s.sqrt(3)
simplify = s.simplify
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 != 0 : raise Exception("D4trend/ " + str(N) +" is not divisible by 2")
    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 != 0 : raise Exception("D4fluct/ " + str(N) +" is not divisible by 2")
    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) :
    return D4trend(f) + D4fluct(f)

####################################################################################

from cdi_haar import haar as D2
import numpy as np
import cdi
N = 1024
F = lambda x: 15**x**2 * 2*(1-x)**4 * np.cos(9*np.pi*x)
f = cdi.sample(F, N-1)

import matplotlib.pyplot as plt
d4 = D4(f)
d2 = D2(f)
plt.figure("D4(f) vs D2(f)")
plt.subplot(2, 2, 1)
plt.title('D4(f) [Daubechies]')
plt.plot(d4, 'b')
plt.xlim(0,N)
plt.grid()

plt.subplot(2, 2, 2)
plt.title('D2(f) [Haar]')
plt.plot(d2, 'r')
plt.xlim(0,N)
plt.grid()

plt.subplot(2, 1, 2)
plt.tight_layout()
plt.title('Overlapped: $f(x) = 15^{x^2}2(1-x)^4 cos(9\pi x)$')
plt.plot(f, 'g', label='Original function f')
plt.plot(d4, 'b', label='D4(f)')
plt.plot(d2, 'r', label='D2(f)')
plt.xlim(0,N)
plt.grid()
plt.legend()

plt.show()