# -*- coding: utf-8 -*-
import numpy as np
from numpy.fft import fft, ifft, ifftshift
from ssqueezepy.utils import padsignal, process_scales
from ssqueezepy.algos import find_closest, indexed_sum
from ssqueezepy.visuals import imshow, plot, scat
from ssqueezepy import Wavelet

#%%# Helpers #################################################################
def _t(min, max, N):
    return np.linspace(min, max, N, endpoint=1)

def cos_f(freqs, N=128, phi=0):
    return np.concatenate([np.cos(2 * np.pi * f * (_t(i, i + 1, N) + phi))
                           for i, f in enumerate(freqs)])

#%%# Define signal & wavelet #################################################
fs = 129
x = cos_f([8], fs)
plot(x, title="Pure sine | N=129, f=8")
scat(x, show=1)

wavelet = Wavelet(('morlet', {'mu': 4}))

#%%# CWT #####################################################################
nv=32; dt=1/fs

n = len(x)  # store original length
x, nup, n1, n2 = padsignal(x, padtype='reflect')
x -= x.mean()
xh = fft(x)

scales = process_scales(scales='log:maximal', len_x=n, wavelet=wavelet, nv=nv)
pn = (-1)**np.arange(nup)

N_orig = wavelet.N
wavelet.N = nup

#%%# cwt ####
Psih = (wavelet(scale=scales, nohalf=False)).astype('complex128')
dPsih = (1j * wavelet.xi / dt) * Psih

Wx  = ifftshift(ifft(pn * Psih  * xh, axis=-1), axes=-1)
dWx = ifftshift(ifft(pn * dPsih * xh, axis=-1), axes=-1)
#%%#
wavelet.N = N_orig
# shorten to pre-padded size
Wx  = Wx[:,  n1:n1 + n]
dWx = dWx[:, n1:n1 + n]

#%%# Phase transform #########################################################
w = np.imag(dWx / Wx) / (2*np.pi)

# clean up tiny-valued Wx that have large `w` values; removing these makes
# no noticeable difference on `Tx` but allows us to see much better
w[np.abs(Wx) < np.abs(Wx).mean()] = 0

#%%# Reassignment frequencies (mapkind='maximal') ############################
na, N = Wx.shape
dT = dt * N
# normalized frequencies to map discrete-domain to physical:
#     f[[cycles/samples]] -> f[[cycles/second]]
# minimum measurable (fundamental) frequency of data
fm = 1 / dT
# maximum measurable (Nyquist) frequency of data
fM = 1 / (2 * dt)

ssq_freqs = fm * np.power(fM / fm, np.arange(na) / (na - 1))

#%%# Reassignment indices
# This step simply finds the index-equivalent of `w`. E.g., for given
# `ssq_freqs` ranging from 2 to 48, if w[5, 2] == 5, then
# `k[5, 2] = np.where(ssq_freqs == 5)` (or if no exact match, then closest to 5)
# `k` thus ranges from 0 to `len(ssq_freqs) - 1`.
k = find_closest(np.log2(w), np.log2(ssq_freqs))

#%%# Synchrosqueeze #########################################################
Tx = indexed_sum(Wx * np.log(2) / nv, k)
Tx = np.flipud(Tx)  # flip for visual aligned with `Wx`

#%%# Visualize ##############################################################
kw = dict(abs=1, cmap='jet', show=1, aspect='auto')
imshow(Wx, title="abs(CWT)", ylabel="scales", yticks=scales, **kw)
imshow(Tx, title="abs(SSQ_CWT)", ylabel="frequencies", yticks=ssq_freqs, **kw)
#%%# Zoom ####
a, b = 50, 158
c, d = 0, None
idxs = np.arange(len(scales)).astype('int64')
imshow(Wx[a:b, c:d], title="abs(CWT), zoomed", yticks=scales[a:b], **kw)
imshow(Tx[a:b, c:d], title="abs(SSQ_CWT), zoomed", yticks=idxs[a:b], **kw)

#%%
plot(w[81:120, 0], xticks=scales[81:120], show=1,
     title="w[81:120, 0] | Phase transform across zoomed scales")
#%%# Repeat for `w` ####
imshow(w, title="Phase transform | (min, max) = (%.3f, %.3f)" % (w.min(), w.max()),
       ylabel="scales", yticks=scales, **kw)
imshow(w[a:b, c:d], title="Phase transform, zoomed", yticks=scales[a:b], **kw)
#%%# zoom
wmn, wmx = w[w > 1e-3].min(), w[w > 1e-3].max()
# wmx += (wmx - wmn)*1.
imshow(w[a:b, c:d], title="Phase transform, magnitude-zoomed",
       yticks=scales[a:b], norm=(wmn, wmx), **kw)
#%%# Repeat for `k`
imshow(k, title=("Phase transform, index-equivalent, (min, max) = "
                 "(%d, %d)" % (k.min(), k.max())),
       ylabel="rows (scale indices)", yticks=idxs, **kw)
imshow(k[a:b, c:d], title="Phase transform, index-equivalent, zoomed",
       yticks=idxs[a:b], **kw)