import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import ShortTimeFFT

sample_rate = 800.0
duration = 2.0
sample_count = int(sample_rate * duration)
t = np.arange(sample_count) / sample_rate

signal = np.where(
    t < 1.0,
    np.sin(2 * np.pi * 50 * t),
    np.sin(2 * np.pi * 150 * t),
)

nperseg = 128
noverlap = 96

sft = ShortTimeFFT.from_window(
    "hann",
    fs=sample_rate,
    nperseg=nperseg,
    noverlap=noverlap,
    scale_to="psd",
)

power = sft.spectrogram(signal)
times = sft.t(signal.size)
freqs = sft.f

early_columns = (times >= 0.1) & (times < 0.8)
late_columns = (times > 1.2) & (times <= 1.9)
early_profile = power[:, early_columns].mean(axis=1)
late_profile = power[:, late_columns].mean(axis=1)

power_db = 10 * np.log10(np.maximum(power, 1e-12))

fig, ax = plt.subplots(figsize=(7, 4))
mesh = ax.pcolormesh(times, freqs, power_db, shading="auto")
ax.set_title("SciPy ShortTimeFFT spectrogram")
ax.set_xlabel("Time (s)")
ax.set_ylabel("Frequency (Hz)")
ax.set_ylim(0, 220)
fig.colorbar(mesh, ax=ax, label="PSD (dB)")
fig.tight_layout()
fig.savefig("spectrogram.png", dpi=150)
plt.close(fig)

print(f"signal samples: {signal.size}")
print(f"frequency bins: {freqs.size}")
print(f"time slices: {times.size}")
print(f"spectrogram shape: {power.shape}")
print(f"dominant early frequency: {freqs[np.argmax(early_profile)]:.1f} Hz")
print(f"dominant late frequency: {freqs[np.argmax(late_profile)]:.1f} Hz")
print("saved plot: spectrogram.png")