"""Chosen rotor run-up: frequency versus order spectra (Python 3.14).
Run in the mirrored media directory; output is a PNG in the working directory.
"""
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

fs, duration = 2048.0, 4.0
t = np.arange(int(fs * duration)) / fs
rotation_hz = 10.0 + 2.5 * t
turns = 10.0 * t + 1.25 * t**2
phase = 2 * np.pi * turns
x = np.cos(3 * phase) + 0.35 * np.cos(7 * phase) + 0.2 * np.sin(2 * np.pi * 90 * t)
samples_per_turn = 128
uniform_turns = np.arange(0, turns[-1], 1 / samples_per_turn)
angle_signal = np.interp(uniform_turns, turns, x)

def spectrum(values, sample_rate):
    window = np.hanning(len(values))
    frequency = np.fft.rfftfreq(len(values), d=1 / sample_rate)
    amplitude = 2 * np.abs(np.fft.rfft(window * values)) / np.sum(window)
    return frequency, amplitude

frequency, time_amplitude = spectrum(x, fs)
orders, angle_amplitude = spectrum(angle_signal, samples_per_turn)
fig, axes = plt.subplots(1, 3, figsize=(10.4, 3.6), dpi=120, layout='constrained')
fig.patch.set_facecolor('white')
for ax in axes:
    ax.set_facecolor('white')
    ax.spines[['top', 'right']].set_visible(False)
    ax.grid(alpha=0.18)
axes[0].plot(t, 60 * rotation_hz, color='#306b87', lw=2)
axes[0].set(xlabel='Time (s)', ylabel='Shaft speed (rpm)', title='Chosen run-up', ylim=(570, 1230))
axes[0].text(0.05, 0.08, '3× and 7× shaft tones\nplus a fixed 90 Hz tone',
             transform=axes[0].transAxes, fontsize=9)
axes[1].plot(frequency, time_amplitude, color='#b26424', lw=1.2)
axes[1].set(xlim=(0, 160), ylim=(0, 0.27), xlabel='Frequency (Hz)',
            ylabel='Windowed amplitude', title='Uniform time samples')
axes[1].annotate('Fixed 90 Hz tone', (90, 0.20), xytext=(45, 0.24), fontsize=9,
                 arrowprops={'arrowstyle': '->', 'color': '0.35'})
axes[2].plot(orders, angle_amplitude, color='#624389', lw=1.2)
axes[2].set(xlim=(0, 12), ylim=(0, 1.12), xlabel='Order (cycles/revolution)',
            ylabel='Windowed amplitude', title='Uniform angular samples')
for order, y in [(3, 1.02), (7, 0.42)]:
    axes[2].text(order, y, f'{order}×', ha='center', fontsize=10)
axes[2].text(0.97, 0.90, '90 Hz spreads\nin order', ha='right',
             transform=axes[2].transAxes, fontsize=9)
fig.suptitle('Angular resampling concentrates speed-locked components', fontsize=12)
fig.savefig(Path('order-tracking-angle-resampling.png'), facecolor='white', transparent=False)
plt.close(fig)