import numpy as np
import matplotlib.pyplot as plt
import pywt

#%%###########################################################################
def allfft(x):
    coef = np.fft.fft(x)
    return coef, np.angle(coef), np.abs(coef)

def wplot(x, y=None, size=0, show=False, ax_equal=False, **kw):
    if y is None:
        plt.plot(x, **kw)
    else:
        plt.plot(x, y, **kw)
    _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)

def wscat(x, y=None, size=0, show=False, ax_equal=False, s=18, **kw):
    if y is None:
        plt.scatter(np.arange(len(x)), x, s=s, **kw)
    else:
        plt.scatter(x, y, s=s, **kw)
    _scale_plot(plt.gcf(), plt.gca(), size=size, show=show, ax_equal=ax_equal)

def imshow(data, norm=None, complex=None, show=1, **kw):
    kw['interpolation'] = kw.get('interpolation', 'none')
    if norm is None:
        mx = np.max(np.abs(data))
        vmin, vmax = -mx, mx
    else:
        vmin, vmax = norm

    if (complex is None and np.sum(np.abs(np.imag(data))) < 1e-8) or (
            complex is False):
        plt.imshow(np.real(data), cmap='bwr', vmin=vmin, vmax=vmax, **kw)
    else:
        fig, axes = plt.subplots(1, 2)
        axes[0].imshow(data.real, cmap='bwr', vmin=vmin, vmax=vmax, **kw)
        axes[1].imshow(data.imag, cmap='bwr', vmin=vmin, vmax=vmax, **kw)

    plt.subplots_adjust(left=0, right=1, bottom=0, top=1, wspace=0, hspace=0)
    if show:
        plt.show()


def _scale_plot(fig, ax, size=True, show=False, ax_equal=False):
    xmin, xmax = ax.get_xlim()
    rng = xmax - xmin
    ax.set_xlim(xmin + .018 * rng, xmax - .018 * rng)
    if size:
        fig.set_size_inches(15, 7)
    if ax_equal:
        yabsmax = max(np.abs([*ax.get_ylim()]))
        mx = max(yabsmax, max(np.abs([xmin, xmax])))
        ax.set_xlim(-mx, mx)
        ax.set_ylim(-mx, mx)
        fig.set_size_inches(8, 8)
    if show:
        plt.show()
#%%###########################################################################
def morlet(sigma=5):
    cs = (1 + np.exp(-sigma ** 2) - 2 * np.exp(-.75 * sigma ** 2)) ** (-.5)
    ks = np.exp(-.5 * sigma ** 2)
    return lambda t: np.exp(-.5 * t ** 2) * (np.exp(1j * sigma * t) - ks
                                             ) * cs * np.pi**(-.25)

def morlet_kernel(N, minmax=(-8, 8), sigma=5, real=False):
    t = np.linspace(*minmax, N, False)
    if real:
        return lambda scale: np.real(morlet(sigma)(t / scale))
    return lambda scale: morlet(sigma)(t / scale)

def cwt_window(N, name, real=False):
    if name == 'morlet':
        return morlet_kernel(N, real=real)

def _transform(x, kernel, scales):
    coef = np.zeros(len(scales), dtype='complex128')

    for i, scale in enumerate(scales):
        psi = np.conj(kernel(scale))
        coef[i] = np.sum(x * psi / np.sqrt(scale))
    return coef

def _scales(N, nv=32):
    noct = np.log2(N) - 1
    n_scales = int(noct * nv)
    return np.power(2 ** (1 / nv), np.arange(1, n_scales + 1))
#%%###########################################################################
def cwt(x, win_len=None, win_inc=None, win='morlet', viz_scale=0, real=False):
    N = len(x)
    win_len = win_len or N // 8
    win_inc = win_inc or win_len
    n_wins = (N - win_len) // win_inc + 1

    scales = _scales(N, nv=32)
    coef = np.zeros((n_wins, len(scales)), dtype='complex128')
    kernel = cwt_window(win_len, win, real=real)

    for tau in range(n_wins):
        start = tau * win_inc
        end   = start + win_len
        coef[tau, :] = _transform(x[start:end], kernel, scales)
        if viz_scale:
            cwt_viz(x, kernel(viz_scale), start, end)
    return coef.T


def cwt2(x, win_len=None, win='morlet', real=False):
    N = len(x)
    win_len = win_len or N // 8

    scales = _scales(N, nv=32)
    coef = np.zeros((len(scales), N), dtype='complex128')
    kernel = cwt_window(win_len, win, real=real)
    wl2 = win_len // 2

    for i, scale in enumerate(scales):
        coef[i, :] = np.convolve(x, kernel(scale)[::-1])[wl2:-(wl2 - 1)]
    return coef


def cwt_viz(x, psi, start, end):
    wplot(x)
    # plt.axvline(start, color='k', linestyle='--')
    # plt.axvline(end,   color='k', linestyle='--')
    wplot(np.pad(psi.real, [start, len(x) - end]), color='purple', show=1)
    plt.show()

def _t(min, max, N):
    return np.linspace(min, max, N, False)

def cos_f(freqs, N=128):
    return np.concatenate([np.cos(2 * np.pi * f * _t(i, i + 1, N))
                           for i, f in enumerate(freqs)])
#%%###########################################################################
N = 128
x = cos_f([1, 4], N)
coef, theta, r = allfft(x)

wplot(x); wscat(x, s=8, show=1)
wplot(coef.real); wplot(coef.imag)
#%%###########################################################################
cwt(x, win_len=96, win_inc=20, real=1, viz_scale=2)
coef  = cwt(x, win_len=96, win_inc=1, real=1)
coef2 = cwt2(x, win_len=96, win='morlet', real=1)
coefp = pywt.cwt(x, _scales(len(x)), 'morl', method='conv')[0]
#%%
imshow(coef)
imshow(coef2)
imshow(coefp)