# ============================================================
# Unit 4: Policy Gradient (REINFORCE) for Pixelcopter-PLE-v0
# Deep Reinforcement Learning Course - Hugging Face
# 支持继续训练版本
# ============================================================

import numpy as np
import matplotlib.pyplot as plt
from collections import deque
import gymnasium as gym
from gymnasium import spaces
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import pickle
import os
from datetime import datetime

# ===== 修正的PLE环境导入和Wrapper =====
from ple.games.pixelcopter import Pixelcopter
from ple import PLE

class PLEWrapper(gym.Env):
    def __init__(self):
        super().__init__()
        self.game = Pixelcopter()
        # 只使用fps参数，移除display参数
        self.env = PLE(self.game, fps=30)
        self.env.init()
        
        # 定义观察和动作空间
        state_dim = len(self.env.getGameState())
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(state_dim,), dtype=np.float32)
        self.action_space = spaces.Discrete(len(self.env.getActionSet()))
        self.actions = self.env.getActionSet()
        
    def reset(self, seed=None):
        self.env.reset_game()
        state = np.array(list(self.env.getGameState().values()), dtype=np.float32)
        return state, {}
    
    def step(self, action):
        reward = self.env.act(self.actions[action])
        state = np.array(list(self.env.getGameState().values()), dtype=np.float32)
        terminated = self.env.game_over()
        return state, reward, terminated, False, {}

# ============================================================
# 设备配置
# ============================================================
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ============================================================
# 环境配置
# ============================================================
env = PLEWrapper()
eval_env = PLEWrapper()
s_size = env.observation_space.shape[0]
a_size = env.action_space.n
print(f"Environment: Pixelcopter-PLE")
print(f"Observation Space: {s_size}, Action Space: {a_size}")

# ============================================================
# 策略网络定义
# ============================================================
class Policy(nn.Module):
    """
    策略网络：输入状态，输出动作概率分布
    """
    def __init__(self, s_size, a_size, h_size=128):
        """
        初始化策略网络
        Args:
            s_size: 状态空间维度
            a_size: 动作空间维度  
            h_size: 隐藏层大小
        """
        super(Policy, self).__init__()
        self.fc1 = nn.Linear(s_size, h_size)
        self.fc2 = nn.Linear(h_size, h_size * 2)
        self.fc3 = nn.Linear(h_size * 2, a_size)
        
        # 添加dropout提高泛化能力
        self.dropout = nn.Dropout(0.1)

    def forward(self, x):
        """
        前向传播
        Args:
            x: 输入状态
        Returns:
            动作概率分布
        """
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return F.softmax(x, dim=1)

    def act(self, state):
        """
        根据当前策略选择动作
        Args:
            state: 当前状态
        Returns:
            action: 选择的动作
            log_prob: 该动作的对数概率（用于梯度计算）
        """
        # 转换状态为tensor并移到正确设备
        state = torch.from_numpy(state).float().unsqueeze(0).to(device)
        
        # 获取动作概率分布（保持在同一设备上）
        probs = self.forward(state)
        
        # 创建分类分布
        m = torch.distributions.Categorical(probs)
        
        # 采样动作（基于概率分布，不是贪心选择）
        action = m.sample()
        
        # 返回动作值和对数概率
        return action.item(), m.log_prob(action)

