import csv
import time
import os
import numpy as np


class DataLogger:
    def __init__(self, filename_prefix="dataset"):
        """
        初始化日志记录器
        :param filename_prefix: 文件名前缀，例如 "train_batch"
        """
        # 1. 自动创建 data 文件夹，保持项目目录整洁
        if not os.path.exists('data'):
            os.makedirs('data')
            print("[Logger] Created 'data' directory.")

        # 2. 生成带时间戳的文件名，防止覆盖旧数据
        # 格式: data/dataset_20240105_143022.csv
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        self.filename = f"data/{filename_prefix}_{timestamp}.csv"

        # 3. 定义标准表头 (Header) - 对应你的12个力 + 3个构型 + 3个标签
        # 传感器命名规则: F(手指ID)_U(单元ID)_(Fn/Ft)
        sensor_headers = []
        for f in range(3):  # 3个手指
            for u in range(1, 3):  # 2个单元
                sensor_headers.append(f"F{f}_U{u}_Fn")
                sensor_headers.append(f"F{f}_U{u}_Ft")

        # 完整的列名列表
        self.header = ['Timestamp', 'Config_ID'] + sensor_headers + ['Label_X', 'Label_Y', 'Label_Theta']

        # 4. 创建文件并写入第一行（表头）
        try:
            with open(self.filename, mode='w', newline='') as f:
                writer = csv.writer(f)
                writer.writerow(self.header)
            print(f"[Logger] Log file initialized: {self.filename}")
        except Exception as e:
            print(f"[Logger] Error creating file: {e}")

    def log(self, sensor_data, config_id, label_vector):
        """
        写入一行数据
        :param sensor_data: 12维 numpy 数组或列表 (触觉数据)
        :param config_id: 整数 (0, 1, 2) (当前手指构型)
        :param label_vector: [x, y, theta] (机械臂的真实偏差值)
        """
        with open(self.filename, mode='a', newline='') as f:
            writer = csv.writer(f)

            # 拼接数据: 时间 + ID + 传感器数值 + 标签
            # time.time() 获取当前精确时间戳
            row = [f"{time.time():.3f}", int(config_id)] + list(sensor_data) + list(label_vector)

            # 格式化: 将浮点数保留4位小数，节省空间且美观
            formatted_row = []
            for item in row:
                if isinstance(item, float):
                    formatted_row.append(f"{item:.4f}")
                else:
                    formatted_row.append(item)

            writer.writerow(formatted_row)


# === 单元测试 (Unit Test) ===
# 直接运行此文件，测试是否能生成 CSV
if __name__ == "__main__":
    print("Testing DataLogger...")
    logger = DataLogger(filename_prefix="test_data")

    # 模拟写入 5 条假数据
    for i in range(5):
        # 模拟 12 个传感器数据 (0-10之间)
        fake_sensor = np.random.rand(12) * 10
        # 模拟机械臂偏差 (Label)
        fake_label = [0.0, 5.5, -2.0]

        logger.log(sensor_data=fake_sensor, config_id=0, label_vector=fake_label)
        print(f"Logged row {i + 1}")
        time.sleep(0.1)

    print("Test done. Please check the 'data' folder.")