import copy
import random
from abc import ABC, abstractmethod
from collections import deque
from typing import Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from zyplib.nn.modules import BaseModule


class ExperienceBuffer:
    """经验回放缓冲区

    用于存储和采样强化学习经验。

    Parameters
    ----------
    - `capacity` : `int`
        - 缓冲区容量
    """

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

    def push(
        self, state: Tensor, action: Tensor, reward: float, next_state: Tensor, done: bool
    ):
        """添加经验"""
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size: int) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        """采样经验"""
        batch = random.sample(self.buffer, batch_size)
        state, action, reward, next_state, done = map(torch.stack, zip(*batch))
        return state, action, reward, next_state, done

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


class RLBase(BaseModule, ABC):
    """强化学习基类

    提供强化学习算法的基本接口和通用功能。

    Parameters
    ----------
    - `state_dim` : `int`
        - 状态维度
    - `action_dim` : `int`
        - 动作维度
    - `lr` : `float`
        - 学习率
    - `gamma` : `float`
        - 折扣因子
    """

    def __init__(
        self, state_dim: int, action_dim: int, lr: float = 1e-3, gamma: float = 0.99
    ):
        super().__init__()
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.lr = lr
        self.gamma = gamma

    @abstractmethod
    def select_action(self, state: Tensor, training: bool = True) -> Tensor:
        """选择动作

        Parameters
        ----------
        - `state` : `Tensor`
            - 当前状态
        - `training` : `bool`, optional
            - 是否为训练模式, by default True

        Returns
        ----------
        - `Tensor`
            - 选择的动作
        """
        pass

    @abstractmethod
    def update_policy(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """更新策略

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 损失字典
        """
        pass

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算损失

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 损失字典
        """
        return self.update_policy(batch_data)


class DQN(RLBase):
    """Deep Q-Network (DQN)

    实现深度 Q 网络算法，使用神经网络来近似 Q 函数。

    References
    ----------
    Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., ... & Hassabis, D. (2015).
    Human-level control through deep reinforcement learning.
    Nature, 518(7540), 529-533.

    Parameters
    ----------
    - `q_network` : `nn.Module`
        - Q 网络
    - `state_dim` : `int`
        - 状态维度
    - `action_dim` : `int`
        - 动作维度
    - `lr` : `float`, optional
        - 学习率, by default 1e-3
    - `gamma` : `float`, optional
        - 折扣因子, by default 0.99
    - `epsilon` : `float`, optional
        - ε-贪婪策略参数, by default 0.1
    - `target_update_freq` : `int`, optional
        - 目标网络更新频率, by default 100
    - `buffer_size` : `int`, optional
        - 经验回放缓冲区大小, by default 10000
    - `batch_size` : `int`, optional
        - 批次大小, by default 32

    Examples
    ----------
    ```python
    # 创建 DQN 模型
    model = DQN(
        q_network=q_net,
        state_dim=4,
        action_dim=2,
        buffer_size=50000,
        batch_size=64
    )

    # 从内部缓冲区训练
    loss_dict = model.update_from_buffer()

    # 或使用外部数据训练
    loss_dict = model.update_policy(batch_data)
    ```
    """

    def __init__(
        self,
        q_network: nn.Module,
        state_dim: int,
        action_dim: int,
        lr: float = 1e-3,
        gamma: float = 0.99,
        epsilon: float = 0.1,
        target_update_freq: int = 100,
        buffer_size: int = 10000,
        batch_size: int = 32,
    ):
        super().__init__(state_dim, action_dim, lr, gamma)

        self.q_network = q_network
        # 使用深拷贝创建目标网络，更加稳健
        try:
            self.target_network = copy.deepcopy(q_network)
        except Exception as e:
            # 如果深拷贝失败，回退到原方法并记录警告
            print(f'Warning: deepcopy failed ({e}), using fallback method')
            self.target_network = type(q_network)(**q_network.__dict__)
            self.target_network.load_state_dict(q_network.state_dict())

        self.epsilon = epsilon
        self.target_update_freq = target_update_freq
        self.update_count = 0
        self.batch_size = batch_size

        self.buffer = ExperienceBuffer(buffer_size)
        self.optimizer = torch.optim.Adam(self.q_network.parameters(), lr=lr)

    def select_action(self, state: Tensor, training: bool = True) -> Tensor:
        """ε-贪婪动作选择

        Parameters
        ----------
        - `state` : `Tensor`
            - 当前状态，形状为 (batch_size, state_dim)
        - `training` : `bool`, optional
            - 是否为训练模式, by default True

        Returns
        ----------
        - `Tensor`
            - 选择的动作，形状为 (batch_size,)
        """
        if training and random.random() < self.epsilon:
            # 随机探索
            return torch.randint(0, self.action_dim, (state.size(0),))
        else:
            # 贪婪选择
            with torch.no_grad():
                q_values = self.q_network(state)
                return q_values.argmax(dim=1)

    def store_experience(
        self, state: Tensor, action: Tensor, reward: float, next_state: Tensor, done: bool
    ):
        """存储经验到缓冲区

        Parameters
        ----------
        - `state` : `Tensor`
            - 当前状态
        - `action` : `Tensor`
            - 执行的动作
        - `reward` : `float`
            - 获得的奖励
        - `next_state` : `Tensor`
            - 下一个状态
        - `done` : `bool`
            - 是否结束
        """
        self.buffer.push(state, action, reward, next_state, done)

    def update_from_buffer(
        self, batch_size: Optional[int] = None
    ) -> Optional[Dict[str, Tensor]]:
        """从内部缓冲区采样并更新网络

        Parameters
        ----------
        - `batch_size` : `int`, optional
            - 批次大小, by default None (使用初始化时的batch_size)

        Returns
        ----------
        - `Dict[str, Tensor]`, optional
            - 包含损失的字典，如果缓冲区数据不足则返回None
        """
        batch_size = batch_size or self.batch_size
        if len(self.buffer) < batch_size:
            return None

        states, actions, rewards, next_states, dones = self.buffer.sample(batch_size)
        batch_data = {
            'states': states,
            'actions': actions,
            'rewards': rewards,
            'next_states': next_states,
            'dones': dones,
        }
        return self.update_policy(batch_data)

    def update_policy(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """更新 Q 网络

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'states', 'actions', 'rewards', 'next_states', 'dones' 的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 Q 损失的字典
        """
        states = batch_data['states']
        actions = batch_data['actions']
        rewards = batch_data['rewards']
        next_states = batch_data['next_states']
        dones = batch_data['dones']

        # 当前 Q 值
        current_q_values = self.q_network(states).gather(1, actions.unsqueeze(1))

        # 目标 Q 值
        with torch.no_grad():
            next_q_values = self.target_network(next_states).max(1)[0]
            target_q_values = rewards + (self.gamma * next_q_values * ~dones)

        # 计算损失
        q_loss = F.mse_loss(current_q_values.squeeze(), target_q_values)

        # 更新网络
        self.optimizer.zero_grad()
        q_loss.backward()
        self.optimizer.step()

        # 更新目标网络
        self.update_count += 1
        if self.update_count % self.target_update_freq == 0:
            self.target_network.load_state_dict(self.q_network.state_dict())

        return {'q_loss': q_loss, 'total_loss': q_loss}


class PolicyGradient(RLBase):
    """策略梯度算法

    实现基础的策略梯度算法，直接优化策略函数。

    References
    ----------
    Williams, R. J. (1992).
    Simple statistical gradient-following algorithms for connectionist reinforcement learning.
    Machine learning, 8(3-4), 229-256.

    Parameters
    ----------
    - `policy_network` : `nn.Module`
        - 策略网络
    - `state_dim` : `int`
        - 状态维度
    - `action_dim` : `int`
        - 动作维度
    - `lr` : `float`, optional
        - 学习率, by default 1e-3
    - `gamma` : `float`, optional
        - 折扣因子, by default 0.99

    Examples
    ----------
    ```python
    # 创建策略梯度模型
    model = PolicyGradient(
        policy_network=policy_net,
        state_dim=4,
        action_dim=2,
        gamma=0.99
    )

    # 方式1：使用内部状态（向后兼容）
    loss_dict = model.update_policy({'rewards': rewards})

    # 方式2：直接传入数据（推荐）
    loss_dict = model.update_policy({'rewards': rewards, 'log_probs': log_probs})
    ```
    """

    def __init__(
        self,
        policy_network: nn.Module,
        state_dim: int,
        action_dim: int,
        lr: float = 1e-3,
        gamma: float = 0.99,
    ):
        super().__init__(state_dim, action_dim, lr, gamma)

        self.policy_network = policy_network
        self.optimizer = torch.optim.Adam(self.policy_network.parameters(), lr=lr)

        # 存储轨迹（用于向后兼容）
        self.log_probs = []
        self.rewards = []

    def select_action(self, state: Tensor, training: bool = True) -> Tensor:
        """根据策略选择动作

        Parameters
        ----------
        - `state` : `Tensor`
            - 当前状态
        - `training` : `bool`, optional
            - 是否为训练模式, by default True

        Returns
        ----------
        - `Tensor`
            - 选择的动作
        """
        probs = self.policy_network(state)
        m = torch.distributions.Categorical(probs)
        action = m.sample()

        if training:
            self.log_probs.append(m.log_prob(action))

        return action

    def update_policy(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """更新策略网络

        支持两种调用方式：
        1. 仅传入rewards，使用内部存储的log_probs（向后兼容）
        2. 同时传入rewards和log_probs（推荐方式，更清晰）

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'rewards' 的字典，可选包含 'log_probs'

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含策略损失的字典
        """
        rewards = batch_data['rewards']

        # 检查是否提供了外部log_probs
        if 'log_probs' in batch_data:
            log_probs = batch_data['log_probs']
        else:
            # 使用内部状态（向后兼容）
            if not self.log_probs:
                raise ValueError(
                    'No log_probs available. Either call select_action first or provide log_probs in batch_data.'
                )
            log_probs = self.log_probs

        # 计算折扣奖励
        discounted_rewards = []
        discounted_reward = 0
        for reward in reversed(rewards):
            discounted_reward = reward + self.gamma * discounted_reward
            discounted_rewards.insert(0, discounted_reward)

        discounted_rewards = torch.tensor(discounted_rewards)
        discounted_rewards = (discounted_rewards - discounted_rewards.mean()) / (
            discounted_rewards.std() + 1e-9
        )

        # 计算策略损失
        policy_loss = []
        for log_prob, reward in zip(log_probs, discounted_rewards):
            policy_loss.append(-log_prob * reward)

        policy_loss = torch.cat(policy_loss).sum()

        # 更新网络
        self.optimizer.zero_grad()
        policy_loss.backward()
        self.optimizer.step()

        # 清空轨迹（如果使用内部状态）
        if 'log_probs' not in batch_data:
            self.log_probs = []
            self.rewards = []

        return {'policy_loss': policy_loss, 'total_loss': policy_loss}


class ActorCritic(RLBase):
    """Actor-Critic 算法

    实现 Actor-Critic 算法，结合策略梯度和价值函数方法。

    References
    ----------
    Konda, V. R., & Tsitsiklis, J. N. (2000).
    Actor-critic algorithms.
    In Advances in neural information processing systems (pp. 1008-1014).

    Parameters
    ----------
    - `actor_network` : `nn.Module`
        - Actor 网络（策略网络）
    - `critic_network` : `nn.Module`
        - Critic 网络（价值网络）
    - `state_dim` : `int`
        - 状态维度
    - `action_dim` : `int`
        - 动作维度
    - `lr_actor` : `float`, optional
        - Actor 学习率, by default 1e-3
    - `lr_critic` : `float`, optional
        - Critic 学习率, by default 1e-3
    - `gamma` : `float`, optional
        - 折扣因子, by default 0.99

    Examples
    ----------
    ```python
    # 创建 Actor-Critic 模型
    model = ActorCritic(
        actor_network=actor_net,
        critic_network=critic_net,
        state_dim=4,
        action_dim=2,
        gamma=0.99
    )

    # 训练
    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        actor_network: nn.Module,
        critic_network: nn.Module,
        state_dim: int,
        action_dim: int,
        lr_actor: float = 1e-3,
        lr_critic: float = 1e-3,
        gamma: float = 0.99,
    ):
        super().__init__(state_dim, action_dim, lr_actor, gamma)

        self.actor = actor_network
        self.critic = critic_network

        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=lr_critic)

    def select_action(self, state: Tensor, training: bool = True) -> Tensor:
        """根据 Actor 选择动作

        Parameters
        ----------
        - `state` : `Tensor`
            - 当前状态
        - `training` : `bool`, optional
            - 是否为训练模式, by default True

        Returns
        ----------
        - `Tensor`
            - 选择的动作
        """
        if training:
            probs = self.actor(state)
            m = torch.distributions.Categorical(probs)
            action = m.sample()
        else:
            with torch.no_grad():
                probs = self.actor(state)
                action = probs.argmax(dim=-1)
        return action

    def update_policy(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """更新 Actor 和 Critic 网络

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'states', 'actions', 'rewards', 'next_states', 'dones' 的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 Actor 和 Critic 损失的字典
        """
        states = batch_data['states']
        actions = batch_data['actions']
        rewards = batch_data['rewards']
        next_states = batch_data['next_states']
        dones = batch_data['dones']

        # 计算价值和下一状态价值
        values = self.critic(states).squeeze()
        next_values = self.critic(next_states).squeeze()

        # 计算 TD 目标和优势
        td_targets = rewards + self.gamma * next_values * ~dones
        advantages = td_targets - values

        # Actor 损失
        probs = self.actor(states)
        m = torch.distributions.Categorical(probs)
        log_probs = m.log_prob(actions)
        actor_loss = -(log_probs * advantages.detach()).mean()

        # Critic 损失
        critic_loss = F.mse_loss(values, td_targets.detach())

        # 更新网络
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        return {
            'actor_loss': actor_loss,
            'critic_loss': critic_loss,
            'total_loss': actor_loss + critic_loss,
        }


class PPO(RLBase):
    """Proximal Policy Optimization (PPO)

    实现近端策略优化算法，通过限制策略更新幅度来提高训练稳定性。
    支持 GAE (Generalized Advantage Estimation) 来更好地估计优势函数。

    References
    ----------
    Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017).
    Proximal policy optimization algorithms.
    arXiv preprint arXiv:1707.06347.

    Parameters
    ----------
    - `actor_network` : `nn.Module`
        - Actor 网络（策略网络）
    - `critic_network` : `nn.Module`
        - Critic 网络（价值网络）
    - `state_dim` : `int`
        - 状态维度
    - `action_dim` : `int`
        - 动作维度
    - `lr_actor` : `float`, optional
        - Actor 学习率, by default 3e-4
    - `lr_critic` : `float`, optional
        - Critic 学习率, by default 1e-3
    - `gamma` : `float`, optional
        - 折扣因子, by default 0.99
    - `clip_epsilon` : `float`, optional
        - PPO 裁剪比率, by default 0.2
    - `k_epochs` : `int`, optional
        - 每次更新的训练轮数, by default 4
    - `gae_lambda` : `float`, optional
        - GAE lambda 参数, by default 0.95
    - `use_gae` : `bool`, optional
        - 是否使用 GAE, by default True

    Examples
    ----------
    ```python
    # 创建 PPO 模型
    model = PPO(
        actor_network=actor_net,
        critic_network=critic_net,
        state_dim=4,
        action_dim=2,
        gamma=0.99,
        clip_epsilon=0.2,
        use_gae=True
    )

    # 训练
    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        actor_network: nn.Module,
        critic_network: nn.Module,
        state_dim: int,
        action_dim: int,
        lr_actor: float = 3e-4,
        lr_critic: float = 1e-3,
        gamma: float = 0.99,
        clip_epsilon: float = 0.2,
        k_epochs: int = 4,
        gae_lambda: float = 0.95,
        use_gae: bool = True,
    ):
        super().__init__(state_dim, action_dim, lr_actor, gamma)

        self.actor = actor_network
        self.critic = critic_network
        self.clip_epsilon = clip_epsilon
        self.k_epochs = k_epochs
        self.gae_lambda = gae_lambda
        self.use_gae = use_gae

        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=lr_critic)

    def select_action(self, state: Tensor, training: bool = True) -> Tensor:
        """根据 Actor 选择动作

        Parameters
        ----------
        - `state` : `Tensor`
            - 当前状态
        - `training` : `bool`, optional
            - 是否为训练模式, by default True

        Returns
        ----------
        - `Tensor`
            - 选择的动作
        """
        if training:
            probs = self.actor(state)
            m = torch.distributions.Categorical(probs)
            action = m.sample()
        else:
            with torch.no_grad():
                probs = self.actor(state)
                action = probs.argmax(dim=-1)
        return action

    def compute_gae(
        self, rewards: Tensor, values: Tensor, next_values: Tensor, dones: Tensor
    ) -> Tuple[Tensor, Tensor]:
        """计算 GAE (Generalized Advantage Estimation)

        Parameters
        ----------
        - `rewards` : `Tensor`
            - 奖励序列
        - `values` : `Tensor`
            - 状态价值序列
        - `next_values` : `Tensor`
            - 下一状态价值序列
        - `dones` : `Tensor`
            - 结束标志序列

        Returns
        ----------
        - `Tuple[Tensor, Tensor]`
            - (advantages, returns) 优势和回报
        """
        # 性能优化: 预分配 advantages 张量以避免 O(N^2) 复杂度的列表插入
        advantages = torch.zeros_like(rewards)
        gae = 0

        # 计算 TD 误差
        deltas = rewards + self.gamma * next_values * ~dones - values

        # 反向计算 GAE, 从后往前填充 advantages 张量
        for t in reversed(range(len(deltas))):
            gae = deltas[t] + self.gamma * self.gae_lambda * gae
            advantages[t] = gae

        returns = advantages + values

        return advantages, returns

    def update_policy(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """使用 PPO 更新策略

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'states', 'actions', 'rewards', 'log_probs', 'values', 'next_states', 'dones' 的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 PPO 损失的字典
        """
        states = batch_data['states']
        actions = batch_data['actions']
        rewards = batch_data['rewards']
        old_log_probs = batch_data['log_probs']
        old_values = batch_data['values']
        next_states = batch_data.get('next_states')
        dones = batch_data.get('dones')

        # 计算优势和回报
        if self.use_gae and next_states is not None and dones is not None:
            # 使用 GAE
            next_values = self.critic(next_states).squeeze().detach()
            advantages, returns = self.compute_gae(
                rewards, old_values, next_values, dones
            )
        else:
            # 使用简单的蒙特卡洛折扣奖励
            discounted_rewards = []
            discounted_reward = 0
            for reward in reversed(rewards):
                discounted_reward = reward + self.gamma * discounted_reward
                discounted_rewards.insert(0, discounted_reward)

            returns = torch.tensor(discounted_rewards)
            advantages = returns - old_values

        # 标准化优势
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # 多轮更新
        total_actor_loss = 0
        total_critic_loss = 0

        for _ in range(self.k_epochs):
            # 当前策略的概率
            probs = self.actor(states)
            m = torch.distributions.Categorical(probs)
            new_log_probs = m.log_prob(actions)

            # 重要性采样比率
            ratio = torch.exp(new_log_probs - old_log_probs.detach())

            # PPO 裁剪目标
            surr1 = ratio * advantages.detach()
            surr2 = (
                torch.clamp(ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon)
                * advantages.detach()
            )
            actor_loss = -torch.min(surr1, surr2).mean()

            # Critic 损失
            values = self.critic(states).squeeze()
            critic_loss = F.mse_loss(values, returns.detach())

            # 更新网络
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            self.critic_optimizer.step()

            total_actor_loss += actor_loss.item()
            total_critic_loss += critic_loss.item()

        return {
            'actor_loss': torch.tensor(total_actor_loss / self.k_epochs),
            'critic_loss': torch.tensor(total_critic_loss / self.k_epochs),
            'total_loss': torch.tensor(
                (total_actor_loss + total_critic_loss) / self.k_epochs
            ),
        }
