# -*- coding: UTF-8 -*-
# 作者：huanhuilong
# 标题：Python 傅立叶变换
# 描述：傅立叶变换是可逆的，请在正向变换+反向变换后打印，曲线应该完全重合

import numpy as np
from scipy.fft import fft, fftfreq, ifft
import matplotlib.pyplot as plt


def wave(duration, sample_rate, freq):
    x = np.linspace(0, duration, sample_rate*duration, endpoint=False)
    w = np.sin((2 * np.pi) * freq*x)
    return x, w


def transform(duration, sample_rate, w):
    N = sample_rate*duration
    xf = fftfreq(N, 1/sample_rate)
    wf = fft(w)
    return xf, wf


def inverse_transform(duration, sample_rate, wf):
    w = ifft(wf)
    return w


if __name__ == "__main__":

    duration = 3
    sample_rate = 30000
    freq = 4

    # 正玄波
    x, w = wave(duration, sample_rate, freq)

    # 傅立叶变换
    xf, wf = transform(duration, sample_rate, w)

    # 傅立叶逆变换
    w2 = inverse_transform(duration, sample_rate, wf)

    # 应该完全重合
    plt.plot(w, 'o')
    plt.plot(w2)

    plt.show()
