"""
测试数据生成工具
"""
import numpy as np
import os
import json

def generate_test_scenarios(num_scenarios=10, workspace_range=(-0.5, 0.5)):
    """
    生成测试场景数据
    
    参数:
        num_scenarios: 要生成的测试场景数量
        workspace_range: 工作空间范围
    返回:
        scenarios: 包含初始位置和目标位置的测试场景列表
    """
    scenarios = []
    for i in range(num_scenarios):
        # 生成随机初始位置和目标位置
        init_pos = np.random.uniform(workspace_range[0], workspace_range[1], 3)
        target_pos = np.random.uniform(workspace_range[0], workspace_range[1], 3)
        
        # 确保目标位置与初始位置有一定距离
        while np.linalg.norm(target_pos - init_pos) < 0.3:
            target_pos = np.random.uniform(workspace_range[0], workspace_range[1], 3)
            
        scenarios.append({
            'scenario_id': i,
            'init_pos': init_pos.tolist(),
            'target_pos': target_pos.tolist()
        })
        
    return scenarios


def save_scenarios_to_json(scenarios, file_path='test_scenarios.json'):
    """
    保存测试场景到JSON文件

    参数:
        scenarios: 测试场景列表
        file_path: 输出文件路径
    """
    # 确保目录存在
    os.makedirs(os.path.dirname(file_path), exist_ok=True)

    with open(file_path, 'w') as f:
        json.dump(scenarios, f, indent=2)
    print(f"成功保存 {len(scenarios)} 个测试场景到 {file_path}")

def load_scenarios_from_json(file_path='test_scenarios.json'):
    """
    从JSON文件加载测试场景
    
    参数:
        file_path: 输入文件路径
    返回:
        scenarios: 加载的测试场景列表
    """
    try:
        with open(file_path, 'r') as f:
            scenarios = json.load(f)
        print(f"成功加载 {len(scenarios)} 个测试场景")
        return scenarios
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到")
        return None

if __name__ == "__main__":
    # 示例用法
    test_scenarios = generate_test_scenarios(5)
    save_scenarios_to_json(test_scenarios, 'data/test_scenarios.json')