import gymnasium as gym
import numpy as np
import torch
from algorithms.sac import SAC
from utils.replay_buffer import ReplayBuffer
import yaml
import os

# 读取 YAML 配置文件
config_path = os.path.join(os.path.dirname(__file__), 'configs', 'pendulum_config.yaml')
with open(config_path, "r", encoding="utf-8") as f:
    config = yaml.safe_load(f)

def main():
    # ---------------------------------------------------------
    # 2. 实例化环境与获取维度信息
    # ---------------------------------------------------------
    env = gym.make('Pendulum-v1')
    
    # 动态获取环境的维度信息，确保算法通用性
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    max_action = float(env.action_space.high[0])
    
    print(f"环境已加载: 状态维度 {state_dim}, 动作维度 {action_dim}, 最大动作界限 {max_action}")

    # ---------------------------------------------------------
    # 3. 实例化 SAC 代理与经验回放池
    # ---------------------------------------------------------
    agent = SAC(state_dim, action_dim, max_action, config)
    replay_buffer = ReplayBuffer(state_dim, action_dim, max_size=config['buffer_size'])

    total_steps = 0 # 记录与环境交互的总步数

    # ---------------------------------------------------------
    # 4. 训练大循环
    # ---------------------------------------------------------
    for episode in range(config['max_episodes']):
        # 重置环境，获取初始状态 (Gymnasium 返回 state 和 info)
        state, _ = env.reset()
        episode_reward = 0
        
        for step in range(config['max_steps']):
            # --- a. 动作选择策略 ---
            # 强化学习工程技巧：在训练初期使用纯随机动作，收集高多样性的初始数据
            if total_steps < config['start_steps']:
                action = env.action_space.sample()
            else:
                # 预热结束后，交由 SAC 的策略网络进行带噪声的采样
                action = agent.select_action(state, evaluate=False)
            
            # --- b. 与环境交互 ---
            next_state, reward, terminated, truncated, _ = env.step(action)
            
            # 判断回合是否真正结束 (超时截断 truncated 不算做环境动力学意义上的 done)
            done = float(terminated)
            
            # --- c. 存入经验池 ---
            replay_buffer.add(state, action, reward, next_state, done)
            
            state = next_state
            episode_reward += reward
            total_steps += 1
            
            # --- d. 核心学习逻辑 ---
            # 只有当经验池里的数据量足够凑齐一个 Batch 时，才开始更新网络
            if replay_buffer.size > config['batch_size']:
                agent.update(replay_buffer, config['batch_size'])
                
            # 如果提前倒地或撞毁，结束当前回合
            if terminated or truncated:
                break
                
        # 打印当前回合的训练结果
        print(f"Episode: {episode+1:03d} | Total Steps: {total_steps:06d} | Reward: {episode_reward:.2f}")

        # --- 阶段性保存模型 (可选) ---
        if (episode + 1) % 50 == 0:
            torch.save(agent.actor.state_dict(), f"sac_actor_pendulum_ep{episode+1}.pth")
            print(f"[*] 已保存第 {episode+1} 回合的模型权重。")

    env.close()
    print("训练结束！")

if __name__ == "__main__":
    main()