import numpy as np
import pylab as pl
import scipy as sc
import time
from matplotlib.pylab import mpl

mpl.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei']  # 显示中文
mpl.rcParams['axes.unicode_minus'] = False  # 显示负号


# 从文件中读取数据，后续直接使用usb省略这一步
def Read_data(data_name):
    # 通过文件读取数据
    # mat_data = sc.io.loadmat(data_name)  # 读取文件
    #
    # mat_data.keys()  # 获得文件分类
    #
    # data = mat_data["X122_DE_time"]  # 提取振动数据

    data = np.loadtxt(data_name)
    # print(data)
    # print(len(data))
    return data


# 对数据进行快速傅里叶变换
def FFT(rate, number, data):
    sampling_rate = rate  # 采样频率
    fft_size = number  # FFT处理的数据样本数
    time_start = time.time()
    # 进行快速傅里叶变换
    xs = data[:fft_size]  # 从波形数据中取样fft_size个点进行运算
    # 利用np.fft.fft()进行FFT计算
    # rfft返回N/2+1个数据。而fft返回N个数据，但这N个数据是左右对称的，因此实际有效的是N/2+1个数据。
    time_end1 = time.time()
    xfft = sc.fft.fftn(xs)
    # a = 0
    # while a < fft_size:
    #     print(a, ":", xfft[a])
    #     a = a + 1
    # 打印显示前5个傅里叶转换结果
    # 可以看出傅里叶转换结果是复数
    # 复数用y=a+bj的形式表示，a为实部，b为虚部
    # print(xfft[:5])

    # 通过下面的np.linspace计算出返回值中每个下标对应的真正的频率：
    # 最后除以1e3将频率单位由Hz转换为KHz
    time_end2 = time.time()
    freqs = np.linspace(0, int(sampling_rate), int(fft_size)) / 1000

    # 对复数取绝对值(取模运算)，复数的模为：√a^2 + b^2
    abs_xfft = abs(xfft)
    time_end3 = time.time()
    # a = 0
    # while a < fft_size:
    #     print(a, ":", abs_xfft[a])
    #     a = a + 1
    # print("Time1", time_end1 - time_start)
    # print("Time2", time_end2 - time_end1)
    # print("Time3", time_end3 - time_end2)
    print("Time4", time_end3 - time_start)
    print()
    return freqs, abs_xfft


# 画出时域信号和频域信号图谱
def draw_figur(data, fft_size, freqs, abs_xfft, sampling_rate):
    fig = pl.figure()

    # 使用鼠标缩放图片
    def call_back(event):
        axtemp = event.inaxes
        x_min, x_max = axtemp.get_xlim()
        fanwei = (x_max - x_min) / 10
        if event.button == 'up':
            axtemp.set(xlim=(x_min + fanwei, x_max - fanwei))
            print('up')
        elif event.button == 'down':
            axtemp.set(xlim=(x_min - fanwei, x_max + fanwei))
            print('down')
        fig.canvas.draw_idle()  # 绘图动作实时反映在图像上

    fig.canvas.mpl_connect('scroll_event', call_back)
    fig.canvas.mpl_connect('button_press_event', call_back)

    pl.subplot(3, 1, 1)  # 截取幕布的一部分
    np.linspace(-5, 5, 10)

    pl.plot(data[:fft_size], color='r')
    pl.xlabel('点数')
    pl.ylabel('加速度值')
    pl.title('时域信号')

    pl.subplot(3, 1, 3)  # 截取幕布的一部分
    np.linspace(-5, 5, 10)

    pl.title(f"频域信号(通道号：{sampling_rate})")
    pl.ylabel("振幅")
    pl.xlabel("频率(KHz)")
    pl.plot(freqs, abs_xfft)

    time_name = time.time()
    pl.show()
    # pl.savefig(f"{time_name}.png")