#开发时间：2024/6/12 14:29
import numpy as np
import pandas as pd
from filterpy.kalman import KalmanFilter
import matplotlib.pyplot as plt
from matplotlib import rcParams

# 卡尔曼滤波去噪
def kalman_denoise(signal):
    kf = KalmanFilter(dim_x=2, dim_z=1)
    kf.x = np.array([0., 0.])
    kf.F = np.array([[1., 1.],
                     [0., 1.]])
    kf.H = np.array([[1., 0.]])
    kf.P *= 1000.
    kf.R = 5
    kf.Q = np.array([[1e-5, 0.],
                     [0., 1e-5]])

    filtered_signal = []
    for z in signal:
        kf.predict()
        kf.update(z)
        filtered_signal.append(kf.x[0])
    return np.array(filtered_signal)

plt.rcParams['xtick.direction'] = 'in'
plt.rcParams['ytick.direction'] = 'in'
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

config = {
    "font.family": 'serif',
    "font.size": 20,
    "mathtext.fontset": 'stix',
    "font.serif": ['Times New Roman'],  # 宋体
    'axes.unicode_minus': False  # 处理负号
}

rcParams.update(config)

# 读取信号
fs = 1000
noisy_signal = pd.read_csv('noisy_signals_time.csv')

# 检查是否包含时间列
if 'time' in noisy_signal.columns:
    time_column = noisy_signal['time'].values.reshape(-1, 1)
    other_columns = noisy_signal.drop(columns='time')
else:
    time_column = None
    other_columns = noisy_signal

# 应用去噪方法
denoised_signals = pd.DataFrame()

for column in other_columns.columns:
    denoised_signals[column] = kalman_denoise(other_columns[column].values)

# 如果有时间列，将其添加回去
if time_column is not None:
    denoised_signals.insert(0, 'time', time_column)

# 保存去噪后的信号到CSV文件
denoised_signals.to_csv('denoise_kalman.csv', index=False)

# 绘图
plt.figure(figsize=(12, 10))

# 绘制原始信号和去噪后信号（这里只绘制第一列作为示例）
plt.subplot(2, 1, 1)
plt.plot(noisy_signal.iloc[:, 1], label='Noisy Signal')
plt.legend()

plt.subplot(2, 1, 2)
plt.plot(denoised_signals.iloc[:, 1], label='Kalman Filtered')
plt.legend()

plt.tight_layout()
plt.show()