class Actor(nn.Module):
    """
    針對(duì)DDPG算法的演員網(wǎng)絡(luò)。
    """
    def __init__(self, state_dim, action_dim, max_action,use_batch_norm):
        """
        初始化演員的策略網(wǎng)絡(luò)

        :參數(shù)state_dim: 狀態(tài)空間的維度
        :參數(shù)action_dim: 動(dòng)作空間的維度
        :參數(shù)max_action: 動(dòng)作的最大值
        """
        super(Actor, self).__init__()
        self.bn1 = nn.LayerNorm(HIDDEN_LAYERS_ACTOR) if use_batch_norm else nn.Identity()
        self.bn2 = nn.LayerNorm(HIDDEN_LAYERS_ACTOR) if use_batch_norm else nn.Identity()

        self.l1 = nn.Linear(state_dim, HIDDEN_LAYERS_ACTOR)
        self.l2 = nn.Linear(HIDDEN_LAYERS_ACTOR, HIDDEN_LAYERS_ACTOR)
        self.l3 = nn.Linear(HIDDEN_LAYERS_ACTOR, action_dim)
        self.max_action = max_action

    def forward(self, state):
        """
        通過網(wǎng)絡(luò)正向傳播。

        :參數(shù)state: 輸入狀態(tài)
        :返回值: 動(dòng)作
        """

        a = torch.relu(self.bn1(self.l1(state)))
        a = torch.relu(self.bn2(self.l2(a)))
        return self.max_action * torch.tanh(self.l3(a))

class Critic(nn.Module):
    """
    針對(duì)DDPG算法的評(píng)論家網(wǎng)絡(luò)。
    """
    def __init__(self, state_dim, action_dim,use_batch_norm):
        """
        初始化評(píng)論家的值網(wǎng)絡(luò)。

        :參數(shù)state_dim: 狀態(tài)空間的維度
        :參數(shù)action_dim: 動(dòng)作空間的維度
        """
        super(Critic, self).__init__()
        self.bn1 = nn.BatchNorm1d(HIDDEN_LAYERS_CRITIC) if use_batch_norm else nn.Identity()
        self.bn2 = nn.BatchNorm1d(HIDDEN_LAYERS_CRITIC) if use_batch_norm else nn.Identity()
        self.l1 = nn.Linear(state_dim + action_dim, HIDDEN_LAYERS_CRITIC)

        self.l2 = nn.Linear(HIDDEN_LAYERS_CRITIC, HIDDEN_LAYERS_CRITIC)
        self.l3 = nn.Linear(HIDDEN_LAYERS_CRITIC, 1)

    def forward(self, state, action):
        """
        通過網(wǎng)絡(luò)的正向傳播。

        :參數(shù)state:輸入狀態(tài)
        :參數(shù)action: 輸入動(dòng)作
        :返回值: “狀態(tài)-動(dòng)作”對(duì)的Q-值
        """
        q = torch.relu(self.bn1(self.l1(torch.cat([state, action], 1))))
        q = torch.relu(self.bn2(self.l2(q)))
        return self.l3(q)

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        return random.sample(self.buffer, batch_size)

    def __len__(self):
        return len(self.buffer)

"""
節(jié)選自??https://github.com/vitchyr/rlkit/blob/master/rlkit/exploration_strategies/ou_strategy.py??
"""
class OUNoise(object):
    def __init__(self, action_space, mu=0.0, theta=0.15, max_sigma=0.3, min_sigma=0.3, decay_period=100000):
        self.mu           = mu
        self.theta        = theta
        self.sigma        = max_sigma
        self.max_sigma    = max_sigma
        self.min_sigma    = min_sigma
        self.decay_period = decay_period
        self.action_dim   = action_space.shape[0]
        self.low          = action_space.low
        self.high         = action_space.high
        self.reset()

    def reset(self):
        self.state = np.ones(self.action_dim) * self.mu

    def evolve_state(self):
        x  = self.state
        dx = self.theta * (self.mu - x) + self.sigma * np.random.randn(self.action_dim)
        self.state = x + dx
        return self.state

    def get_action(self, action, t=0): 
        ou_state = self.evolve_state()
        self.sigma = self.max_sigma - (self.max_sigma - self.min_sigma) * min(1.0, t / self.decay_period)
        return np.clip(action + ou_state, self.low, self.high)

