from math import gcd

import numpy as np
from scipy.fft import rfft, rfftfreq
from scipy.signal import resample_poly

sample_rate = 800
target_rate = 1200
seconds = 0.25

sample_count = int(sample_rate * seconds)
t = np.arange(sample_count) / sample_rate
signal = np.sin(2 * np.pi * 40 * t) + 0.25 * np.sin(2 * np.pi * 120 * t)

factor = gcd(sample_rate, target_rate)
up = target_rate // factor
down = sample_rate // factor

resampled = resample_poly(signal, up=up, down=down)
t_resampled = np.arange(resampled.size) / target_rate


def dominant_frequency(values, rate):
    spectrum = np.abs(rfft(values))
    frequencies = rfftfreq(values.size, d=1 / rate)
    return frequencies[np.argmax(spectrum[1:]) + 1]


expected_samples = int(np.ceil(signal.size * up / down))

print(f"input samples: {signal.size} at {sample_rate} Hz")
print(f"output samples: {resampled.size} at {target_rate} Hz")
print(f"input duration: {signal.size / sample_rate:.3f} s")
print(f"output duration: {resampled.size / target_rate:.3f} s")
print(f"dominant frequency before: {dominant_frequency(signal, sample_rate):.1f} Hz")
print(f"dominant frequency after: {dominant_frequency(resampled, target_rate):.1f} Hz")
print("first five resampled values:", np.round(resampled[:5], 4).tolist())
print("length check:", resampled.size == expected_samples)
print("duration check:", np.isclose(t_resampled[-1] + 1 / target_rate, seconds))