class FeedForwardNN(nn.Module):

    def __init__(self, in_dim, out_dim):
        
        super(FeedForwardNN, self).__init__()

        self.layer1 = nn.Linear(in_dim, 64)
        self.layer2 = nn.Linear(64, 64)
        self.layer3 = nn.Linear(64, out_dim)

    def forward(self, obs):
        
        if isinstance(obs, np.ndarray):
            obs = torch.tensor(obs, dtype=torch.float)

        activation1 = F.relu(self.layer1(obs))
        activation2 = F.relu(self.layer2(activation1))
        output = self.layer3(activation2)

        return output

class PPO:

    def __init__(self, policy_class, env, **hyperparameters):

        # PPO 初始化用于训练的超参数
        self._init_hyperparameters(hyperparameters)

        # 提取环境信息
        self.env = env
        self.obs_dim = env.observation_space.shape[0]
        self.act_dim = env.action_space.shape[0]
        
        # 初始化演员和评论家网络
        self.actor = policy_class(self.obs_dim, self.act_dim)                                                
        self.critic = policy_class(self.obs_dim, 1)

        # 为演员和评论家初始化优化器
        self.actor_optim = Adam(self.actor.parameters(), lr=self.lr)
        self.critic_optim = Adam(self.critic.parameters(), lr=self.lr)

        # 初始化协方差矩阵，用于查询actor网络的action
        self.cov_var = torch.full(size=(self.act_dim,), fill_value=0.5)
        self.cov_mat = torch.diag(self.cov_var)

        # 这个记录器将帮助我们打印出每个迭代的摘要
        self.logger = {
            'delta_t': time.time_ns(),
            't_so_far': 0,          # 到目前为止的时间步数
            'i_so_far': 0,          # 到目前为止的迭代次数
            'batch_lens': [],       # 批次中的episodic长度
            'batch_rews': [],       # 批次中的rews回报
            'actor_losses': [],     # 当前迭代中演员网络的损失
        }

    def learn(self, total_timesteps):

        print(f"Learning... Running {self.max_timesteps_per_episode} timesteps per episode, ", end='')
        print(f"{self.timesteps_per_batch} timesteps per batch for a total of {total_timesteps} timesteps")
        t_so_far = 0 # 到目前为止仿真的时间步数
        i_so_far = 0 # 到目前为止，已运行的迭代次数
        while t_so_far < total_timesteps:                                                                  
    
            # 收集批量实验数据
            batch_obs, batch_acts, batch_log_probs, batch_rtgs, batch_lens = self.rollout()                    

            # 计算收集这一批数据的时间步数
            t_so_far += np.sum(batch_lens)

            # 增加迭代次数
            i_so_far += 1

            # 记录到目前为止的时间步数和到目前为止的迭代次数
            self.logger['t_so_far'] = t_so_far
            self.logger['i_so_far'] = i_so_far

            # 计算第k次迭代的advantage
            V, _ = self.evaluate(batch_obs, batch_acts)
            A_k = batch_rtgs - V.detach()                                                                 

            # 将优势归一化 在理论上不是必须的，但在实践中，它减少了我们优势的方差，使收敛更加稳定和快速。
            # 添加这个是因为在没有这个的情况下，解决一些环境的问题太不稳定了。
            A_k = (A_k - A_k.mean()) / (A_k.std() + 1e-10)
            
            # 在其中更新我们的网络。
            for _ in range(self.n_updates_per_iteration):  
  
                V, curr_log_probs = self.evaluate(batch_obs, batch_acts)

                # 重要性采样的权重
                ratios = torch.exp(curr_log_probs - batch_log_probs)

                surr1 = ratios * A_k
                surr2 = torch.clamp(ratios, 1 - self.clip, 1 + self.clip) * A_k

                # 计算两个网络的损失。
                actor_loss = (-torch.min(surr1, surr2)).mean()
                critic_loss = nn.MSELoss()(V, batch_rtgs)

                # 计算梯度并对actor网络进行反向传播
                # 梯度清零
                self.actor_optim.zero_grad()
                # 反向传播，产生梯度
                actor_loss.backward(retain_graph=True)
                # 通过梯度下降进行优化
                self.actor_optim.step()

                # 计算梯度并对critic网络进行反向传播
                self.critic_optim.zero_grad()
                critic_loss.backward()
                self.critic_optim.step()

                self.logger['actor_losses'].append(actor_loss.detach())
                
            self._log_summary()

            if i_so_far % self.save_freq == 0:
                torch.save(self.actor.state_dict(), './ppo_actor.pth')
                torch.save(self.critic.state_dict(), './ppo_critic.pth')

    def rollout(self):
        """
            这就是我们从实验中收集一批数据的地方。由于这是一个on-policy的算法，我们需要在每次迭代行为者/批评者网络时收集一批新的数据。
        """
        batch_obs = []
        batch_acts = []
        batch_log_probs = []
        batch_rews = []
        batch_rtgs = []
        batch_lens = []

        # 一回合的数据。追踪每一回合的奖励，在回合结束的时候会被清空，开始新的回合。
        ep_rews = []

        # 追踪到目前为止这批程序我们已经运行了多少个时间段
        t = 0 

        # 继续实验，直到我们每批运行超过或等于指定的时间步数
        while t < self.timesteps_per_batch:
            ep_rews = []  每回合收集的奖励

            # 重置环境
            obs = self.env.reset()
            done = False
            
            # 运行一个回合的最大时间为max_timesteps_per_episode的时间步数
            for ep_t in range(self.max_timesteps_per_episode):
            
                if self.render and (self.logger['i_so_far'] % self.render_every_i == 0) and len(batch_lens) == 0:
                    self.env.render()

                # 递增时间步数，到目前为止已经运行了这批程序
                t += 1

                #  追踪本批中的观察结果
                batch_obs.append(obs)

                # 计算action，并在env中执行一次step。
                # 注意，rew是奖励的简称。
                action, log_prob = self.get_action(obs)
                obs, rew, done, _ = self.env.step(action)

                # 追踪最近的奖励、action和action的对数概率
                ep_rews.append(rew)
                batch_acts.append(action)
                batch_log_probs.append(log_prob)

                if done:
                    break
                    
            # 追踪本回合的长度和奖励
            batch_lens.append(ep_t + 1)
            batch_rews.append(ep_rews)

        # 将数据重塑为函数描述中指定形状的张量，然后返回
        batch_obs = torch.tensor(batch_obs, dtype=torch.float)
        batch_acts = torch.tensor(batch_acts, dtype=torch.float)
        batch_log_probs = torch.tensor(batch_log_probs, dtype=torch.float)
        batch_rtgs = self.compute_rtgs(batch_rews)                                                              

        # 在这批中记录回合的回报和回合的长度。
        self.logger['batch_rews'] = batch_rews
        self.logger['batch_lens'] = batch_lens

        return batch_obs, batch_acts, batch_log_probs, batch_rtgs, batch_lens

    def compute_rtgs(self, batch_rews):

        batch_rtgs = []
        
        # 遍历每一回合，一个回合有一批奖励
        for ep_rews in reversed(batch_rews):
            # 到目前为止的折扣奖励
            discounted_reward = 0

            # 遍历这一回合的所有奖励。我们向后退，以便更顺利地计算每一个折现的回报
            for rew in reversed(ep_rews):
                
                discounted_reward = rew + discounted_reward * self.gamma
                batch_rtgs.insert(0, discounted_reward)

        # 将每个回合的折扣奖励的数据转换成张量
        batch_rtgs = torch.tensor(batch_rtgs, dtype=torch.float)

        return batch_rtgs

    def get_action(self, obs):
    
        mean = self.actor(obs)

        # 用上述协方差矩阵中的平均行动和标准差创建一个分布。
        dist = MultivariateNormal(mean, self.cov_mat)
        action = dist.sample()
        log_prob = dist.log_prob(action)

        return action.detach().numpy(), log_prob.detach()

    def evaluate(self, batch_obs, batch_acts):
        """
            估算每个观察值，以及最近一批actor网络迭代中的每个action的对数prob。
        """
        
        # 为每个batch_obs查询critic网络的V值。V的形状应与batch_rtgs相同。
        V = self.critic(batch_obs).squeeze()

        # 使用最近的actor网络计算批量action的对数概率。
        mean = self.actor(batch_obs)
        dist = MultivariateNormal(mean, self.cov_mat)
        log_probs = dist.log_prob(batch_acts)

        # 返回批次中每个观察值的值向量V和批次中每个动作的对数概率log_probs
        return V, log_probs