import copy
import numpy as np
import random
import torch
import torch.nn as nn
import torch.nn.functional as F

from collections import deque
from torch.optim import Adam
from torch.utils.tensorboard.writer import SummaryWriter
from typing import NamedTuple, Optional, TypeVar

from rl_algo_impls.dqn.policy import DQNPolicy
from rl_algo_impls.shared.algorithm import Algorithm
from rl_algo_impls.shared.callbacks.callback import Callback
from rl_algo_impls.shared.schedule import linear_schedule
from rl_algo_impls.wrappers.vectorable_wrapper import VecEnv, VecEnvObs


class Transition(NamedTuple):
    obs: np.ndarray
    action: np.ndarray
    reward: float
    done: bool
    next_obs: np.ndarray


class Batch(NamedTuple):
    obs: np.ndarray
    actions: np.ndarray
    rewards: np.ndarray
    dones: np.ndarray
    next_obs: np.ndarray


class ReplayBuffer:
    def __init__(self, num_envs: int, maxlen: int) -> None:
        self.num_envs = num_envs
        self.buffer = deque(maxlen=maxlen)

    def add(
        self,
        obs: VecEnvObs,
        action: np.ndarray,
        reward: np.ndarray,
        done: np.ndarray,
        next_obs: VecEnvObs,
    ) -> None:
        assert isinstance(obs, np.ndarray)
        assert isinstance(next_obs, np.ndarray)
        for i in range(self.num_envs):
            self.buffer.append(
                Transition(obs[i], action[i], reward[i], done[i], next_obs[i])
            )

    def sample(self, batch_size: int) -> Batch:
        ts = random.sample(self.buffer, batch_size)
        return Batch(
            obs=np.array([t.obs for t in ts]),
            actions=np.array([t.action for t in ts]),
            rewards=np.array([t.reward for t in ts]),
            dones=np.array([t.done for t in ts]),
            next_obs=np.array([t.next_obs for t in ts]),
        )

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


DQNSelf = TypeVar("DQNSelf", bound="DQN")


class DQN(Algorithm):
    def __init__(
        self,
        policy: DQNPolicy,
        env: VecEnv,
        device: torch.device,
        tb_writer: SummaryWriter,
        learning_rate: float = 1e-4,
        buffer_size: int = 1_000_000,
        learning_starts: int = 50_000,
        batch_size: int = 32,
        tau: float = 1.0,
        gamma: float = 0.99,
        train_freq: int = 4,
        gradient_steps: int = 1,
        target_update_interval: int = 10_000,
        exploration_fraction: float = 0.1,
        exploration_initial_eps: float = 1.0,
        exploration_final_eps: float = 0.05,
        max_grad_norm: float = 10.0,
    ) -> None:
        super().__init__(policy, env, device, tb_writer)
        self.policy = policy

        self.optimizer = Adam(self.policy.q_net.parameters(), lr=learning_rate)

        self.target_q_net = copy.deepcopy(self.policy.q_net).to(self.device)
        self.target_q_net.train(False)
        self.tau = tau
        self.target_update_interval = target_update_interval

        self.replay_buffer = ReplayBuffer(self.env.num_envs, buffer_size)
        self.batch_size = batch_size

        self.learning_starts = learning_starts
        self.train_freq = train_freq
        self.gradient_steps = gradient_steps

        self.gamma = gamma
        self.exploration_eps_schedule = linear_schedule(
            exploration_initial_eps,
            exploration_final_eps,
            end_fraction=exploration_fraction,
        )

        self.max_grad_norm = max_grad_norm

    def learn(
        self: DQNSelf, total_timesteps: int, callback: Optional[Callback] = None
    ) -> DQNSelf:
        self.policy.train(True)
        obs = self.env.reset()
        obs = self._collect_rollout(self.learning_starts, obs, 1)
        learning_steps = total_timesteps - self.learning_starts
        timesteps_elapsed = 0
        steps_since_target_update = 0
        while timesteps_elapsed < learning_steps:
            progress = timesteps_elapsed / learning_steps
            eps = self.exploration_eps_schedule(progress)
            obs = self._collect_rollout(self.train_freq, obs, eps)
            rollout_steps = self.train_freq
            timesteps_elapsed += rollout_steps
            for _ in range(
                self.gradient_steps if self.gradient_steps > 0 else self.train_freq
            ):
                self.train()
            steps_since_target_update += rollout_steps
            if steps_since_target_update >= self.target_update_interval:
                self._update_target()
                steps_since_target_update = 0
            if callback:
                callback.on_step(timesteps_elapsed=rollout_steps)
        return self

    def train(self) -> None:
        if len(self.replay_buffer) < self.batch_size:
            return
        o, a, r, d, next_o = self.replay_buffer.sample(self.batch_size)
        o = torch.as_tensor(o, device=self.device)
        a = torch.as_tensor(a, device=self.device).unsqueeze(1)
        r = torch.as_tensor(r, dtype=torch.float32, device=self.device)
        d = torch.as_tensor(d, dtype=torch.long, device=self.device)
        next_o = torch.as_tensor(next_o, device=self.device)

        with torch.no_grad():
            target = r + (1 - d) * self.gamma * self.target_q_net(next_o).max(1).values
        current = self.policy.q_net(o).gather(dim=1, index=a).squeeze(1)
        loss = F.smooth_l1_loss(current, target)

        self.optimizer.zero_grad()
        loss.backward()
        if self.max_grad_norm:
            nn.utils.clip_grad_norm_(self.policy.q_net.parameters(), self.max_grad_norm)
        self.optimizer.step()

    def _collect_rollout(self, timesteps: int, obs: VecEnvObs, eps: float) -> VecEnvObs:
        for _ in range(0, timesteps, self.env.num_envs):
            action = self.policy.act(obs, eps, deterministic=False)
            next_obs, reward, done, _ = self.env.step(action)
            self.replay_buffer.add(obs, action, reward, done, next_obs)
            obs = next_obs
        return obs

    def _update_target(self) -> None:
        for target_param, param in zip(
            self.target_q_net.parameters(), self.policy.q_net.parameters()
        ):
            target_param.data.copy_(
                self.tau * param.data + (1 - self.tau) * target_param.data
            )