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

        :參數state_dim: 狀態空間的維度
        :參數action_dim: 動作空間的維度
        :參數max_action: 動作的最大值
        """
        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):
        """
        通過網絡正向傳播。

        :參數state: 輸入狀態
        :返回值: 動作
        """

        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):
    """
    針對DDPG算法的評論家網絡。
    """
    def __init__(self, state_dim, action_dim,use_batch_norm):
        """
        初始化評論家的值網絡。

        :參數state_dim: 狀態空間的維度
        :參數action_dim: 動作空間的維度
        """
        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):
        """
        通過網絡的正向傳播。

        :參數state:輸入狀態
        :參數action: 輸入動作
        :返回值: “狀態-動作”對的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)

"""
節選自??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算法代理。

        :參數state_dim: 狀態空間的維度
        :參數action_dim: 動作空間的維度
        :參數max_action: 動作的最大值
        """
        # [第0步]
        #初始化演員的策略網絡
        self.actor = Actor(state_dim, action_dim, max_action,use_batch_norm)
        # 使用與演員的策略網絡相同的權重初始化演員目標網絡
        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)

        #初始化評論家的值網絡
        self.critic = Critic(state_dim, action_dim,use_batch_norm)
        #使用與評論家的值網絡相同的權重初始化評論家的目標網絡
        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)

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

def select_action(self, state):
        """
        根據當前狀態選擇一個動作。

        :參數state：當前狀態
        :返回值：選擇的動作
        """
        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):
        """
        訓練DDPG代理

        :參數use_target_network: 是否使用目標網絡
        :參數use_batch_norm: 是否使用批量歸一化
        """
        if len(self.replay_buffer) < BATCH_SIZE:
            return

        # [第4步]. 從重放緩沖區中抽取一批樣本
        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))

        #評論家網絡更新#
        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步]. 計算目標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步]. 使用梯度下降來更新評論家網絡的權重
        #以最小化損失函數
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        #更新演員網絡#
        actor_loss = -self.critic(state, self.actor(state)).mean()

        # [第7步]. 使用梯度下降來更新演員網絡的權重
        #以最小化損失函數和最大化Q-value => 選擇產生最高累積獎勵的動作
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # [第8步]. 更新目標網絡
        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):
    """
    訓練DDPG代理

    :參數use_target_network: 是否使用目標網絡
    :參數use_batch_norm: 是否使用批量歸一化
    :參數num_episodes: 需要訓練的回合數
    :返回值: 回合獎勵列表
    """
    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