#範例12-4：把雜訊濾波：D.T.F.T：使用scipy做離散非週期性數據的傅立葉轉換（FFT， Fast Fourier Transform)，有範例
#目的：我們給定一些以sin(x)為基礎的訊號，但隨機加上干擾噪音雜訊，當轉成頻譜後，把雜訊波消除，然後再inverse回原波
from sympy import *
import matplotlib.pyplot as plt
from scipy import fftpack
import numpy as np
#(1).每隔0.05秒測量一點，共200個數據
time_step = 0.05
time_vec = np.arange(0, 10, time_step)
#print(time_vec)

period = 5
#print(2*np.pi*time_vec/period) 
#(2).以sin(2πt/p)為基礎的訊號，但隨機加上干擾噪音雜訊
sig = (np.sin(2*np.pi*time_vec/period) + 0.25*np.random.randn(time_vec.size))
#print(sig)
#print(np.round(sig,2))

#(3).把訊號用FFT轉成頻譜sig_fft複數
sig_fft = fftpack.fft(sig)
Amplitude = np.abs(sig_fft)
Power = Amplitude**2
Angle = np.angle(sig_fft)

#(4).計算出頻率軸(由sig_fft複數來計算)（不需要再shift重新排序了，因為plot時會自動調整）
sample_freq = fftpack.fftfreq(sig.size, d= time_step)
#print('sig.size=',sig.size,'len(sig)=',len(sig))
#print(Amplitude)
#print(sample_freq)

#(5).計算頻譜圖中的最大振幅處的頻率值
#原理：因為原訊號是sin(x)波，轉成頻域後，只會出現一個頻率的振幅，其它都是雜訊
#方法一：稍微複雜
Amp_Freq = np.array([Amplitude, sample_freq])
#Amp_Freq[0] = Amplitude
#Amp_Freq[1] = sample_freq
#numpy指令：argmax()：求出陣列最大值的所在索引index
#Amp_Freq[0,:].argmax()：表示在Amplitude裡面找出最大振幅的所在索引index
Amp_position = Amp_Freq[0,:].argmax()
#然後再到#Amp_Freq[1] （=sample_freq）裡面去查它對應的頻率值
peak_freq = Amp_Freq[1, Amp_position]

#方法二：更簡單更直觀的方法
#Amp_position = Amplitude.argmax()
#peak_freq = sample_freq[Amp_position]
print('Amp_position=', Amp_position)
print('peak_freq=', peak_freq)

#(6).濾波：把最大振幅波頻率以外的全部的波，其振幅都設定為0（沒有雜訊）
#最大漲幅波頻率=0.2， -0.2，所以f > 0.2 and f <-0.2的波，其sig_fft，都設定為0
high_freq_fft = sig_fft.copy()
#不需要迴圈，直接對陣列進行判別與設定：arrayb1[i > 5] =0
#目的：把sig_fft裡面，所有頻率
high_freq_fft[np.abs(sample_freq) > peak_freq] = 0
#print(high_freq_fft)

#(7).反向推算原函數（由已經濾波後的fft=high_freq_fft來反推）
filter_sig = fftpack.ifft(high_freq_fft)

#(8).畫圖: 點數據list
#用matplotlib可以直接畫處list數據
#在同一圖fig，畫出二條線的方法：最後再plt.show()即可
plt.subplot(2,1,1)
plt.plot(time_vec, sig, 'r-',label='raw data')
plt.legend()
#濾波後，畫圖
plt.subplot(2,1,1)
plt.plot(time_vec, filter_sig, 'b-',label='del noise')
plt.legend()
#plt.show();

#頻域畫圖
#在下方畫出sub plot的方法
plt.subplot(2,1,2)
plt.plot(sample_freq, Amplitude, 'r-',label='FFT(frequency spectrum')
plt.legend()
plt.show();