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 SarsaConfig:
    """Configuration for SARSA algorithm."""

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

    # Training config
    gamma: float = 0.999
    max_training_steps: int = 500000
    """The maximum number of environment steps to train for."""
    learning_rate: float = 2.5e-4
    """The learning rate for the optimizer."""

    # Logging config
    log_wandb: bool = False
    """Whether to log the training process to Weights and Biases."""


class PolicyNet(nn.Module):
    def __init__(
        self,
        env_dim: int,
        action_num: int,
    ) -> None:
        super().__init__()
        self.env_dim = env_dim
        self.action_num = action_num

        layers = [
            nn.Linear(self.env_dim, 128),
            nn.ReLU(),
            nn.Linear(128, self.action_num),
        ]
        self.model = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)


@dataclass
class Experience:
    terminated: bool
    truncated: bool
    observation: Any  # S
    action: Any  # A
    reward: float  # R
    next_observation: Any = None  # S'
    next_action: Any = None  # A'


@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, with_next_sa: bool = True) -> list[Experience]:
        if with_next_sa is False:
            return self.buffer
        else:
            res = []
            for i in range(len(self.buffer) - 1):
                experience = self.buffer[i]
                next_experience = self.buffer[i + 1]
                res.append(
                    Experience(
                        observation=experience.observation,
                        action=experience.action,
                        reward=experience.reward,
                        next_observation=next_experience.observation,
                        next_action=next_experience.action,
                        terminated=next_experience.terminated,
                        truncated=next_experience.truncated,
                    )
                )
            return res


class Agent:
    def __init__(self, policy_net: PolicyNet, optimizer: torch.optim.Optimizer) -> None:
        self._policy_net = policy_net
        self._optimizer = optimizer
        self._replay_buffer = ReplayBuffer()
        self._action_num = policy_net.action_num

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

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

    def act(self, observation: np.floating, epsilon: float) -> int:
        if np.random.rand() < epsilon:
            action = np.random.randint(self._action_num)
            return action
        x = torch.from_numpy(observation.astype(np.float32))
        with torch.no_grad():
            logits = self._policy_net(x)
        action = int(torch.argmax(logits).item())
        return action

    def policy_update(self, gamma: float) -> float:
        experiences = self._replay_buffer.sample()

        observations = torch.tensor([experience.observation for experience in experiences])
        actions = torch.tensor([experience.action for experience in experiences], dtype=torch.float32)
        next_observations = torch.tensor([experience.next_observation for experience in experiences])
        next_actions = torch.tensor([experience.next_action for experience in experiences])
        rewards = torch.tensor([experience.reward for experience in experiences])
        terminated = torch.tensor(
            [experience.terminated for experience in experiences],
            dtype=torch.float32,
        )

        # q preds
        action_q_preds = self._policy_net(observations).gather(1, actions.long().unsqueeze(1)).squeeze(1)
        with torch.no_grad():
            next_action_q_preds = self._policy_net(next_observations).gather(1, next_actions.unsqueeze(1)).squeeze(1)
            q_targets = rewards + gamma * (1 - terminated) * next_action_q_preds
        loss = torch.nn.functional.mse_loss(action_q_preds, q_targets)
        # update
        self._optimizer.zero_grad()
        loss.backward()
        # clip grad
        torch.nn.utils.clip_grad_norm_(self._policy_net.parameters(), max_norm=1.0)
        self._optimizer.step()
        return loss.item()


class SarsaTrainer:
    def __init__(self, config: SarsaConfig) -> None:
        self.config = config
        self.env = self._make_env()
        if isinstance(self.env.action_space, gym.spaces.Discrete) is False:
            raise ValueError("Only discrete action space is supported.")
        env_dim = self.env.observation_space.shape[0]
        action_num = self.env.action_space.n
        policy_net = PolicyNet(env_dim=env_dim, action_num=action_num)
        optimizer = optim.Adam(policy_net.parameters(), lr=config.learning_rate)
        self.agent = Agent(policy_net=policy_net, optimizer=optimizer)

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

    def _make_env(self):
        env = gym.make(self.config.env_name, render_mode=self.config.render_mode)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env = gym.wrappers.Autoreset(env)
        return env

    def train(self) -> None:
        epsilon = 1.0
        global_step = 0

        observation, _ = self.env.reset()
        while global_step < self.config.max_training_steps:
            global_step += 1
            epsilon = max(0.05, epsilon * 0.9999)

            action = self.agent.act(observation, epsilon)
            next_observation, reward, terminated, truncated, info = self.env.step(action)
            experience = Experience(
                observation=observation,
                action=action,
                reward=float(reward),
                next_observation=next_observation,
                terminated=terminated,
                truncated=truncated,
            )
            self.agent.add_experience(experience)
            observation = next_observation
            if self.env.render_mode is not None:
                self.env.render()

            if terminated or truncated:
                if info and "episode" in info:
                    episode_reward = info["episode"]["r"]
                    loss = self.agent.policy_update(gamma=self.gamma)
                    self.agent.onpolicy_reset()
                    print(
                        f"global_step={global_step}, epsilon={epsilon}, episodic_return={episode_reward}, loss={loss}"
                    )
                    if self.config.log_wandb:
                        wandb.log(
                            {
                                "global_step": global_step,
                                "episode_reward": episode_reward,
                                "loss": loss,
                            }
                        )


if __name__ == "__main__":
    default_config = SarsaConfig(
        env_name="CartPole-v1",
        render_mode=None,
        solved_threshold=475.0,
        max_training_steps=2000000,
        learning_rate=0.01,
        log_wandb=True,
    )
    trainer = SarsaTrainer(default_config)
    trainer.train()
