preprocess_data.py

"""数据预处理与特征工程脚本。

将后台收集的原始鼠标数据处理为神经网络训练可用的样本。
主要功能包括：
1. 加载原始CSV数据
2. 根据时间间隔分割独立轨迹段
3. 计算速度、加速度等衍生特征
4. 采样轨迹点并构建训练样本
5. 进行Z-score归一化
6. 划分训练集和测试集
"""

import os
import csv
import json
import numpy as np
import yaml



def load_config():
    """加载配置文件。

    Returns:
        dict: 配置参数字典。
    """
    config_path = os.path.join(os.path.dirname(__file__), 'config.yaml')
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


config = load_config()
RAW_DIR = config['data']['raw_dir']
PROCESSED_DIR = config['data']['processed_dir']
MAX_TIME_GAP = config['data']['max_time_gap']
MAX_DIRECTION_CHANGE = config['data'].get('max_direction_change', np.pi / 3)
MIN_POINTS_PER_TRAJECTORY = config['data']['min_points_per_trajectory']
NUM_KEY_POINTS = config['data']['num_key_points']
MIN_KEY_POINTS = config['data'].get('min_key_points', 5)
MAX_KEY_POINTS = config['data'].get('max_key_points', 100)
TRAIN_TEST_SPLIT = config['data']['train_test_split']
TRAIN_DATA_PATH = config['data']['train_data_path']
TEST_DATA_PATH = config['data']['test_data_path']

if not MIN_KEY_POINTS <= NUM_KEY_POINTS <= MAX_KEY_POINTS:
    raise ValueError(f"num_key_points must be between {MIN_KEY_POINTS} and {MAX_KEY_POINTS}")

MIN_POINTS_PER_TRAJECTORY = max(MIN_POINTS_PER_TRAJECTORY, NUM_KEY_POINTS)


def load_raw_data():
    """加载所有原始鼠标数据。

    遍历 RAW_DIR 目录下所有CSV文件，读取时间戳、X、Y坐标，
    并按时间戳排序。

    Returns:
        list: 包含所有点的列表，每个点为(timestamp, x, y)元组。
    """
    all_points = []
    for filename in os.listdir(RAW_DIR):
        if filename.endswith(".csv"):
            filepath = os.path.join(RAW_DIR, filename)
            try:
                with open(filepath, "r") as f:
                    reader = csv.DictReader(f)
                    for row in reader:
                        try:
                            timestamp = float(row["timestamp"])
                            x = int(row["x"])
                            y = int(row["y"])
                            all_points.append((timestamp, x, y))
                        except (ValueError, KeyError):
                            continue
            except Exception as e:
                print(f"Error reading {filepath}: {e}")
    all_points.sort(key=lambda p: p[0])
    return all_points


def split_trajectories(all_points):
    """将连续的点序列分割为独立的轨迹段。

    根据 MAX_TIME_GAP 和 MAX_DIRECTION_CHANGE 阈值进行分割：
    1. 当两个连续点之间的时间间隔超过 MAX_TIME_GAP 时，认为是新轨迹的开始
    2. 当相邻向量间的夹角超过 MAX_DIRECTION_CHANGE 时，强行分割为新轨迹
    只保留点数足够的轨迹。

    Args:
        all_points (list): 所有点的列表，每个点为(timestamp, x, y)元组。

    Returns:
        list: 轨迹列表，每个轨迹是一个点列表。
    """
    trajectories = []
    current_trajectory = []

    for point in all_points:
        if not current_trajectory:
            current_trajectory.append(point)
        else:
            time_gap = point[0] - current_trajectory[-1][0]
            
            need_split = time_gap > MAX_TIME_GAP
            
            if not need_split and len(current_trajectory) >= 2:
                prev_point = current_trajectory[-2]
                curr_point = current_trajectory[-1]
                
                dx1 = curr_point[1] - prev_point[1]
                dy1 = curr_point[2] - prev_point[2]
                dx2 = point[1] - curr_point[1]
                dy2 = point[2] - curr_point[2]
                
                len1 = np.sqrt(dx1 ** 2 + dy1 ** 2)
                len2 = np.sqrt(dx2 ** 2 + dy2 ** 2)
                
                if len1 > 0 and len2 > 0:
                    cos_angle = (dx1 * dx2 + dy1 * dy2) / (len1 * len2)
                    cos_angle = np.clip(cos_angle, -1.0, 1.0)
                    angle = np.arccos(cos_angle)
                    if angle > MAX_DIRECTION_CHANGE:
                        need_split = True
            
            if need_split:
                if len(current_trajectory) >= MIN_POINTS_PER_TRAJECTORY:
                    trajectories.append(current_trajectory)
                current_trajectory = [point]
            else:
                current_trajectory.append(point)

    if len(current_trajectory) >= MIN_POINTS_PER_TRAJECTORY:
        trajectories.append(current_trajectory)

    print(f"Split {len(all_points)} points into {len(trajectories)} trajectories")
    return trajectories


