# MDP五元组
class MDP:
    S: List[State]      # 状态空间
    A: List[Action]     # 动作空间
    P: Dict             # 转移概率 P(s'|s,a)
    R: Callable         # 奖励函数 R(s,a,s')
    gamma: float         # 折扣因子 γ ∈ [0,1]

def value_iteration(mdp, theta=1e-6, gamma=0.9):
    V = {s: 0 for s in mdp.S}
    
    while True:
        delta = 0
        for s in mdp.S:
            v = V[s]
            V[s] = max(sum(
                mdp.P(s'|s,a) * (mdp.R(s,a,s') + gamma * V[s'])
                for s' in mdp.S
            ) for a in mdp.A)
            delta = max(delta, abs(v - V[s]))
        
        if delta < theta:
            break
    
    # 提取最优策略
    policy = {}
    for s in mdp.S:
        policy[s] = argmax_a sum(
            mdp.P(s'|s,a) * (mdp.R(s,a,s') + gamma * V[s'])
            for s' in mdp.S
        )
    
    return V, policy

def policy_iteration(mdp, gamma=0.9):
    # 随机初始化策略
    policy = {s: random.choice(mdp.A) for s in mdp.S}
    
    while True:
        # 策略评估：求解 V^π
        V = evaluate_policy(policy, mdp, gamma)
        
        # 策略改进
        policy_stable = True
        for s in mdp.S:
            old_action = policy[s]
            policy[s] = argmax_a sum(
                mdp.P(s'|s,a) * (mdp.R(s,a,s') + gamma * V[s'])
                for s' in mdp.S
            )
            if old_action != policy[s]:
                policy_stable = False
        
        if policy_stable:
            break
    
    return V, policy

def first_visit_mc(env, num_episodes, gamma=0.9, epsilon=0.1):
    Q = {s: {a: 0 for a in env.A} for s in env.S}
    returns = {s: {a: [] for a in env.A} for s in env.S}
    
    for episode in range(num_episodes):
        episode_history = []
        state = env.reset()
        
        while not env.done:
            # Epsilon-greedy策略
            if random.random() < epsilon:
                action = random.choice(env.A)
            else:
                action = argmax_a Q[state][a]
            
            next_state, reward, done = env.step(action)
            episode_history.append((state, action, reward))
            state = next_state
        
        # 计算回报
        G = 0
        for t in reversed(range(len(episode_history))):
            s, a, r = episode_history[t]
            G = r + gamma * G
            returns[s][a].append(G)
            Q[s][a] = mean(returns[s][a])
    
    return Q

def td_zero(env, num_steps, alpha=0.1, gamma=0.9, epsilon=0.1):
    V = {s: 0 for s in env.S}
    
    for _ in range(num_steps):
        state = env.reset()
        done = False
        
        while not done:
            if random.random() < epsilon:
                action = random.choice(env.A)
            else:
                action = argmax_a Q(state, a)  # 需要Q值
    
    return V

def sarsa(env, num_episodes, alpha=0.1, gamma=0.9, epsilon=0.1):
    Q = defaultdict(lambda: defaultdict(float))
    
    for episode in range(num_episodes):
        state = env.reset()
        action = epsilon_greedy(Q, state, epsilon)
        done = False
        
        while not done:
            next_state, reward, done = env.step(action)
            next_action = epsilon_greedy(Q, next_state, epsilon)
            
            if done:
                td_target = reward
            else:
                td_target = reward + gamma * Q[next_state][next_action]
            
            td_error = td_target - Q[state][action]
            Q[state][action] += alpha * td_error
            
            state = next_state
            action = next_action
    
    return Q

def q_learning(env, num_episodes, alpha=0.1, gamma=0.9, epsilon=0.1):
    Q = defaultdict(lambda: defaultdict(float))
    
    for episode in range(num_episodes):
        state = env.reset()
        done = False
        
        while not done:
            action = epsilon_greedy(Q, state, epsilon)
            next_state, reward, done = env.step(action)
            
            # 关键区别：使用max_a' Q(s',a')
            td_target = reward + gamma * max(Q[next_state].values())
            td_error = td_target - Q[state][action]
            Q[state][action] += alpha * td_error
            
            state = next_state
    
    return Q

class DQN:
    def __init__(self, state_dim, action_dim):
        self.q_network = NeuralNetwork(state_dim, action_dim)
        self.target_network = NeuralNetwork(state_dim, action_dim)
        self.replay_buffer = ReplayBuffer(capacity=100000)
        self.gamma = 0.99
        self.epsilon = 1.0
    
    def train_step(self, batch_size=32):
        states, actions, rewards, next_states, dones = self.replay_buffer.sample(batch_size)
        
        # 使用目标网络计算目标Q值
        with torch.no_grad():
            next_q = self.target_network(next_states).max(1)[0]
            target_q = rewards + (1 - dones) * self.gamma * next_q
        
        # 更新Q网络
        current_q = self.q_network(states).gather(1, actions)
        loss = MSE(current_q, target_q)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
    
    def update_target_network(self):
        # 定期同步目标网络
        self.target_network.load_state_dict(self.q_network.state_dict())

def reinforce(env, num_episodes, gamma=0.99, lr=0.001):
    policy_network = PolicyNetwork()
    optimizer = torch.optim.Adam(policy_network.parameters(), lr=lr)
    
    for episode in range(num_episodes):
        log_probs = []
        rewards = []
        state = env.reset()
        done = False
        
        while not done:
            action, log_prob = policy_network.sample(state)
            log_probs.append(log_prob)
            state, reward, done = env.step(action)
            rewards.append(reward)
        
        # 计算回报
        G = 0
        returns = []
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)
        
        # 策略梯度更新
        policy_loss = []
        for log_prob, G in zip(log_probs, returns):
            policy_loss.append(-log_prob * G)
        
        optimizer.zero_grad()
        sum(policy_loss).backward()
        optimizer.step()

class ActorCritic:
    def __init__(self, state_dim, action_dim):
        self.actor = PolicyNetwork(state_dim, action_dim)
        self.critic = ValueNetwork(state_dim)
    
    def update(self, state, action, reward, next_state, done, gamma=0.99):
        # Critic更新：TD(0)
        value = self.critic(state)
        with torch.no_grad():
            next_value = self.critic(next_state)
            td_target = reward + (1 - done) * gamma * next_value
        critic_loss = F.mse_loss(value, td_target)
        
        # Actor更新：使用TD误差作为advantage
        advantage = td_target - value
        action_log_prob = self.actor.log_prob(state, action)
        actor_loss = -(action_log_prob * advantage.detach())
        
        # 联合更新
        self.actor_optimizer.zero_grad()
        self.critic_optimizer.zero_grad()
        (actor_loss + critic_loss).backward()
        self.actor_optimizer.step()
        self.critic_optimizer.step()