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():
    csv_path = os.path.join(current_dir, "pcc_dataset_smooth.csv") # 新文件名
    model_save_path = os.path.join(current_dir, "best_model_smooth.pth")

    if not os.path.exists(csv_path):
        print("❌ 请先运行 dataset_generator.py (新版)！")
        return

    print("加载数据中...")
    df = pd.read_csv(csv_path)
    
    # 前3列是输入 (x,y,z)
    X_raw = df.iloc[:, 0:3].values.astype(np.float32)
    # 后9列是标签 (t1, c1, s1, ...)
    y_raw = df.iloc[:, 3:].values.astype(np.float32)

    # --- 归一化 ---
    # 标签里的 cos/sin 本身就在 -1到1 之间，其实不需要强行归一化
    # 但 Theta 需要。为了简单，我们还是整体归一化，推理时反归一化即可。
    X_mean = X_raw.mean(axis=0); X_std = X_raw.std(axis=0) + 1e-6
    y_mean = y_raw.mean(axis=0); y_std = y_raw.std(axis=0) + 1e-6

    X_norm = (X_raw - X_mean) / X_std
    y_norm = (y_raw - y_mean) / y_std

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")
    
    dataset = TensorDataset(torch.from_numpy(X_norm), torch.from_numpy(y_norm))
    loader = DataLoader(dataset, batch_size=512, shuffle=True) # 大 Batch Size

    model = PINN_IK(output_dim=9).to(device)
    criterion = nn.MSELoss()
    optimizer = optim.AdamW(model.parameters(), lr=0.0005) # 稍微降低初始 LR
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=10, factor=0.5)

    epochs = 200
    print("开始训练...")
    
    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()
        
        avg_loss = total_loss / len(loader)
        scheduler.step(avg_loss) # 根据 Loss 调整 LR
        
        if (epoch+1) % 10 == 0:
            lr = optimizer.param_groups[0]['lr']
            print(f"Epoch [{epoch+1}/{epochs}] | Loss: {avg_loss:.6f} | LR: {lr:.6f}")

    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("✅ 训练完成")

if __name__ == "__main__":
    train_model()