import os
os.environ['MKL_THREADING_LAYER'] = 'GNU'
os.environ["OMP_NUM_THREADS"] = "8"
os.environ["MKL_NUM_THREADS"] = "8"
import torch
import numpy as np
from torch.nn import Module
import gymnasium as gym
from gym_env_250326 import CustomEnv
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import SubprocVecEnv
from stable_baselines3.common.vec_env import DummyVecEnv
from sb3_contrib import RecurrentPPO
from torch.utils.tensorboard import SummaryWriter
import pickle

# 自定义模块导入
from torch import nn
from stable_baselines3.common.policies import ActorCriticPolicy
from stable_baselines3.common.utils import obs_as_tensor
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback
from stable_baselines3.common.buffers import RolloutBuffer
from gymnasium import spaces
import types

torch.backends.cudnn.enabled = False

current_dir = os.path.dirname(os.path.abspath("__file__"))
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))

class Sin(Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.sin(x)
    
class RewardAwareEnvironmentWrapper(gym.Wrapper):
    """
    环境包装器，跟踪奖励历史并传递给特征提取器
    """
    def __init__(self, env):
        super().__init__(env)
        self.reward_history = []
        self.max_reward_history = 60 
        
    def reset(self, **kwargs):
        self.reward_history = []
        return self.env.reset(**kwargs)
    
    def step(self, action):
        # 修复：适配新的Gymnasium API (5个返回值)
        result = self.env.step(action)
        
        # 检查返回值的数量以兼容不同版本
        if len(result) == 5:
            # 新版本 Gymnasium: obs, reward, terminated, truncated, info
            obs, reward, terminated, truncated, info = result
            done = terminated or truncated  # 合并terminated和truncated为done
        else:
            # 旧版本 Gym: obs, reward, done, info
            obs, reward, done, info = result
        
        # 记录奖励历史
        self.reward_history.append(reward)
        if len(self.reward_history) > self.max_reward_history:
            self.reward_history.pop(0)
        
        # 将奖励历史添加到info中，供特征提取器使用
        info['reward_history'] = self.reward_history.copy()
        info['current_reward'] = reward
        
        # 返回与原环境相同格式的值
        if len(result) == 5:
            return obs, reward, terminated, truncated, info
        else:
            return obs, reward, done, info

class MultiTimeScaleLSTMExtractor(nn.Module):
    def __init__(self, observation_space, features_dim=32):
        super().__init__()
        self.n_obs = observation_space.shape[0]  # 总共14个观测量
        
        self.delayed_indices = list(range(0, 8))
        self.current_indices = list(range(2, 8))           
        self.leading_indices = list(range(8, self.n_obs))
        # self.delayed_indices = []
        # self.current_indices = list(range(0, 6))
        # self.leading_indices = list(range(6, self.n_obs))

        if len(self.leading_indices) > 0:
            self.leading_seq_length = 30
            self.leading_lstm = nn.LSTM(
                input_size=len(self.leading_indices),
                hidden_size=16,
                num_layers=1,
                batch_first=True,
                dropout=0.0
            )
            self.leading_mlp = nn.Sequential(
                nn.Linear(16, 16),
                Sin()
            )

        # LSTM分支 - 处理时间延迟观测量
        if len(self.delayed_indices) > 0:
            self.delayed_seq_length = 60
            self.delayed_lstm = nn.LSTM(
                input_size=len(self.delayed_indices),
                hidden_size=8,
                num_layers=1,
                batch_first=True,
                dropout=0.0
            )
            self.delayed_mlp = nn.Sequential(
                nn.Linear(8, 8),
                Sin()
            )

        # MLP分支 - 处理当前观测量
        if len(self.current_indices) > 0:
            current_obs_count = len(self.current_indices)
            self.current_mlp = nn.Sequential(
                nn.Linear(current_obs_count, 16),
                Sin(),
            )

        # 奖励历史LSTM - 新增
        self.reward_seq_length = 30
        self.reward_lstm = nn.LSTM(
            input_size=1,  # 奖励是标量
            hidden_size=8,
            num_layers=1,
            batch_first=True
        )
        self.reward_mlp = nn.Sequential(
            nn.Linear(8, 8),
            Sin()
        )

        # 简化注意力机制 - 降低复杂度
        attention_dim = 16  # 统一注意力维度
        
        # 将不同分支的输出投影到统一维度
        self.leading_proj = nn.Linear(16, attention_dim) if len(self.leading_indices) > 0 else None
        self.delayed_proj = nn.Linear(8, attention_dim) if len(self.delayed_indices) > 0 else None  
        self.current_proj = nn.Linear(16, attention_dim) if len(self.current_indices) > 0 else None
        self.reward_proj = nn.Linear(8, attention_dim)  # 奖励分支投影

        # 时间注意力机制 - 学习不同时间尺度的重要性
        self.temporal_attention = nn.MultiheadAttention(
            embed_dim=attention_dim,
            num_heads=2,
            batch_first=True
        )

        # 融合层 - 将所有分支的特征融合
        num_branches = sum([len(self.leading_indices) > 0, 
                           len(self.delayed_indices) > 0, 
                           len(self.current_indices) > 0]) + 1
        combined_size = num_branches * attention_dim  # 每个分支16维
        
        self.fusion = nn.Sequential(
            nn.Linear(combined_size, features_dim),  # 直接输出到目标维度
            Sin()                                    # 只用一层
        )
        
        self.features_dim = features_dim
        
        # 添加记忆缓冲区
        self.leading_memory = None
        self.delayed_memory = None
        self.reward_memory = None

        # 当前奖励存储（用于传递给update_reward_memory）
        self.current_reward = 0.0

    def update_reward_memory(self, reward, batch_size, device):
        """更新奖励记忆"""
        reward_tensor = torch.full((batch_size, 1), reward, device=device, dtype=torch.float32)
        
        if self.reward_memory is None or self.reward_memory.shape[0] != batch_size:
            self.reward_memory = torch.zeros(
                (batch_size, self.reward_seq_length, 1),
                device=device, dtype=torch.float32
            )
            for i in range(self.reward_seq_length):
                self.reward_memory[:, i, :] = reward_tensor
        else:
            self.reward_memory = torch.roll(self.reward_memory, shifts=-1, dims=1)
            self.reward_memory[:, -1, :] = reward_tensor

    def forward(self, observations):
        # 处理观测值，创建或更新记忆缓冲区
        if len(observations.shape) == 2:
            batch_size, n_obs = observations.shape
            
            # 管理超前信号记忆
            if self.leading_memory is None or self.leading_memory.shape[0] != batch_size:
                self.leading_memory = torch.zeros(
                    (batch_size, self.leading_seq_length, len(self.leading_indices)),
                    device=observations.device, dtype=observations.dtype
                )
                for i in range(self.leading_seq_length):
                    self.leading_memory[:, i, :] = observations[:, self.leading_indices]
            else:
                self.leading_memory = torch.roll(self.leading_memory, shifts=-1, dims=1)
                self.leading_memory[:, -1, :] = observations[:, self.leading_indices]
            
            # 管理滞后信号记忆
            if self.delayed_memory is None or self.delayed_memory.shape[0] != batch_size:
                self.delayed_memory = torch.zeros(
                    (batch_size, self.delayed_seq_length, len(self.delayed_indices)),
                    device=observations.device, dtype=observations.dtype
                )
                for i in range(self.delayed_seq_length):
                    self.delayed_memory[:, i, :] = observations[:, self.delayed_indices]
            else:
                self.delayed_memory = torch.roll(self.delayed_memory, shifts=-1, dims=1)
                self.delayed_memory[:, -1, :] = observations[:, self.delayed_indices]
        
            # 管理奖励记忆 - 使用存储的当前奖励
            self.update_reward_memory(self.current_reward, batch_size, observations.device)

        features = []
        
        # 处理超前观测量
        if len(self.leading_indices) > 0:
            _, (leading_hidden, _) = self.leading_lstm(self.leading_memory)
            leading_features = self.leading_mlp(leading_hidden[-1])  # 取最后一层的隐状态
            leading_features = self.leading_proj(leading_features)   # 投影到统一维度
            features.append(leading_features)
        
        # 处理滞后观测量
        if len(self.delayed_indices) > 0:
            _, (delayed_hidden, _) = self.delayed_lstm(self.delayed_memory)
            delayed_features = self.delayed_mlp(delayed_hidden[-1])
            delayed_features = self.delayed_proj(delayed_features)   # 投影到统一维度
            features.append(delayed_features)
        
        # 处理当前观测量
        if len(self.current_indices) > 0:
            current_obs = observations[:, self.current_indices]
            current_features = self.current_mlp(current_obs)
            current_features = self.current_proj(current_features)   # 投影到统一维度
            features.append(current_features)

        # 处理奖励历史特征
        if self.reward_memory is not None:
            _, (reward_hidden, _) = self.reward_lstm(self.reward_memory)
            reward_features = self.reward_mlp(reward_hidden[-1])
            reward_features = self.reward_proj(reward_features)
            features.append(reward_features)
        
        # 应用时间注意力机制
        if len(features) > 1:
            # 将特征重新排列为注意力机制的输入格式
            stacked_features = torch.stack(features, dim=1)  # [batch_size, num_branches, feature_dim]
            attended_features, _ = self.temporal_attention(
                stacked_features, stacked_features, stacked_features
            )
            # 修复：使用reshape而不是view，或使用contiguous().view()
            combined_features = attended_features.reshape(attended_features.shape[0], -1)
        else:
            combined_features = torch.cat(features, dim=1)
        
        return self.fusion(combined_features)

class MlpExtractor(nn.Module):
    """
    自定义的MLP特征提取器，添加了SB3所需的forward_actor和forward_critic方法
    """
    def __init__(self, feature_dim, latent_dim_pi=16, latent_dim_vf=16):
        super().__init__()
        self.latent_dim_pi = latent_dim_pi
        self.latent_dim_vf = latent_dim_vf
        
        # 创建actor和critic网络
        self.policy_net = nn.Sequential(
            nn.Linear(feature_dim, 32),
            Sin(),
            nn.Linear(32, latent_dim_pi),
            Sin()
        )
        
        self.value_net = nn.Sequential(
            nn.Linear(feature_dim, 32),
            Sin(),
            nn.Linear(32, latent_dim_vf),
            Sin()
        )
    
    def forward(self, features):
        """同时提取actor和critic特征"""
        return self.policy_net(features), self.value_net(features)
    
    def forward_actor(self, features):
        """仅提取actor特征"""
        return self.policy_net(features)
    
    def forward_critic(self, features):
        """仅提取critic特征"""
        return self.value_net(features)

class CustomActorCriticPolicy(ActorCriticPolicy):
    def __init__(self, observation_space, action_space, lr_schedule, **kwargs):
        # 移除网络相关的关键字参数
        features_extractor_kwargs = kwargs.pop("features_extractor_kwargs", {})
        features_extractor_kwargs["features_dim"] = 32
        
        super().__init__(
            observation_space,
            action_space,
            lr_schedule,
            net_arch=[],  # 使用空列表而不是None
            activation_fn=Sin,
            features_extractor_class=MultiTimeScaleLSTMExtractor,
            features_extractor_kwargs=features_extractor_kwargs,
            **kwargs
        )
    
    def _build_mlp_extractor(self):
        # 使用自定义的MlpExtractor替代默认的
        features_dim = self.features_extractor.features_dim
        self.mlp_extractor = MlpExtractor(
            feature_dim=features_dim,
            latent_dim_pi=16,
            latent_dim_vf=16

        )
class RewardTrackingPPO(PPO):
    """
    扩展PPO以传递奖励信息给特征提取器
    """
    def collect_rollouts(
        self,
        env: GymEnv,
        callback: MaybeCallback,
        rollout_buffer: RolloutBuffer,
        n_rollout_steps: int,
    ) -> bool:
        """
        重写collect_rollouts方法以传递奖励信息
        """
        # 在每次收集rollout之前重置奖励
        if hasattr(self.policy.features_extractor, 'current_reward'):
            self.policy.features_extractor.current_reward = 0.0
        
        assert self._last_obs is not None, "No previous observation was provided"
        
        # Switch to eval mode (this affects batch norm / dropout)
        self.policy.set_training_mode(False)

        n_steps = 0
        rollout_buffer.reset()
        
        # Sample new weights for the state dependent exploration
        if self.use_sde:
            self.policy.reset_noise(env.num_envs)

        callback.on_rollout_start()

        while n_steps < n_rollout_steps:
            if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
                # Sample a new noise matrix
                self.policy.reset_noise(env.num_envs)

            with torch.no_grad():
                # Convert to pytorch tensor or to TensorDict
                obs_tensor = obs_as_tensor(self._last_obs, self.device)
                actions, values, log_probs = self.policy(obs_tensor)
            actions = actions.cpu().numpy()

            # Rescale and perform action
            clipped_actions = actions

            if isinstance(self.action_space, spaces.Box):
                if self.policy.squash_output:
                    # Unscale the actions to match env bounds
                    # if they were previously squashed (scaled in [-1, 1])
                    clipped_actions = self.policy.unscale_action(clipped_actions)
                else:
                    # Otherwise, clip the actions to avoid out of bound error
                    # as we are sampling from an unbounded Gaussian distribution
                    clipped_actions = np.clip(actions, self.action_space.low, self.action_space.high)

            new_obs, rewards, dones, infos = env.step(clipped_actions)
            
            # 更新特征提取器中的奖励信息
            if hasattr(self.policy.features_extractor, 'current_reward'):
                # 如果是向量化环境，取第一个环境的奖励
                reward_to_update = rewards[0] if isinstance(rewards, np.ndarray) else rewards
                self.policy.features_extractor.current_reward = float(reward_to_update)

            self.num_timesteps += env.num_envs

            # Give access to local variables
            callback.on_step()
            if callback.on_step() is False:
                return False

            self._update_info_buffer(infos, dones)
            n_steps += 1

            if isinstance(self.action_space, spaces.Discrete):
                # Reshape in case of discrete action
                actions = actions.reshape(-1, 1)

            # Handle timeout by bootstraping with value function
            # see GitHub issue #633
            for idx, done in enumerate(dones):
                if (
                    done
                    and infos[idx].get("terminal_observation") is not None
                    and infos[idx].get("TimeLimit.truncated", False)
                ):
                    terminal_obs = self.policy.obs_to_tensor(infos[idx]["terminal_observation"])[0]
                    with torch.no_grad():
                        terminal_value = self.policy.predict_values(terminal_obs)[0]
                    rewards[idx] += self.gamma * terminal_value

            rollout_buffer.add(
                self._last_obs,
                actions,
                rewards,
                self._last_episode_starts,
                values,
                log_probs,
            )
            self._last_obs = new_obs
            self._last_episode_starts = dones

        with torch.no_grad():
            # Compute value for the last timestep
            values = self.policy.predict_values(obs_as_tensor(new_obs, self.device))

        rollout_buffer.compute_returns_and_advantage(last_values=values, dones=dones)

        callback.on_rollout_end()

        return True
    
    def _update_reward_in_extractor(self, reward):
        """
        更新特征提取器中的当前奖励
        """
        if hasattr(self.policy.features_extractor, 'current_reward'):
            # 修复：正确处理NumPy数组和标量值
            if isinstance(reward, np.ndarray):
                # 如果是数组，取第一个元素（或者平均值，根据需要）
                reward_value = float(reward.item()) if reward.size == 1 else float(reward[0])
            else:
                # 如果是标量，直接转换
                reward_value = float(reward)
            
            self.policy.features_extractor.current_reward = reward_value
            
if __name__ == '__main__':
    # 包装环境以跟踪奖励历史
    base_env = CustomEnv(device_id=0)
    vec_env = RewardAwareEnvironmentWrapper(base_env)
    name = "d1a3o14_cloak_lstm"

    # model = PPO.load(os.path.join(parent_dir, "models", "250729", "d1a3o12_cloak_lstm.zip"), env=vec_env, device=torch.device("cuda:0"))

    model = RewardTrackingPPO(
        policy=CustomActorCriticPolicy, 
        env=vec_env, 
        device=torch.device("cuda:0"),
        n_steps=1024,
        batch_size=128,
        learning_rate=5e-4,
        gamma=0.99,
        gae_lambda=0.95,
        clip_range=0.2,
        ent_coef=0.01,
        max_grad_norm=0.5,
        verbose=0)

    writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
    max_reward = 0

    history_data = []

    for i in range(500):
        model.learn(total_timesteps=1200)
        test_env = model.get_env()
        test_obs = test_env.reset()
        list_reward = []
        episolde_data = {'actions': [], 'observations': [], 'rewards': []}

        for step in range(200):
            test_action, _states = model.predict(observation=test_obs)

            # 修复：处理环境step的返回值
            result = test_env.step(test_action)
            if len(result) == 5:
                test_obs, test_rewards, terminated, truncated, info = result
                test_dones = terminated or truncated
            else:
                test_obs, test_rewards, test_dones, info = result
            
            # 更新特征提取器中的奖励信息
            model._update_reward_in_extractor(test_rewards)

            list_reward.append(test_rewards)
            episolde_data['actions'].append(test_action[0, :])
            episolde_data['observations'].append(np.array(test_obs))
            episolde_data['rewards'].append(test_rewards)

        history_data.append(episolde_data)

        avg_reward = np.mean(list_reward[-100:])
        writer.add_scalar('Reward', np.mean(avg_reward), i)
        if avg_reward > max_reward:
            max_reward = avg_reward
            model.save(os.path.join(parent_dir, "models", "250729", name + ".zip"))