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=150000, save_path=None): # 增加到 15万
    print(f"🚀 开始生成数据集 (平滑约束版)，目标: {num_samples} ...")
    start_time = time.time()

    kinematics = SoftArmKinematics(3, 0.24, 3, 0.033)

    # 1. 生成协同的姿态
    # 我们不再让三段完全独立，而是让它们倾向于向同一个方向弯曲
    # 这模拟了“能量最小”的自然状态
    
    y_data = np.zeros((num_samples, 9)) # [t1, c1, s1, t2, c2, s2, t3, c3, s3]
    
    # --- 策略：全局趋势 + 局部微调 ---
    # 生成一个“主弯曲角”和“主方向”
    main_theta = np.random.uniform(0.0, np.pi/2.2, num_samples)
    main_phi = np.random.uniform(-np.pi, np.pi, num_samples)
    
    for i in range(3): # 3个段
        # 每一段的弯曲程度在主弯曲角基础上浮动
        # 这样的数据生成的机器人是平滑的 C 形或轻微 S 形，而不是剧烈的乱麻形
        # 这种“平滑先验”让逆运动学变得有唯一解
        noise_theta = np.random.normal(0, 0.1, num_samples) # 小噪声
        noise_phi = np.random.normal(0, 0.2, num_samples)   # 小噪声
        
        thetas = np.clip(main_theta + noise_theta, 0, np.pi/2)
        phis = main_phi + noise_phi # 允许 phi 稍微偏转
        
        # 填充
        y_data[:, i*3 + 0] = thetas
        y_data[:, i*3 + 1] = np.cos(phis)
        y_data[:, i*3 + 2] = np.sin(phis)
        
        # 暂存用于 FK
        if i == 0: t1, p1 = thetas, phis
        if i == 1: t2, p2 = thetas, phis
        if i == 2: t3, p3 = thetas, phis

    # 2. 计算 FK
    pos_data = np.zeros((num_samples, 3))
    print("🔄 正在计算正运动学...")
    
    for i in range(num_samples):
        q_input = [
            (t1[i], p1[i], 0.24),
            (t2[i], p2[i], 0.24),
            (t3[i], p3[i], 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}")

    # 3. 保存
    if save_path is None:
        save_path = os.path.join(current_dir, "pcc_dataset_smooth.csv") # 改名区分

    combined_data = np.hstack((pos_data, y_data))
    header = "x,y,z,t1,c1,s1,t2,c2,s2,t3,c3,s3"
    np.savetxt(save_path, combined_data, delimiter=',', header=header, comments='', fmt='%.6f')
    
    print(f"✅ 完成！耗时: {time.time()-start_time:.2f}s")

if __name__ == "__main__":
    generate_dataset()