# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt
import librosa

from ssqueezepy import ssq_cwt, ssq_stft, stft, cwt, Wavelet
from ssqueezepy.experimental import scale_to_freq
from ssqueezepy.visuals import plot, imshow


IDX = 0
path = [
  r"vhs_noisy.mp3",
  r"vhs_clean.mp3",
  r"vhs_pure_noise.mp3",
][IDX]
offset = [
    0,
    20*22050,
    0
][IDX]
ecwt_factor = [
    1/500,
    1/500,
    1/500,
][IDX]
zoom_size = [
    16384*2,
    16384*2,
    16384*2
][IDX]
last_norm_scaling = [
    .7,
    .7,
    .5,
][IDX]
second_norm_scaling = [
    .3,
    .4,
    1,
][IDX]

#%%
x, fs = librosa.load(path)
print(len(x), fs)

N = len(x)
t = np.linspace(0, N/fs, N, 0)

zoom = zoom_size
x_slc = x[offset:offset+zoom]
t_slc = t[offset:offset+zoom]

#%%
wavelet = Wavelet()
Tsx, Sx, tsx_freqs, *_ = ssq_stft(x_slc-x_slc.mean(), fs=fs, n_fft=2048, window='hann')
Twx, Wx, twx_freqs, *_ = ssq_cwt(x_slc, wavelet, fs=fs, scales='log')
Tsx, Sx = Tsx[::-1], Sx[::-1]
Tsxo, Sxo = Tsx.copy(), Sx.copy()
Twxo, Wxo = Twx.copy(), Wx.copy()

#%%
escale = np.linspace(0, 1, len(Tsx), 1)[:, None][::-1]
escale_cwt = np.logspace(np.log10(ecwt_factor), np.log10(1), len(Wx), 1, base=10
                         )[:, None][::-1]
Tsx, Sx, Twx, Wx = [g.copy() for g in (Tsx, Sx, Twx, Wx)]
Tsx, Sx = Tsxo * escale, np.log10(1 + Sxo * escale)
Twx, Wx = Twxo * escale_cwt, Wxo * escale_cwt

#%%
fig, axes = plt.subplots(5, 1, layout='constrained', figsize=(7, 30))
pkw = dict(fig=fig)
ikw = dict(abs=1, fig=fig, show=0, xticks=t_slc)

# plot(t, x, **pkw, ax=axes[0])
plot(t_slc, x_slc, **pkw, ax=axes[0])
imshow(Sxo, **ikw, ax=axes[1], title="|STFT|", yticks=tsx_freqs[::-1])
imshow(Sx, **ikw, ax=axes[2],  title="|STFT|, 1/f^2 energy-scaled",
       yticks=tsx_freqs[::-1], norm_scaling=second_norm_scaling)
imshow(Wxo, **ikw, ax=axes[3], title="|CWT|", yticks=twx_freqs)
imshow(Wx, **ikw, ax=axes[4],  title="|CWT|,  1/f^2 energy-scaled",
       yticks=twx_freqs, xlabel="time [sec]", norm_scaling=last_norm_scaling)