import torch
import numpy as np
import os
from ..models.mlp_network import SimpleIKNet
from ..layers.differentiable_pcc import DifferentiablePCC

def train():
    # 1. 配置
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Running on {device}")
    
    # 初始化模型
    ik_model = SimpleIKNet().to(device)
    # 初始化物理层 (不需要优化参数，只用于计算)
    fk_layer = DifferentiablePCC().to(device)
    
    optimizer = torch.optim.Adam(ik_model.parameters(), lr=1e-4)
    
    # 2. 训练循环
    epochs = 20000
    batch_size = 64
    
    # 定义关节限制 (Constraint)
    # 假设 theta 范围 [0, pi/2], phi 范围 [-pi, pi]
    theta_limit = 3.14159 / 2
    
    for epoch in range(epochs):
        # --- A. 随机采样目标点 (Unsupervised) ---
        # 我们需要在机械臂可达的工作空间内采样
        # 简单起见，我们在一个大致的半球内采样，或者随机生成关节角求正解来得到目标点
        # 方法：Teacher Forcing (用随机关节角生成的真实 Pos 作为 Target) 
        # 这样能保证 Target 一定是可达的，训练效率最高
        
        with torch.no_grad():
            # 随机生成合法的关节角作为"答案"
            random_thetas = torch.rand(batch_size, 3).to(device) * theta_limit
            random_phis = (torch.rand(batch_size, 3).to(device) * 2 - 1) * torch.pi
            
            # 拼成 (Batch, 6)
            gt_joints = torch.zeros(batch_size, 6).to(device)
            gt_joints[:, 0::2] = random_thetas
            gt_joints[:, 1::2] = random_phis
            
            # 计算对应的 Target Position
            target_pos = fk_layer(gt_joints) 
        
        # --- B. 前向传播 ---
        # 输入：Target Position
        # 输出：Predicted Joints
        pred_joints = ik_model(target_pos)
        
        # --- C. 物理一致性 Loss (PINN Loss) ---
        # 将预测的关节角代入物理层，看算出的位置在哪里
        est_pos = fk_layer(pred_joints)
        
        # 1. 位置误差 (Task Loss)
        loss_pos = torch.mean((est_pos - target_pos)**2)
        
        # 2. 关节限制 Loss (Physical Constraint)
        # 惩罚 theta < 0 或 theta > limit
        # 提取 theta: 0, 2, 4 列
        pred_thetas = pred_joints[:, 0::2]
        loss_limits = torch.mean(torch.relu(-pred_thetas)**2) + \
                      torch.mean(torch.relu(pred_thetas - theta_limit)**2)
                      
        # 总 Loss
        total_loss = loss_pos + 0.1 * loss_limits
        
        # --- D. 反向传播 ---
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
        
        if epoch % 1000 == 0:
            print(f"Epoch {epoch} | Pos Loss: {loss_pos.item():.6f} | Limit Loss: {loss_limits.item():.6f}")

    # 3. 保存
    save_path = "../checkpoints/pinn_basic.pth"
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    torch.save(ik_model.state_dict(), save_path)
    print(f"Model saved to {save_path}")

if __name__ == "__main__":
    train()