import numpy as np
import scipy.signal as signal
import scipy.fftpack as fft
import wave

def stretch(signal, rate, window_size, h):
    # Параметры алгоритма
    factor = rate
    window_size = int(window_size)
    hop_size = int(window_size // h)

    # Фазовая информация
    phase = np.zeros(window_size)
    hanning_window = np.hanning(window_size)

    # Получение фазовой разности
    def get_phase_diff(x_fft, y_fft):
        angle_x = np.angle(x_fft)
        angle_y = np.angle(y_fft)
        return angle_y - angle_x

    # Реалзация алгоритма phase vocoder
    resampled = signal
    time_steps = np.arange(0, resampled.size - window_size, hop_size) / float(resampled.size)
    w = 2 * np.pi * signal.hamming(window_size)

    # Разбиваем аудиозапись на окна
    for t in time_steps:
        # Применяем оконную функцию и получаем FFT
        fft_window = fft.fft(w * resampled[int(t * resampled.size): int(t * resampled.size) + window_size])
        fft_window = np.abs(fft_window) * np.exp(1j * phase)

        # Получаем амплитуду и фазу нового сигнала
        ifft_window = fft.ifft(fft_window).real
        ifft_window *= hanning_window
        resampled[int(t * resampled.size) + hop_size: int(t * resampled.size) + window_size] += ifft_window[hop_size:]

        # Обновляем фазовую информацию
        phase_diff = get_phase_diff(fft_window, fft.fft(w * resampled[int(t * resampled.size) + hop_size: int(t * resampled.size) + window_size + hop_size]))
        phase_diff = np.angle(phase_diff)
        phase = (phase + phase_diff) % (2 * np.pi)

    # Применяем растяжение/сжатие
    resampled = signal.resample(resampled, int(len(resampled) // factor))
    return resampled.astype(np.float32)

def change_speed(filename, rate):
    # Чтение wav-файла
    with wave.open(filename, 'rb') as wave_file:
        signal = np.frombuffer(wave_file.readframes(-1), dtype='int16')

    # Определение параметров алгоритма
    window_size = 2**13
    h = 4

    # Применение алгоритма phase vocoder
    stretched_signal = stretch(signal, rate, window_size, h)

    # Запись нового wav-файла
    with wave.open('output.wav', 'wb') as wave_file:
        wave_file.setparams(wave_file.getparams())
        wave_file.writeframes(stretched_signal.tobytes())