from __future__ import annotations

import numpy as np

from pytfd_compat import pwd, sm, stft, wd
from pytfd_compat.windows import get_window


def main() -> None:
    sample_rate = 24_000
    t = np.arange(256) / sample_rate
    signal = np.sin(2 * np.pi * (2_000 + 4_000 * t) * t)
    analysis_window = get_window("gaussian", 31, a=1.5)
    smoothing_kernel = np.array([0.2, 0.6, 0.2])

    stft_matrix = stft(signal, analysis_window, hop=4, n_fft=256)
    wd_matrix = wd(signal)
    pwd_matrix = pwd(signal, analysis_window, n_fft=256)
    sm_matrix = sm(signal, analysis_window, smoothing_kernel, hop=4, n_fft=256)

    print("STFT:", stft_matrix.shape)
    print("WD:", wd_matrix.shape)
    print("PWD:", pwd_matrix.shape)
    print("SM:", sm_matrix.shape)


if __name__ == "__main__":
    main()