def calculate_features(trajectory):
    """计算轨迹的速度、加速度和曲率特征。

    Args:
        trajectory (list): 单个轨迹的点列表。

    Returns:
        tuple: (timestamps, x_coords, y_coords, speeds, accelerations, curvatures)
    """
    timestamps = np.array([p[0] for p in trajectory])
    x_coords = np.array([p[1] for p in trajectory])
    y_coords = np.array([p[2] for p in trajectory])

    time_diffs = np.diff(timestamps)
    time_diffs[time_diffs == 0] = 1e-6

    dx = np.diff(x_coords)
    dy = np.diff(y_coords)
    distances = np.sqrt(dx ** 2 + dy ** 2)
    speeds = distances / time_diffs

    acc_diffs = np.diff(speeds)
    acc_time_diffs = time_diffs[1:]
    acc_time_diffs[acc_time_diffs == 0] = 1e-6
    accelerations = acc_diffs / acc_time_diffs

    directions = np.arctan2(dy, dx)
    dir_diffs = np.diff(directions)
    dir_diffs = np.arctan2(np.sin(dir_diffs), np.cos(dir_diffs))
    curv_time_diffs = (time_diffs[:-1] + time_diffs[1:]) / 2
    curv_time_diffs[curv_time_diffs == 0] = 1e-6
    curvatures = dir_diffs / curv_time_diffs

    return timestamps, x_coords, y_coords, speeds, accelerations, curvatures


