import numpy as np
import torch
import torch.nn.functional as F
from torch.distributions import Normal
from torch.nn.utils import parameters_to_vector, vector_to_parameters
from networks import PolicyNetwork, ValueNetwork

class TRPOAgent:
    def __init__(self, state_dim, action_dim, action_bound, hidden_dim=128, kl_margin=0.01, gamma=0.99, tau=0.95, cg_iters=10):
        self.gamma = gamma
        self.tau = tau
        self.kl_margin = kl_margin 
        self.cg_iters = cg_iters

        self.actor = PolicyNetwork(state_dim, action_dim, action_bound, hidden_dim)
        self.critic = ValueNetwork(state_dim, hidden_dim)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=1e-3)

    def get_action(self, state):
        """
        根据当前状态采样动作
        """
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        with torch.no_grad():
            dist = self.actor.evaluate(state_tensor)
            action = dist.sample()
        return action.squeeze(0).numpy()

    def get_value(self, state):
        with torch.no_grad():
            return self.critic(torch.FloatTensor(state).unsqueeze(0)).squeeze().item()

    def _compute_advantages(self, rewards, values, masks):
        """
        使用广义优势估计 (GAE) 计算优势函数
        """
        returns = []
        gae = 0
        for i in reversed(range(len(rewards))):
            delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
            gae = delta + self.gamma * self.tau * masks[i] * gae
            returns.insert(0, gae + values[i])
        return returns

    # --- 关键修复 1：将固定的 old_dist 作为参数传入 ---
    def _hessian_vector_product(self, states, old_dist, vector, damping=0.1):
        """
        计算海森矩阵与向量的乘积 (Hvp)
        这里的 old_dist 必须是不带梯度的历史常数分布
        """
        dist = self.actor.evaluate(states)
        # 计算当前分布与固定的旧分布之间的 KL 散度
        kl = torch.distributions.kl_divergence(old_dist, dist).mean()

        # 一阶导数
        grads = torch.autograd.grad(kl, self.actor.parameters(), create_graph=True)
        flat_grad_kl = torch.cat([grad.view(-1) for grad in grads])

        # 与给定向量点乘
        kl_v = (flat_grad_kl * vector).sum()
        
        # 二阶导数
        grads = torch.autograd.grad(kl_v, self.actor.parameters())
        flat_grad_grad_kl = torch.cat([grad.contiguous().view(-1) for grad in grads])

        # 加入阻尼系数，保证矩阵正定，避免数值不稳定发散
        return flat_grad_grad_kl + vector * damping

    # --- 关键修复 2：共轭梯度法同样接收 old_dist ---
    def _conjugate_gradient(self, states, old_dist, b, nsteps, residual_tol=1e-10):
        """
        共轭梯度法，近似求解 Hx = b
        """
        x = torch.zeros_like(b)
        r = b.clone()
        p = b.clone()
        rdotr = torch.dot(r, r)
        for _ in range(nsteps):
            Hp = self._hessian_vector_product(states, old_dist, p)
            alpha = rdotr / torch.dot(p, Hp)
            x += alpha * p
            r -= alpha * Hp
            new_rdotr = torch.dot(r, r)
            if new_rdotr < residual_tol:
                break
            p = r + new_rdotr / rdotr * p
            rdotr = new_rdotr
        return x

    def update(self, memory):
        states = torch.FloatTensor(np.array([m[0] for m in memory]))
        actions = torch.FloatTensor(np.array([m[1] for m in memory]))
        rewards = [m[2] for m in memory]
        next_states = torch.FloatTensor(np.array([m[3] for m in memory]))
        masks = [m[4] for m in memory]

        with torch.no_grad():
            values = self.critic(states).squeeze().numpy().tolist()
            next_value = self.critic(next_states[-1].unsqueeze(0)).squeeze().item()
            values.append(next_value)
        
        returns = self._compute_advantages(rewards, values, masks)
        returns = torch.FloatTensor(returns)
        values = torch.FloatTensor(values[:-1])
        advantages = returns - values
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # 【修复】：加大 Critic 的训练力度，从 10 提升到 40 Epochs
        # 确保裁判的眼光足够准确
        for _ in range(40): 
            critic_loss = F.mse_loss(self.critic(states).squeeze(), returns)
            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            self.critic_optimizer.step()

        with torch.no_grad():
            old_mean, old_std = self.actor(states)
            old_dist = Normal(old_mean, old_std)
            old_log_probs = old_dist.log_prob(actions).sum(dim=1)

        def compute_surrogate_loss():
            dist = self.actor.evaluate(states)
            log_probs = dist.log_prob(actions).sum(dim=1)
            ratio = torch.exp(log_probs - old_log_probs)
            surrogate_loss = (ratio * advantages).mean()
            return surrogate_loss, dist

        surrogate_loss, dist = compute_surrogate_loss()
        
        loss_grad = torch.autograd.grad(surrogate_loss, self.actor.parameters())
        loss_grad_flat = torch.cat([grad.view(-1) for grad in loss_grad])

        step_dir = self._conjugate_gradient(states, old_dist, loss_grad_flat, self.cg_iters)

        shs = 0.5 * torch.dot(step_dir, self._hessian_vector_product(states, old_dist, step_dir))
        
        if shs < 1e-8:
            return
            
        lm = torch.sqrt(shs / self.kl_margin)
        fullstep = step_dir / lm
        
        old_params = parameters_to_vector(self.actor.parameters())

        # 线性搜索
        success = False
        step_size = 1.0
        for _ in range(10):
            new_params = old_params + step_size * fullstep
            vector_to_parameters(new_params, self.actor.parameters())
            
            with torch.no_grad():
                new_surrogate_loss, new_dist = compute_surrogate_loss()
                kl = torch.distributions.kl_divergence(old_dist, new_dist).mean()
            
            # 【修复】：增加极小的浮点数宽容度，防止在极小提升时被误判失败而拒绝更新
            if new_surrogate_loss >= surrogate_loss - 1e-8 and kl <= self.kl_margin * 1.5:
                success = True
                break
            
            step_size *= 0.5
            
        if not success:
            vector_to_parameters(old_params, self.actor.parameters())