import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import glob


class KalmanFilter:
    def __init__(self, process_variance, measurement_variance, initial_estimate, initial_error_covariance):
        self.process_variance = process_variance
        self.measurement_variance = measurement_variance
        self.estimate = initial_estimate
        self.error_covariance = initial_error_covariance

    def update(self, measurement):
        predicted_estimate = self.estimate
        predicted_error_covariance = self.error_covariance + self.process_variance
        kalman_gain = predicted_error_covariance / (predicted_error_covariance + self.measurement_variance)
        self.estimate = predicted_estimate + kalman_gain * (measurement - predicted_estimate)
        self.error_covariance = (1 - kalman_gain) * predicted_error_covariance
        return self.estimate


def read_csv_columns(csv_file_path, header_names):
    df = pd.read_csv(csv_file_path)
    return df[header_names].values


def save_to_csv(data, header_names, output_file):
    df = pd.DataFrame(data, columns=header_names)
    df.to_csv(output_file, index=False)


def plot_comparison_and_save(original_data, denoised_data, header_names, save_path):
    plt.figure(figsize=(14, 10))
    for i, col_name in enumerate(header_names):
        plt.plot(original_data[:, i], label='Original ' + col_name)
        plt.plot(denoised_data[:, i], label='Denoised ' + col_name, linestyle='--', linewidth=2)
    plt.title('Kalman Filter Denoising Comparison')
    plt.xlabel('Time')
    plt.ylabel('Value')
    plt.legend(loc='lower left')
    plt.tight_layout()
    plt.savefig(save_path)


if __name__ == "__main__":
    input_directory = r'D:\pycharmProjects\python_server\Wavelet_DenoisedData'  # 输入目录路径
    output_csv_directory = r'D:\pycharmProjects\python_server\KalManData'  # 输出目录路径
    output_png_directory = r'D:\pycharmProjects\python_server\KalManPng'  # 输出png目录路径
    columns_to_process = ['accX', 'accY', 'accZ', 'gyroX', 'gyroY', 'gyroZ', 'magX', 'magY', 'magZ']

    # 创建输出目录（如果不存在）
    if not os.path.exists(output_csv_directory):
        os.makedirs(output_csv_directory)

    if not os.path.exists(output_png_directory):
        os.makedirs(output_png_directory)

    csv_files = glob.glob(os.path.join(input_directory, '*.csv'))

    for input_csv_file in csv_files:
        file_name = os.path.basename(input_csv_file)
        output_csv_file = os.path.join(output_csv_directory,
                                       os.path.splitext(file_name)[0] + '_kalman_denoised.csv')
        output_plot_file = os.path.join(output_png_directory,
                                        os.path.splitext(file_name)[0] + '_kalman_denoising_comparison.png')

        data = read_csv_columns(input_csv_file, columns_to_process)

        denoised_data = np.zeros_like(data)
        for i, col in enumerate(columns_to_process):
            kf = KalmanFilter(process_variance=1e-5, measurement_variance=0.1 ** 2, initial_estimate=0.0,
                              initial_error_covariance=1.0)
            denoised_data[:, i] = [kf.update(measurement) for measurement in data[:, i]]

        save_to_csv(denoised_data, columns_to_process, output_csv_file)
        plot_comparison_and_save(data, denoised_data, columns_to_process, output_plot_file)