# ============================================================
# 学习率调度器
# ============================================================
class LearningRateScheduler:
    def __init__(self, optimizer, initial_lr, decay_rate=0.95, decay_episodes=5000):
        self.optimizer = optimizer
        self.initial_lr = initial_lr
        self.decay_rate = decay_rate
        self.decay_episodes = decay_episodes
        
    def step(self, episode):
        if episode > 0 and episode % self.decay_episodes == 0:
            new_lr = self.initial_lr * (self.decay_rate ** (episode // self.decay_episodes))
            for param_group in self.optimizer.param_groups:
                param_group['lr'] = new_lr
            print(f"📉 Learning rate decayed to: {new_lr:.2e}")

# ============================================================
# 改进的REINFORCE算法实现
# ============================================================
def reinforce_continued(policy, optimizer, n_training_episodes, max_t, gamma, print_every, 
                       previous_scores=[], model_path=None, lr_scheduler=None):
    """
    支持继续训练的REINFORCE算法
    Args:
        policy: 策略网络
        optimizer: 优化器
        n_training_episodes: 新增训练轮数
        max_t: 每轮最大步数
        gamma: 折扣因子
        print_every: 打印间隔
        previous_scores: 之前的训练分数
        model_path: 模型保存路径
        lr_scheduler: 学习率调度器
    Returns:
        scores: 所有得分列表（包括之前的）
    """
    scores_deque = deque(maxlen=100)  # 保存最近100轮得分
    scores = previous_scores.copy()   # 保留之前的训练历史
    
    # 如果有之前的分数，用最近的分数初始化deque
    if previous_scores:
        recent_scores = previous_scores[-100:] if len(previous_scores) >= 100 else previous_scores
        scores_deque.extend(recent_scores)
        print(f"📈 Resuming with recent average score: {np.mean(scores_deque):.2f}")
        print(f"📊 Previous best score: {max(previous_scores):.2f}")
    
    start_episode = len(previous_scores) + 1
    best_avg_score = max([np.mean(previous_scores[max(0, i-99):i+1]) for i in range(len(previous_scores))]) if previous_scores else -float('inf')
    
    print(f"🚀 Starting continued training from episode {start_episode}")
    print(f"🎯 Target: Beat previous best average score of {best_avg_score:.2f}")
    print()
    
    for i_episode in range(start_episode, start_episode + n_training_episodes):
        saved_log_probs = []  # 保存每步的log概率
        rewards = []          # 保存每步的奖励
        state, _ = env.reset()
        
        # --- 1. 收集一条完整轨迹 ---
        for t in range(max_t):
            # 根据当前策略选择动作
            action, log_prob = policy.act(state)
            saved_log_probs.append(log_prob)
            
            # 执行动作，获取下一状态和奖励
            state, reward, terminated, truncated, _ = env.step(action)
            rewards.append(reward)
            
            # 检查是否结束
            if terminated or truncated:
                break
        
        # 记录本轮总得分
        episode_score = sum(rewards)
        scores_deque.append(episode_score)
        scores.append(episode_score)
        
        # --- 2. 计算折扣回报 (Discounted Returns) ---
        returns = deque(maxlen=max_t)
        n_steps = len(rewards)
        
        # 从后往前计算累计折扣回报：G_t = r_t + γ*r_{t+1} + γ²*r_{t+2} + ...
        G = 0
        for r in reversed(rewards):
            G = r + gamma * G
            returns.appendleft(G)
        
        # 标准化回报（重要的工程技巧，提高训练稳定性）
        returns = torch.tensor(returns).to(device)
        if len(returns) > 1:  # 避免标准差为0的情况
            returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        
        # --- 3. 计算策略梯度损失 ---
        # 策略梯度定理：∇J(θ) = E[∇log π(a|s) * G_t]
        # 损失函数：L = -∑(log_prob * return) （负号因为要最大化回报）
        policy_loss = []
        for log_prob, return_val in zip(saved_log_probs, returns):
            policy_loss.append(-log_prob * return_val)
        
        # 合并所有损失
        policy_loss = torch.cat(policy_loss).sum()
        
        # --- 4. 反向传播更新参数 ---
        optimizer.zero_grad()
        policy_loss.backward()
        
        # 添加梯度裁剪以提高训练稳定性
        torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=1.0)
        
        optimizer.step()
        
        # 学习率调度
        if lr_scheduler:
            lr_scheduler.step(i_episode)
        
        # 打印训练进度
        if i_episode % print_every == 0:
            current_avg = np.mean(scores_deque)
            current_lr = optimizer.param_groups[0]['lr']
            print(f'Episode {i_episode:6d} | Avg Score: {current_avg:7.2f} | Last Score: {episode_score:7.2f} | Steps: {len(rewards):4d} | LR: {current_lr:.2e}')
            
            # 检查是否创造新纪录
            if current_avg > best_avg_score:
                best_avg_score = current_avg
                print(f"🎉 New best average score: {best_avg_score:.2f}")
                
                # 保存最佳模型
                if model_path:
                    best_model_path = model_path.replace('.pth', '_best.pth')
                    torch.save({
                        'policy_state_dict': policy.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        's_size': s_size,
                        'a_size': a_size,
                        'hidden_size': policy.fc1.out_features,
                        'scores': scores,
                        'episode': i_episode,
                        'best_avg_score': best_avg_score,
                        'timestamp': datetime.now().isoformat()
                    }, best_model_path)
                    print(f"💾 Best model saved: {best_model_path}")
            
            # 定期保存检查点
            if model_path and i_episode % (print_every * 2) == 0:
                checkpoint_path = model_path.replace('.pth', f'_checkpoint_{i_episode}.pth')
                torch.save({
                    'policy_state_dict': policy.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    's_size': s_size,
                    'a_size': a_size,
                    'hidden_size': policy.fc1.out_features,
                    'scores': scores,
                    'episode': i_episode,
                    'timestamp': datetime.now().isoformat()
                }, checkpoint_path)
                print(f"💾 Checkpoint saved: {checkpoint_path}")
            
            print()
            
    return scores

# ============================================================
# 评估函数
# ============================================================
def evaluate_policy(policy, eval_env, n_eval_episodes=10):
    """
    评估策略性能
    Args:
        policy: 训练好的策略网络
        eval_env: 评估环境
        n_eval_episodes: 评估轮数
    Returns:
        episode_rewards: 每轮奖励列表
        mean_reward: 平均奖励
        std_reward: 奖励标准差
    """
    episode_rewards = []
    
    # 设置为评估模式
    policy.eval()
    
    for i in range(n_eval_episodes):
        state, _ = eval_env.reset()
        episode_reward = 0
        done = False
        steps = 0
        
        while not done and steps < 10000:  # 添加最大步数限制
            # 评估时使用确定性策略（不采样，选择概率最大的动作）
            with torch.no_grad():
                state_tensor = torch.from_numpy(state).float().unsqueeze(0).to(device)
                probs = policy.forward(state_tensor)
                action = torch.argmax(probs, dim=1).item()
            
            state, reward, terminated, truncated, _ = eval_env.step(action)
            episode_reward += reward
            done = terminated or truncated
            steps += 1
            
        episode_rewards.append(episode_reward)
        print(f"Eval Episode {i+1:2d}: Reward = {episode_reward:7.2f} | Steps = {steps:4d}")
    
    # 恢复训练模式
    policy.train()
    
    mean_reward = np.mean(episode_rewards)
    std_reward = np.std(episode_rewards)
    
    print(f"\n{'='*50}")
    print(f"Evaluation Results:")
    print(f"Mean Reward: {mean_reward:.2f}")
    print(f"Std Reward: {std_reward:.2f}")
    print(f"Score (mean - std): {mean_reward - std_reward:.2f}")
    print(f"Required for Pixelcopter-PLE-v0: 5.0")
    print(f"{'='*50}")
    
    return episode_rewards, mean_reward, std_reward

# ============================================================
# 模型加载函数
# ============================================================
def load_model(model_path, policy, optimizer):
    """
    加载已保存的模型
    """
    if os.path.exists(model_path):
        print(f"🔄 Loading existing model from {model_path}")
        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
        
        # 加载模型参数
        policy.load_state_dict(checkpoint['policy_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        
        # 加载训练历史
        previous_scores = checkpoint.get('scores', [])
        episode = checkpoint.get('episode', 0)
        best_avg_score = checkpoint.get('best_avg_score', -float('inf'))
        
        print(f"✅ Model loaded successfully!")
        print(f"📊 Loaded {len(previous_scores)} previous training episodes")
        if previous_scores:
            print(f"🎯 Previous best score: {max(previous_scores):.2f}")
            print(f"🏆 Previous best average score: {best_avg_score:.2f}")
        
        return previous_scores, episode, best_avg_score
    else:
        print("🆕 No existing model found, starting fresh training")
        return [], 0, -float('inf')

# ============================================================
# 主训练流程
# ============================================================
if __name__ == "__main__":
    # 超参数设置 - 针对继续训练优化
    HIDDEN_SIZE = 256
    INITIAL_LEARNING_RATE = 2e-5  # 继续训练时使用较小的学习率
    N_TRAINING_EPISODES = 20000   # 继续训练的轮数
    MAX_T = 10000
    GAMMA = 0.995                 # 稍微提高折扣因子
    PRINT_EVERY = 1000
    
    # 模型路径
    MODEL_PATH = "/home/eason/Workspace/Result_DRL/reinforce_pixelcopter.pth"
    
    print("="*60)
    print("REINFORCE Continued Training for Pixelcopter-PLE-v0")
    print("="*60)
    print(f"📅 Training started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print()
    
    # 初始化策略网络和优化器
    policy = Policy(s_size, a_size, HIDDEN_SIZE).to(device)
    optimizer = optim.Adam(policy.parameters(), lr=INITIAL_LEARNING_RATE)
    
    # 初始化学习率调度器
    lr_scheduler = LearningRateScheduler(optimizer, INITIAL_LEARNING_RATE, decay_rate=0.95, decay_episodes=5000)
    
    print(f"🧠 Policy Network: {policy}")
    print(f"⚙️  Optimizer: Adam (initial_lr={INITIAL_LEARNING_RATE:.2e})")
    print(f"📈 Training Episodes: {N_TRAINING_EPISODES}")
    print(f"⏱️  Max Steps per Episode: {MAX_T}")
    print(f"💰 Discount Factor: {GAMMA}")
    print()
    
    # 尝试加载已有模型
    previous_scores, last_episode, best_avg_score = load_model(MODEL_PATH, policy, optimizer)
    
    # 更新优化器学习率（确保使用当前设定的学习率）
    for param_group in optimizer.param_groups:
        param_group['lr'] = INITIAL_LEARNING_RATE
    
    print(f"📚 Current learning rate: {INITIAL_LEARNING_RATE:.2e}")
    print()
    
    # 开始训练
    print("🚀 Starting training...")
    print("-" * 80)
    
    scores = reinforce_continued(
        policy=policy,
        optimizer=optimizer,
        n_training_episodes=N_TRAINING_EPISODES,
        max_t=MAX_T,
        gamma=GAMMA,
        print_every=PRINT_EVERY,
        previous_scores=previous_scores,
        model_path=MODEL_PATH,
        lr_scheduler=lr_scheduler
    )
    
    print("\n" + "="*60)
    print("Training Completed!")
    print("="*60)
    
    # 保存最终模型
    final_model_data = {
        'policy_state_dict': policy.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        's_size': s_size,
        'a_size': a_size,
        'hidden_size': HIDDEN_SIZE,
        'scores': scores,
        'total_episodes': len(scores),
        'final_avg_score': np.mean(scores[-100:]) if len(scores) >= 100 else np.mean(scores),
        'training_completed_at': datetime.now().isoformat(),
        'hyperparameters': {
            'learning_rate': INITIAL_LEARNING_RATE,
            'gamma': GAMMA,
            'hidden_size': HIDDEN_SIZE,
            'max_t': MAX_T
        }
    }
    
    torch.save(final_model_data, MODEL_PATH)
    print(f"✅ Final model saved to {MODEL_PATH}")
    
    # 评估训练好的模型
    print("\n" + "="*60)
    print("Evaluating Final Policy")
    print("="*60)
    
    episode_rewards, mean_reward, std_reward = evaluate_policy(policy, eval_env, n_eval_episodes=10)
    
    # 训练结果总结
    print(f"\n🎉 Final Training Results:")
    print(f"   Total Episodes Trained: {len(scores)}")
    print(f"   Final Average Score (last 100): {np.mean(scores[-100:]) if len(scores) >= 100 else np.mean(scores):.2f}")
    print(f"   Best Single Episode Score: {max(scores):.2f}")
    print(f"   Evaluation Mean Reward: {mean_reward:.2f}")
    print(f"   Evaluation Std Reward: {std_reward:.2f}")
    print(f"   Final Score (mean - std): {mean_reward - std_reward:.2f}")
    print(f"   Required for Pixelcopter-PLE-v0: 5.0")
    
    if mean_reward - std_reward >= 5.0:
        print(f"   Status: ✅ PASSED! Congratulations!")
    else:
        needed_improvement = 5.0 - (mean_reward - std_reward)
        print(f"   Status: ❌ Need {needed_improvement:.2f} more points")
        print(f"   Suggestion: Continue training with lower learning rate or adjust network architecture")
    
    print(f"\n📅 Training completed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")