def sample_trajectory(trajectory):
    """从完整轨迹中采样关键点并构建训练样本。

    使用自适应采样策略：根据方向变化（曲率代理）加权分配采样点，
    在方向变化大的区域分配更多采样点。

    Args:
        trajectory (list): 单个轨迹的点列表。

    Returns:
        tuple: (dx, dy, distance, direction, avg_curvature, sampled_x, sampled_y, 
                sampled_speeds, sampled_accelerations, sampled_curvatures)
    """
    timestamps, x_coords, y_coords, speeds, accelerations, curvatures = calculate_features(trajectory)

    start_idx = 0
    end_idx = len(trajectory) - 1
    start_x, start_y = x_coords[start_idx], y_coords[start_idx]
    end_x, end_y = x_coords[end_idx], y_coords[end_idx]

    dx = end_x - start_x
    dy = end_y - start_y
    distance = np.sqrt(dx ** 2 + dy ** 2)
    direction = np.arctan2(dy, dx)

    avg_curvature = np.mean(np.abs(curvatures)) if len(curvatures) > 0 else 0.0

    dx_segments = np.diff(x_coords)
    dy_segments = np.diff(y_coords)
    directions = np.arctan2(dy_segments, dx_segments)
    dir_changes = np.abs(np.diff(directions))
    dir_changes = np.arctan2(np.sin(dir_changes), np.cos(dir_changes))

    weights = np.ones(len(trajectory))
    if len(dir_changes) > 0 and np.max(dir_changes) > 0:
        dir_change_weights = 1 + dir_changes / np.max(dir_changes)
        for i in range(len(dir_change_weights)):
            weights[i + 1] += dir_change_weights[i]

    weights = weights / weights.sum()
    cumsum_weights = np.cumsum(weights)

    sample_indices = np.zeros(NUM_KEY_POINTS, dtype=int)
    sample_indices[0] = start_idx
    sample_indices[-1] = end_idx

    for i in range(1, NUM_KEY_POINTS - 1):
        target = i / (NUM_KEY_POINTS - 1)
        sample_indices[i] = np.argmin(np.abs(cumsum_weights - target))

    sample_indices = np.sort(np.unique(sample_indices))

    if len(sample_indices) < NUM_KEY_POINTS:
        remaining = NUM_KEY_POINTS - len(sample_indices)
        extra_indices = np.linspace(start_idx, end_idx, remaining + 2, dtype=int)[1:-1]
        combined = np.concatenate([sample_indices, extra_indices])
        combined = np.sort(np.unique(combined))
        
        if len(combined) > NUM_KEY_POINTS:
            sample_indices = combined[np.linspace(0, len(combined) - 1, NUM_KEY_POINTS, dtype=int)]
        else:
            padding = NUM_KEY_POINTS - len(combined)
            if padding > 0:
                padding_indices = np.linspace(start_idx, end_idx, padding, dtype=int)
                combined = np.sort(np.unique(np.concatenate([combined, padding_indices])))
                if len(combined) > NUM_KEY_POINTS:
                    sample_indices = combined[np.linspace(0, len(combined) - 1, NUM_KEY_POINTS, dtype=int)]
                else:
                    sample_indices = np.linspace(start_idx, end_idx, NUM_KEY_POINTS, dtype=int)
            else:
                sample_indices = combined
    elif len(sample_indices) > NUM_KEY_POINTS:
        sample_indices = sample_indices[np.linspace(0, len(sample_indices) - 1, NUM_KEY_POINTS, dtype=int)]

    sampled_x = x_coords[sample_indices] - start_x
    sampled_y = y_coords[sample_indices] - start_y

    sampled_speeds = np.zeros(NUM_KEY_POINTS)
    for i, idx in enumerate(sample_indices[:-1]):
        next_idx = sample_indices[i + 1]
        if idx < len(speeds) and next_idx <= len(speeds):
            sampled_speeds[i] = np.mean(speeds[idx:next_idx])
        elif idx < len(speeds):
            sampled_speeds[i] = speeds[idx]
    if len(speeds) > 0:
        sampled_speeds[-1] = speeds[-1]

    sampled_accelerations = np.zeros(NUM_KEY_POINTS - 1)
    for i, idx in enumerate(sample_indices[:-2]):
        next_idx = sample_indices[i + 1]
        if idx < len(accelerations) and next_idx <= len(accelerations):
            sampled_accelerations[i] = np.mean(accelerations[idx:next_idx])
        elif idx < len(accelerations):
            sampled_accelerations[i] = accelerations[idx]
    if len(accelerations) > 0:
        sampled_accelerations[-1] = accelerations[-1]

    sampled_curvatures = np.zeros(NUM_KEY_POINTS - 1)
    for i, idx in enumerate(sample_indices[:-2]):
        next_idx = sample_indices[i + 1]
        if idx < len(curvatures) and next_idx <= len(curvatures):
            sampled_curvatures[i] = np.mean(curvatures[idx:next_idx])
        elif idx < len(curvatures):
            sampled_curvatures[i] = curvatures[idx]
    if len(curvatures) > 0:
        sampled_curvatures[-1] = curvatures[-1]

    return dx, dy, distance, direction, avg_curvature, sampled_x, sampled_y, sampled_speeds, sampled_accelerations, sampled_curvatures


def process_all_trajectories(trajectories):
    """处理所有轨迹，生成训练样本列表。

    Args:
        trajectories (list): 轨迹列表。

    Returns:
        list: 样本列表。
    """
    samples = []
    for traj in trajectories:
        try:
            dx, dy, distance, direction, avg_curvature, x, y, speeds, accelerations, curvatures = sample_trajectory(traj)
            if abs(dx) > 0 or abs(dy) > 0:
                samples.append({
                    "dx": dx,
                    "dy": dy,
                    "distance": distance,
                    "direction": direction,
                    "avg_curvature": avg_curvature,
                    "x": x,
                    "y": y,
                    "speeds": speeds,
                    "accelerations": accelerations,
                    "curvatures": curvatures
                })
        except Exception as e:
            print(f"Error processing trajectory: {e}")
    print(f"Generated {len(samples)} valid training samples")
    return samples


