from dataclasses import asdict, dataclass, field
from typing import Any

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import wandb


@dataclass
class A2CConfig:
    """Configuration for A2C algorithm."""

    # Environment config
    env_name: str = "CartPole-v1"
    render_mode: str | None = None
    solved_threshold: float = 475.0

    # Training config
    gamma: float = 0.999
    lambda_: float = 0.98
    value_loss_coef: float = 0.5
    policy_loss_coef: float = 0.5
    entropy_coef: float = 0.01
    num_episodes: int = 500
    learning_rate: float = 0.01

    # Evaluation config
    eval_episodes: int = 10
    eval_interval: int = 100
    log_wandb: bool = False


class ActorCriticNet(nn.Module):
    def __init__(self, env_dim: int, action_num: int) -> None:
        super().__init__()
        self.env_dim = env_dim
        self.action_num = action_num
        self.shared_layers = nn.Sequential(
            nn.Linear(env_dim, 64),
            nn.ReLU(),
        )
        self.policy_head = nn.Linear(64, action_num)
        self.value_head = nn.Linear(64, 1)

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        x = self.shared_layers(x)
        policy_action_logits = self.policy_head(x)
        v_value = self.value_head(x)
        return policy_action_logits, v_value


@dataclass
class Experience:
    observation: Any
    action: Any
    reward: float
    next_observation: Any
    terminated: bool
    truncated: bool


@dataclass
class ReplayBuffer:
    buffer: list[Experience] = field(default_factory=list)

    def __len__(self) -> int:
        return len(self.buffer)

    def add_experience(self, experience: Experience) -> None:
        self.buffer.append(experience)

    def reset(self) -> None:
        self.buffer = []

    def sample(self) -> list[Experience]:
        return self.buffer

    def total_reward(self) -> float:
        return sum(experience.reward for experience in self.buffer)


class Agent:
    def __init__(self, net: nn.Module, optimizer: torch.optim.Optimizer) -> None:
        self._net = net
        self._optimizer = optimizer
        self._replay_buffer = ReplayBuffer()

    def onpolicy_reset(self) -> None:
        self._replay_buffer.reset()

    def add_experience(self, experience: Experience) -> None:
        self._replay_buffer.add_experience(experience)

    def get_buffer_total_reward(self) -> float:
        return self._replay_buffer.total_reward()

    def act(self, observation: np.ndarray, eval: bool = False) -> int:
        x = torch.from_numpy(observation.astype(np.float32))
        with torch.no_grad():
            action_logits, _ = self._net(x)
        next_action_dist = torch.distributions.Categorical(logits=action_logits)
        action = next_action_dist.sample()
        if eval:
            action = next_action_dist.probs.argmax(dim=-1)
        return action.item()

    def net_update(
        self, gamma: float, lambda_: float, value_loss_coef: float, policy_loss_coef: float, entropy_coef: float
    ) -> tuple[float, float, float]:
        experiences = self._replay_buffer.sample()

        observations = torch.tensor([exp.observation for exp in experiences])
        actions = torch.tensor([exp.action for exp in experiences])
        rewards = torch.tensor([exp.reward for exp in experiences]).unsqueeze(1)
        next_observations = torch.tensor([exp.next_observation for exp in experiences])
        terminateds = torch.tensor([exp.terminated for exp in experiences], dtype=torch.float32).unsqueeze(1)

        # calculate predicted V-values
        policy_action_logits, v_values = self._net(observations)
        # n-step return
        v_targets = torch.zeros_like(rewards)
        for t in reversed(range(len(experiences) - 1)):
            v_targets[t] = rewards[t] + gamma * v_targets[t + 1] * (1 - terminateds[t])
        # calculate value loss
        value_loss = nn.functional.mse_loss(v_values, v_targets)

        # calculate advantages by GAE
        with torch.no_grad():
            _, v_values_next = self._net(next_observations)
        deltas = rewards + gamma * v_values_next * (1 - terminateds) - v_values
        advantages = deltas.clone()
        for t in reversed(range(len(experiences) - 1)):
            advantages[t] = deltas[t] + gamma * lambda_ * advantages[t + 1] * (1 - terminateds[t])
        advantages = advantages / (advantages.std() + 1e-8)
        advantages = advantages.detach()

        action_dist = torch.distributions.Categorical(logits=policy_action_logits)
        action_entropy = action_dist.entropy().mean()
        action_log_probs = action_dist.log_prob(actions)
        # calculate policy loss
        policy_loss = -action_log_probs * advantages
        policy_loss = torch.mean(policy_loss)

        loss = value_loss * value_loss_coef + policy_loss * policy_loss_coef - entropy_coef * action_entropy

        # update
        self._optimizer.zero_grad()
        loss.backward()
        self._optimizer.step()
        return loss.item(), action_entropy.item(), advantages.mean().item()


