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


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


def normalize(data):
    mean = np.mean(data, axis=0)
    std = np.std(data, axis=0)
    return (data - mean) / std


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, normalized_data, header_names, save_path):
    plt.figure(figsize=(14, 10))
    for i, col_name in enumerate(header_names):
        plt.plot(normalized_data[:, i], label='Normalized ' + col_name)
    plt.title('Normalization')
    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\KalManData'  # 输入目录路径
    output_csv_directory = r'D:\pycharmProjects\python_server\NormalizeData'  # 输出目录路径
    output_png_directory = r'D:\pycharmProjects\python_server\NormalizePng'  # 输出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] + '_normalized.csv')
        output_plot_file = os.path.join(output_png_directory,
                                        os.path.splitext(file_name)[0] + '_normalization_comparison.png')

        data = read_csv_columns(input_csv_file, columns_to_process)
        normalized_data = normalize(data)

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