def augment_samples(samples):
    """对训练样本进行数据增强。

    增强方法：
    1. 翻转增强：水平翻转（x坐标取反）
    2. 缩放增强：按比例缩放轨迹（0.8倍）
    3. 垂直翻转增强：y坐标取反

    Args:
        samples (list): 原始样本列表

    Returns:
        list: 增强后的样本列表（包含原始样本）
    """
    augmented = list(samples)

    for s in samples:
        flipped_x = {
            "dx": -s["dx"],
            "dy": s["dy"],
            "distance": s["distance"],
            "direction": np.pi - s["direction"],
            "avg_curvature": s["avg_curvature"],
            "x": -s["x"].copy(),
            "y": s["y"].copy(),
            "speeds": s["speeds"].copy(),
            "accelerations": -s["accelerations"].copy(),
            "curvatures": -s["curvatures"].copy()
        }
        augmented.append(flipped_x)

        flipped_y = {
            "dx": s["dx"],
            "dy": -s["dy"],
            "distance": s["distance"],
            "direction": -s["direction"],
            "avg_curvature": s["avg_curvature"],
            "x": s["x"].copy(),
            "y": -s["y"].copy(),
            "speeds": s["speeds"].copy(),
            "accelerations": -s["accelerations"].copy(),
            "curvatures": -s["curvatures"].copy()
        }
        augmented.append(flipped_y)

        scaled = {
            "dx": s["dx"] * 0.8,
            "dy": s["dy"] * 0.8,
            "distance": s["distance"] * 0.8,
            "direction": s["direction"],
            "avg_curvature": s["avg_curvature"],
            "x": s["x"].copy() * 0.8,
            "y": s["y"].copy() * 0.8,
            "speeds": s["speeds"].copy() * 0.8,
            "accelerations": s["accelerations"].copy() * 0.8,
            "curvatures": s["curvatures"].copy()
        }
        augmented.append(scaled)

    print(f"Augmented {len(samples)} samples to {len(augmented)} samples")
    return augmented


def compute_normalization_params(samples):
    """计算数据归一化参数（Z-score标准化）。

    Args:
        samples (list): 样本列表。

    Returns:
        dict: 归一化参数字典。
    """
    all_dx = []
    all_dy = []
    all_distance = []
    all_direction = []
    all_avg_curvature = []
    all_x = []
    all_y = []
    all_speeds = []
    all_acc = []
    all_curvatures = []

    for s in samples:
        all_dx.append(s["dx"])
        all_dy.append(s["dy"])
        all_distance.append(s["distance"])
        all_direction.append(s["direction"])
        all_avg_curvature.append(s["avg_curvature"])
        for i in range(NUM_KEY_POINTS):
            all_x.append(s["x"][i])
            all_y.append(s["y"][i])
            all_speeds.append(s["speeds"][i])
        for i in range(NUM_KEY_POINTS - 1):
            all_acc.append(s["accelerations"][i])
            all_curvatures.append(s["curvatures"][i])

    dx_array = np.array(all_dx)
    dy_array = np.array(all_dy)
    distance_array = np.array(all_distance)
    direction_array = np.array(all_direction)
    avg_curvature_array = np.array(all_avg_curvature)
    x_array = np.array(all_x)
    y_array = np.array(all_y)
    speeds_array = np.array(all_speeds)
    acc_array = np.array(all_acc)
    curvatures_array = np.array(all_curvatures)

    params = {
        "dx_mean": float(np.mean(dx_array)),
        "dx_std": float(np.std(dx_array)),
        "dy_mean": float(np.mean(dy_array)),
        "dy_std": float(np.std(dy_array)),
        "distance_mean": float(np.mean(distance_array)),
        "distance_std": float(np.std(distance_array)),
        "direction_mean": float(np.mean(direction_array)),
        "direction_std": float(np.std(direction_array)),
        "avg_curvature_mean": float(np.mean(avg_curvature_array)),
        "avg_curvature_std": float(np.std(avg_curvature_array)),
        "x_mean": float(np.mean(x_array)),
        "x_std": float(np.std(x_array)),
        "y_mean": float(np.mean(y_array)),
        "y_std": float(np.std(y_array)),
        "speeds_mean": float(np.mean(speeds_array)),
        "speeds_std": float(np.std(speeds_array)),
        "acc_mean": float(np.mean(acc_array)),
        "acc_std": float(np.std(acc_array)),
        "curvatures_mean": float(np.mean(curvatures_array)),
        "curvatures_std": float(np.std(curvatures_array))
    }

    return params


