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
#%%##########################################################################
def l1(x):
    return np.sum(np.abs(x))

def l2(x):
    return np.sqrt(np.sum(np.abs(x ** 2)))

def plot(scales, y1, y2, logx=1):
    scales = np.log2(scales) if logx else scales
    if logx == 'lin':
        scales = np.arange(len(scales))
    plt.plot(scales, y1)
    plt.plot(scales, y2)
    plt.show()
#%%##########################################################################
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 = []

l1res, l2res = [], []
l1rec, l2rec = [], []
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 /= scale

    l1res.append(l1(np.diff(int_psi_scale)))
    l2res.append(l2(np.diff(int_psi_scale) * np.sqrt(scale)))

    l1rec.append(l1(w))
    l2rec.append(l2(w * np.sqrt(scale)))

#%%##########################################################################
k = None
plot(scales[:k], l1res[:k], l2res[:k], logx=1)
plot(scales[:k], l1rec[:k], l2rec[:k], logx=1)