class DDPG():
    """
    深度確定性策略梯度（DDPG）代理。
    """
    def __init__(self, state_dim, action_dim, max_action,use_batch_norm):
        """
        初始化DDPG算法代理。

        :參數(shù)state_dim: 狀態(tài)空間的維度
        :參數(shù)action_dim: 動(dòng)作空間的維度
        :參數(shù)max_action: 動(dòng)作的最大值
        """
        # [第0步]
        #初始化演員的策略網(wǎng)絡(luò)
        self.actor = Actor(state_dim, action_dim, max_action,use_batch_norm)
        # 使用與演員的策略網(wǎng)絡(luò)相同的權(quán)重初始化演員目標(biāo)網(wǎng)絡(luò)
        self.actor_target = Actor(state_dim, action_dim, max_action,use_batch_norm)
        self.actor_target.load_state_dict(self.actor.state_dict())
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=ACTOR_LR)

        #初始化評(píng)論家的值網(wǎng)絡(luò)
        self.critic = Critic(state_dim, action_dim,use_batch_norm)
        #使用與評(píng)論家的值網(wǎng)絡(luò)相同的權(quán)重初始化評(píng)論家的目標(biāo)網(wǎng)絡(luò)
        self.critic_target = Critic(state_dim, action_dim,use_batch_norm)
        self.critic_target.load_state_dict(self.critic.state_dict())
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=CRITIC_LR)

        #初始化重放緩沖區(qū)
        self.replay_buffer = ReplayBuffer(BUFFER_SIZE)

def select_action(self, state):
        """
        根據(jù)當(dāng)前狀態(tài)選擇一個(gè)動(dòng)作。

        :參數(shù)state：當(dāng)前狀態(tài)
        :返回值：選擇的動(dòng)作
        """
        state = torch.FloatTensor(state.reshape(1, -1))
        action = self.actor(state).cpu().data.numpy().flatten()
        return action

def train(self, use_target_network,use_batch_norm):
        """
        訓(xùn)練DDPG代理

        :參數(shù)use_target_network: 是否使用目標(biāo)網(wǎng)絡(luò)
        :參數(shù)use_batch_norm: 是否使用批量歸一化
        """
        if len(self.replay_buffer) < BATCH_SIZE:
            return

        # [第4步]. 從重放緩沖區(qū)中抽取一批樣本
        batch = self.replay_buffer.sample(BATCH_SIZE)
        state, action, reward, next_state, done = map(np.stack, zip(*batch))

        state = torch.FloatTensor(state)
        action = torch.FloatTensor(action)
        next_state = torch.FloatTensor(next_state)
        reward = torch.FloatTensor(reward.reshape(-1, 1))
        done = torch.FloatTensor(done.reshape(-1, 1))

        #評(píng)論家網(wǎng)絡(luò)更新#
        if use_target_network:
            target_Q = self.critic_target(next_state, self.actor_target(next_state))
        else:
            target_Q = self.critic(next_state, self.actor(next_state))

        # [第5步]. 計(jì)算目標(biāo)Q-value (y_i)
        target_Q = reward + (1 - done) * GAMMA * target_Q
        current_Q = self.critic(state, action)
        critic_loss = nn.MSELoss()(current_Q, target_Q.detach())

        # [第6步]. 使用梯度下降來更新評(píng)論家網(wǎng)絡(luò)的權(quán)重
        #以最小化損失函數(shù)
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        #更新演員網(wǎng)絡(luò)#
        actor_loss = -self.critic(state, self.actor(state)).mean()

        # [第7步]. 使用梯度下降來更新演員網(wǎng)絡(luò)的權(quán)重
        #以最小化損失函數(shù)和最大化Q-value => 選擇產(chǎn)生最高累積獎(jiǎng)勵(lì)的動(dòng)作
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # [第8步]. 更新目標(biāo)網(wǎng)絡(luò)
        if use_target_network:
            for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
                target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)

            for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
                target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)

def train_ddpg(use_target_network, use_batch_norm, num_episodes=NUM_EPISODES):
    """
    訓(xùn)練DDPG代理

    :參數(shù)use_target_network: 是否使用目標(biāo)網(wǎng)絡(luò)
    :參數(shù)use_batch_norm: 是否使用批量歸一化
    :參數(shù)num_episodes: 需要訓(xùn)練的回合數(shù)
    :返回值: 回合獎(jiǎng)勵(lì)列表
    """
    agent = DDPG(state_dim, action_dim, 1,use_batch_norm)

    episode_rewards = []
    noise = OUNoise(env.action_space)

    for episode in range(num_episodes):
        state= env.reset()
        noise.reset()
        episode_reward = 0
        done = False
        step=0
        while not done:
            action_actor = agent.select_action(state)
            action = noise.get_action(action_actor,step) # Add noise for exploration
            next_state, reward, done,_= env.step(action)
            done = float(done) if isinstance(done, (bool, int)) else float(done[0])
            agent.replay_buffer.push(state, action, reward, next_state, done)

            if len(agent.replay_buffer) > BATCH_SIZE:
                agent.train(use_target_network,use_batch_norm)

            state = next_state
            episode_reward += reward
            step+=1

        episode_rewards.append(episode_reward)

        if (episode + 1) % 10 == 0:
            print(f"Episode {episode + 1}: Reward = {episode_reward}")

    return agent, episode_rewards