from __future__ import annotations

import numpy as np

from pytfd_compat import spec, stft
from pytfd_compat.windows import get_window


def main() -> None:
    sample_rate = 48_000
    duration = 0.01
    t = np.arange(int(sample_rate * duration)) / sample_rate
    signal = np.sin(2 * np.pi * 6_000 * t) + 0.35 * np.sin(2 * np.pi * 11_000 * t)
    window = get_window("hanning", 63)

    spectrum = stft(signal, window, hop=8, n_fft=256)
    magnitude = spec(signal, window, hop=8, n_fft=256)

    print(f"signal shape: {signal.shape}")
    print(f"stft shape: {spectrum.shape}")
    print(f"magnitude range: {magnitude.min():.4f} .. {magnitude.max():.4f}")

    try:
        import matplotlib.pyplot as plt
    except ModuleNotFoundError:
        return

    plt.imshow(20 * np.log10(np.maximum(magnitude, 1e-8)), aspect="auto", origin="lower")
    plt.title("Basic STFT Example")
    plt.xlabel("Frame")
    plt.ylabel("Frequency Bin")
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()