"""A short-time Fourier transform spectrogram in decibels."""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import scipy.fft


@dataclass
class Spectrogram:
    times: np.ndarray
    """Center time of each column [s]."""
    freqs: np.ndarray
    """Frequency of each row [Hz]."""
    db: np.ndarray
    """Magnitude in dB relative to the peak, ``(len(freqs), len(times))``."""

    def plot(self, ax=None, **kwargs):
        """Draws the spectrogram with matplotlib; see ``cymsim.plot.spectrogram``."""
        from .plot import spectrogram

        return spectrogram(self, ax=ax, **kwargs)


def spectrogram(signal: np.ndarray, fs: float, window: int = 2048, hop: int = 256) -> Spectrogram:
    """The STFT of ``signal`` with a periodic Hann window of ``window`` samples
    and a hop of ``hop`` samples, in dB relative to the largest magnitude."""
    x = np.asarray(signal, dtype=float)
    if len(x) < window:
        x = np.pad(x, (0, window - len(x)))
    n_cols = (len(x) - window) // hop + 1
    idx = np.arange(window)[None, :] + hop * np.arange(n_cols)[:, None]
    win = 0.5 - 0.5 * np.cos(2 * np.pi * np.arange(window) / window)
    X = scipy.fft.rfft(x[idx] * win[None, :], axis=1)
    p = X.real**2 + X.imag**2
    db = 10 * np.log10(p + 1e-30)
    db -= db.max()
    times = (hop * np.arange(n_cols) + window / 2) / fs
    freqs = np.arange(window // 2 + 1) * fs / window
    return Spectrogram(times, freqs, db.T)