def normalize_samples(samples, params):
    """对样本进行Z-score归一化。

    Args:
        samples (list): 待归一化的样本列表。
        params (dict): 归一化参数字典。

    Returns:
        list: 归一化后的样本列表。
    """
    for s in samples:
        s["dx"] = (s["dx"] - params["dx_mean"]) / params["dx_std"]
        s["dy"] = (s["dy"] - params["dy_mean"]) / params["dy_std"]
        s["distance"] = (s["distance"] - params["distance_mean"]) / params["distance_std"]
        s["direction"] = (s["direction"] - params["direction_mean"]) / params["direction_std"]
        s["avg_curvature"] = (s["avg_curvature"] - params["avg_curvature_mean"]) / params["avg_curvature_std"]
        s["x"] = (s["x"] - params["x_mean"]) / params["x_std"]
        s["y"] = (s["y"] - params["y_mean"]) / params["y_std"]
        s["speeds"] = (s["speeds"] - params["speeds_mean"]) / params["speeds_std"]
        s["accelerations"] = (s["accelerations"] - params["acc_mean"]) / params["acc_std"]
        s["curvatures"] = (s["curvatures"] - params["curvatures_mean"]) / params["curvatures_std"]
    return samples


def write_csv(filepath, samples):
    """将样本写入CSV文件。

    Args:
        filepath (str): 输出CSV文件路径。
        samples (list): 样本列表。
    """
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w", newline="") as f:
        writer = csv.writer(f)
        for sample in samples:
            input_str = f"{sample['dx']},{sample['dy']},{sample['distance']},{sample['direction']},{sample['avg_curvature']}"
            points_str = [f"{sample['x'][i]:.6f},{sample['y'][i]:.6f}" for i in range(NUM_KEY_POINTS)]
            speeds_str = [f"{s:.6f}" for s in sample["speeds"]]
            acc_str = [f"{a:.6f}" for a in sample["accelerations"]]
            curv_str = [f"{c:.6f}" for c in sample["curvatures"]]
            writer.writerow([input_str] + points_str + speeds_str + acc_str + curv_str)


def main():
    """主函数：执行完整的数据预处理流程。"""
    all_points = load_raw_data()
    if not all_points:
        print("No raw data found. Please run background_collector.py first.")
        return

    trajectories = split_trajectories(all_points)
    if not trajectories:
        print("No valid trajectories found.")
        return

    num_trajectories = len(trajectories)
    train_split_idx = int(num_trajectories * TRAIN_TEST_SPLIT)
    
    train_trajectories = trajectories[:train_split_idx]
    test_trajectories = trajectories[train_split_idx:]
    
    print(f"Split {num_trajectories} trajectories into {len(train_trajectories)} train and {len(test_trajectories)} test")

    train_samples = process_all_trajectories(train_trajectories)
    test_samples = process_all_trajectories(test_trajectories)
    
    if not train_samples:
        print("No valid training samples generated.")
        return
    if not test_samples:
        print("No valid test samples generated.")
        return

    train_samples = augment_samples(train_samples)

    normalization_params = compute_normalization_params(train_samples)
    
    normalized_train_samples = normalize_samples(train_samples, normalization_params)
    normalized_test_samples = normalize_samples(test_samples, normalization_params)

    params_path = os.path.join(PROCESSED_DIR, "normalization_params.json")
    with open(params_path, "w") as f:
        json.dump(normalization_params, f, indent=2)
    print(f"Normalization params saved to: {params_path}")

    write_csv(TRAIN_DATA_PATH, normalized_train_samples)
    write_csv(TEST_DATA_PATH, normalized_test_samples)

    print(f"Training data saved to: {TRAIN_DATA_PATH} ({len(normalized_train_samples)} samples)")
    print(f"Test data saved to: {TEST_DATA_PATH} ({len(normalized_test_samples)} samples)")
    print(f"Total: {len(normalized_train_samples) + len(normalized_test_samples)} samples")


if __name__ == "__main__":
    main()