import sys
import os
import numpy as np
import time

# --- 自动路径修正 ---
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, "../../"))
if project_root not in sys.path:
    sys.path.append(project_root)

from soft_arm_sim.model.pcc_kinematics import SoftArmKinematics

def generate_dataset(num_samples=50000, save_path=None):
    """
    生成 PCC 机械臂正运动学数据并保存为 CSV
    """
    print(f"🚀 开始生成数据集，目标数量: {num_samples} ...")
    start_time = time.time()

    # 1. 初始化运动学模型
    kinematics = SoftArmKinematics(
        num_sections=3,
        section_length=0.24, 
        disks_per_section=3,
        disk_radius=0.033
    )

    # 2. 采样范围
    theta_min, theta_max = 0.0, np.pi / 2.5
    phi_min, phi_max = -np.pi, np.pi

    # 3. 随机生成 Joint Space (Theta, Phi)
    # Shape: (N, 6) -> [theta1, phi1, theta2, phi2, theta3, phi3]
    q_data = np.zeros((num_samples, 6))
    for i in range(3): 
        q_data[:, i*2] = np.random.uniform(theta_min, theta_max, num_samples)
        q_data[:, i*2+1] = np.random.uniform(phi_min, phi_max, num_samples)

    # 4. 计算 Task Space (x, y, z)
    # Shape: (N, 3)
    pos_data = np.zeros((num_samples, 3))

    print("🔄 正在计算正运动学 (FK)...")
    
    for i in range(num_samples):
        # 构造输入
        raw_q = q_data[i]
        q_input = [
            (raw_q[0], raw_q[1], 0.24),
            (raw_q[2], raw_q[3], 0.24),
            (raw_q[4], raw_q[5], 0.24)
        ]
        
        # 正运动学解算
        transforms, _ = kinematics.forward(q_input)
        
        # 取末端位置
        pos_data[i] = transforms[-1][:3, 3]
        
        if (i + 1) % 10000 == 0:
            print(f"   进度: {i + 1}/{num_samples}")

    # 5. 保存为 CSV
    if save_path is None:
        save_path = os.path.join(current_dir, "pcc_dataset_3sec.csv")

    print(f"💾 正在保存 CSV 文件...")

    # 将 位置(3列) 和 角度(6列) 拼接到一起 -> 总共9列
    # 列顺序: x, y, z, theta1, phi1, theta2, phi2, theta3, phi3
    combined_data = np.hstack((pos_data, q_data))
    
    # 定义表头
    header_str = "x,y,z,theta1,phi1,theta2,phi2,theta3,phi3"
    
    # 保存
    np.savetxt(save_path, combined_data, delimiter=',', header=header_str, comments='', fmt='%.6f')
    
    duration = time.time() - start_time
    print(f"✅ 完成！")
    print(f"   耗时: {duration:.2f} 秒")
    print(f"   保存路径: {save_path}")

def main():
    generate_dataset(num_samples=50000)

if __name__ == "__main__":
    main()