import numpy as np
import matplotlib.pyplot as plt
import csv
import os


def read_csv_column(csv_file_path, header_name, step=10):
    with open(csv_file_path, 'r') as file:
        reader = csv.DictReader(file)
        column = []
        for i, row in enumerate(reader):
            if i % step == 0:
                column.append(round(float(row[header_name]), 4))
    return column


def plotCurves(csv_file_path, save_path, col_names, step=10):
    # plt.clf()
    plt.figure(figsize=(12, 8))
    fig, ax = plt.subplots()
    for col_name in col_names:
        y_data = read_csv_column(csv_file_path, col_name, step=step)
        x_data = list(range(0, len(y_data)))
        ax.plot(x_data, y_data, label=col_name)
    plt.xlabel('Time')
    plt.ylabel('Value')
    plt.title('Raw Data')
    ax.legend()
    plt.savefig(save_path)


def process_all_csv_files(data_directory, output_directory, col_names, step=10):
    for filename in os.listdir(data_directory):
        if filename.endswith('.csv'):
            csv_file_path = os.path.join(data_directory, filename)
            save_path = os.path.join(output_directory, filename.split('.')[0] + '.png')
            plotCurves(csv_file_path, save_path, col_names, step)


if __name__ == '__main__':
    data_directory = r'D:\pycharmProjects\python_server\data'
    output_directory = r'D:\pycharmProjects\python_server\raw_png'
    col_names = ('accX', 'accY', 'accZ', 'gyroX', 'gyroY', 'gyroZ', 'magX', 'magY', 'magZ')
    process_all_csv_files(data_directory, output_directory, col_names, step=1)