import numpy as np
import matplotlib.pyplot as plt
import scipy.signal as sig

from pywt._extensions._pywt import DiscreteContinuousWavelet
from pywt._functions import integrate_wavelet

#%%##########################################################################
scales = np.power(2 ** (1 / 32), np.arange(1, 256 + 1))

wavelet = DiscreteContinuousWavelet('morl')
int_psi, x = integrate_wavelet(wavelet, precision=10)

linrange = max(x) - min(x)
step = linrange / len(x)  # == x[1] - x[0]
#%%##########################################################################
mae = []
for scale in scales:
    j = np.arange(scale * linrange + 1) / (scale * step)
    j = j.astype(int)  # floor
    if j[-1] >= int_psi.size:
        j = np.extract(j < int_psi.size, j)
    int_psi_scale = int_psi[j][::-1].real

    Ns = len(int_psi_scale)
    w = np.real(sig.morlet2(Ns, Ns / 16) * np.sqrt(Ns / 16) * np.pi**(.25))
    w_int = np.cumsum(w)[::-1] / scale

    mae.append(np.mean(np.abs(w_int - int_psi_scale)))
#%%##########################################################################
plt.plot(np.log2(scales), mae)
plt.title("MAE vs log2(scale)", fontsize=16, fontweight='bold')
plt.show()

plt.plot(np.log2(scales), np.log10(mae) * 20)
plt.title("MAE [dB] vs log2(scale)", fontsize=16, fontweight='bold')
plt.show()