class A2CTrainer:
    def __init__(self, config: A2CConfig) -> None:
        self.config = config
        self.env = gym.make(config.env_name, render_mode=config.render_mode)
        env_dim = self.env.observation_space.shape[0]  # type: ignore[index]
        action_num = self.env.action_space.n  # type: ignore[attr-defined]
        net = ActorCriticNet(env_dim=env_dim, action_num=action_num)
        optimizer = optim.Adam(net.parameters(), lr=config.learning_rate)
        self.agent = Agent(net=net, optimizer=optimizer)

        self.num_episodes = config.num_episodes
        self.gamma = config.gamma
        self.lambda_ = config.lambda_
        self.value_loss_coef = config.value_loss_coef
        self.policy_loss_coef = config.policy_loss_coef
        self.entropy_coef = config.entropy_coef
        self.solved_threshold = config.solved_threshold
        if config.log_wandb:
            wandb.init(
                # set the wandb project where this run will be logged
                project="A2C",
                name=f"[{config.env_name}]lr={config.learning_rate}",
                # track hyperparameters and run metadata
                config=asdict(config),
            )

    def train(self) -> None:
        for i, episode in enumerate(range(self.num_episodes)):
            observation, _ = self.env.reset()
            terminated, truncated = False, False
            while not (terminated or truncated):
                action = self.agent.act(observation)
                next_observation, reward, terminated, truncated, _ = self.env.step(action)
                experience = Experience(
                    observation=observation,
                    action=action,
                    reward=float(reward),
                    terminated=terminated,
                    truncated=truncated,
                    next_observation=next_observation,
                )
                self.agent.add_experience(experience)
                observation = next_observation
                if self.config.render_mode is not None:
                    self.env.render()
            loss, action_entropy, advantages_mean = self.agent.net_update(
                gamma=self.gamma,
                lambda_=self.lambda_,
                value_loss_coef=self.value_loss_coef,
                policy_loss_coef=self.policy_loss_coef,
                entropy_coef=self.entropy_coef,
            )
            total_reward = self.agent.get_buffer_total_reward()
            solved = total_reward > self.solved_threshold
            self.agent.onpolicy_reset()
            print(
                f"Episode {episode}, total_reward: {total_reward}, solved: {solved}, loss: {loss}, "
                f"action_entropy: {action_entropy}, advantages_mean: {advantages_mean}"
            )
            if self.config.log_wandb:
                wandb.log(
                    {
                        "episode": episode,
                        "loss": loss,
                        "total_reward": total_reward,
                        "action_entropy": action_entropy,
                        "advantages_mean": advantages_mean,
                    }
                )

            if i % self.config.eval_interval == 0:
                eval_reward = self.evaluate(self.config.eval_episodes)
                print(f"Episode {episode}, Eval reward: {eval_reward}")
                if self.config.log_wandb:
                    wandb.log({"eval_reward": eval_reward, "episode": episode})

    def evaluate(self, num_episodes: int) -> float:
        total_reward = 0.0
        for _ in range(num_episodes):
            observation, _ = self.env.reset()
            terminated, truncated = False, False
            while not (terminated or truncated):
                action = self.agent.act(observation, eval=True)
                next_observation, reward, terminated, truncated, _ = self.env.step(action)
                observation = next_observation
                total_reward += float(reward)
        return total_reward / num_episodes


if __name__ == "__main__":
    default_config = A2CConfig(
        env_name="CartPole-v1",
        render_mode=None,
        solved_threshold=475.0,
        num_episodes=100000,
        learning_rate=0.002,
        log_wandb=True,
    )
    trainer = A2CTrainer(default_config)
    trainer.train()
