import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from sklearn.decomposition import PCA
import os

file_path = './data1/npy/ALL/'  # 数据当前目录
store_path = './data1/ProcessedData2/'  # 处理好的数据的存储目录
components_number = 20  # PCA降维后的维度
pca = PCA(n_components=components_number)  # 主成分分析方法，返回降维后的数据，但是这个代码没有用到
fre = 60  # 采样频率，每秒的采样数


def filt(data):
    """
    低通滤波器
    :param data: 输入数据, shape=(通道/载波数, 时间帧), 例如(30, 120)，即30个载波，2s的数据，每秒60帧
    :return: 滤波后的数据, shape=(通道/载波数, 时间帧)
    """
    # signal.butter 返回滤波器的分子和分母系数，lowpass表示低通滤波器
    # 配置滤波器 8 表示滤波器的阶数, 0.5为截止频率，即滤除0.5倍采样频率以上的信号
    b, a = signal.butter(8, 0.5, 'lowpass')  
    # filtfilt函数是直接对数据进行滤波，不改变数据的相位，即不改变数据的时间信息
    # data为要过滤的信号, axis=1表示对每一列进行滤波
    filtedData = signal.filtfilt(b, a, data, axis=1) 
    return filtedData

def example_data_and_plot(file_path):
    for file_path, sub_dirs, filenames in os.walk(file_path):
        if filenames:
            filenames = sorted(np.array(filenames))
            for filename in filenames:
                print(filename)
                data = np.load(os.path.join(file_path, filename), allow_pickle=True)

                ori_data_5s = data[:, :5 * fre]
                ori_data_2s = ori_data_5s[:, fre:3 * fre]
                ori_data_2s = np.transpose(ori_data_2s)
                plt.subplot(2, 1, 1)
                plt.plot(ori_data_2s.reshape(-1))
                plt.title('before low pass filter')

                filtedData = filt(data)
                filted_data_5s = filtedData[:, :5 * fre]
                filted_data_2s = filted_data_5s[:, fre:3 * fre]
                filted_data_2s = np.transpose(filted_data_2s)
                plt.subplot(2, 1, 2)
                plt.plot(filted_data_2s.reshape(-1))
                plt.title('after low pass filter')
                plt.suptitle("data process example")
                plt.show()
                exit()


# example_data_and_plot(file_path) # 如果想看数据处理前后的效果，可以取消注释
for file_path, sub_dirs, filenames in os.walk(file_path):
    if filenames:
        filenames = sorted(np.array(filenames))
        for filename in filenames:
            print(filename)
            data = np.load(os.path.join(file_path, filename), allow_pickle=True)  # shape=(通道/载波数, 时间帧)
            filtedData = filt(data)  # 低通滤波, shape=(通道/载波数, 时间帧)

            for i in range(9000 // (5 * fre)):  # 取样长度5s，再切片2s有效数据

                # 选取所有行，5 * fre是5秒的采样数, shape=(通道/载波数, 5 * fre)
                data_5s = filtedData[:, 5 * fre * i:5 * fre * (i + 1)]
                
                # 选取第fre列到第3 * fre列的数据，即2s的数据，注意切片是左闭右开区间, shape=(通道/载波数, 2 * fre)
                data_2s = data_5s[:, fre:3 * fre]
                data_pca = data_2s
                data_pca = np.transpose(data_pca)  # 转置，shape=(2 * fre, 通道/载波数)，即(120, 30)
                newFileName = filename[:-4] + '_' + str(i + 1)
                np.save(os.path.join(store_path, newFileName), data_pca)