import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np
import os
import sys
from torch.utils.data import DataLoader, TensorDataset

# --- 路径修复 ---
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.append(current_dir)

from pinn_model import PINN_IK 

def train_model():
    # 1. 路径处理
    csv_path = os.path.join(current_dir, "pcc_dataset_3sec.csv")
    model_save_path = os.path.join(current_dir, "best_model_v2.pth")

    if not os.path.exists(csv_path):
        print("❌ 错误：找不到数据集")
        return

    # 2. 加载数据
    print("正在加载数据...")
    df = pd.read_csv(csv_path)
    
    # 分离输入输出
    X_raw = df.iloc[:, 0:3].values.astype(np.float32) # x, y, z
    y_raw = df.iloc[:, 3:].values.astype(np.float32) # angles

    # --- 关键步骤：数据标准化 (Standardization) ---
    # 计算均值和标准差
    X_mean = X_raw.mean(axis=0)
    X_std = X_raw.std(axis=0) + 1e-6 # 加一点点防止除以0
    
    y_mean = y_raw.mean(axis=0)
    y_std = y_raw.std(axis=0) + 1e-6

    # 归一化公式: (x - mean) / std
    X_norm = (X_raw - X_mean) / X_std
    y_norm = (y_raw - y_mean) / y_std

    print("数据归一化完成。")
    print(f"输入 Mean: {X_mean}, Std: {X_std}")

    # 3. 准备 Tensor
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Training on: {device}")
    
    dataset = TensorDataset(torch.from_numpy(X_norm), torch.from_numpy(y_norm))
    # 增加 Batch Size 提高稳定性
    loader = DataLoader(dataset, batch_size=256, shuffle=True)

    # 4. 初始化模型
    model = PINN_IK().to(device)
    criterion = nn.MSELoss()
    # 使用 AdamW 优化器，效果通常更好
    optimizer = optim.AdamW(model.parameters(), lr=0.001)
    # 学习率衰减策略：每100个epoch学习率乘0.5
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)

    # 5. 训练循环 (增加 Epoch 到 300)
    epochs = 300 
    print(f"开始训练 {epochs} 个 Epoch...")
    
    for epoch in range(epochs):
        total_loss = 0
        model.train()
        
        for batch_x, batch_y in loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            
            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        
        scheduler.step()
        
        if (epoch+1) % 10 == 0:
            avg_loss = total_loss / len(loader)
            current_lr = optimizer.param_groups[0]['lr']
            print(f"Epoch [{epoch+1}/{epochs}] | Loss: {avg_loss:.6f} | LR: {current_lr:.6f}")

    # 6. 保存模型 + 归一化参数
    # 我们保存一个字典，包含模型权重和数据的统计信息
    save_dict = {
        'model_state': model.state_dict(),
        'stats': {
            'X_mean': X_mean, 'X_std': X_std,
            'y_mean': y_mean, 'y_std': y_std
        }
    }
    torch.save(save_dict, model_save_path)
    print(f"✅ 模型及元数据已保存至: {model_save_path}")

if __name__ == "__main__":
    train_model()