from abc import ABC, abstractmethod
import os
import numpy as np
import torch
from torch.utils.tensorboard import SummaryWriter

from fqf_iqn_qrdqn.memory import LazyMultiStepMemory, \
    LazyPrioritizedMultiStepMemory
from fqf_iqn_qrdqn.utils import RunningMeanStats, LinearAnneaer


class BaseAgent(ABC):

    def __init__(self, env, test_env, log_dir, num_steps=5 * (10 ** 7),
                 batch_size=32, memory_size=10 ** 6, gamma=0.99, multi_step=1,
                 update_interval=4, target_update_interval=10000,
                 start_steps=50000, epsilon_train=0.01, epsilon_eval=0.001,
                 epsilon_decay_steps=250000, double_q_learning=False,
                 dueling_net=False, noisy_net=False, use_per=False,
                 log_interval=100, eval_interval=250000, num_eval_steps=125000,
                 max_episode_steps=27000, grad_cliping=5.0, cuda=True, seed=0):

        self.env = env
        self.test_env = test_env

        torch.manual_seed(seed)
        np.random.seed(seed)
        self.env.seed(seed)
        self.test_env.seed(2 ** 31 - 1 - seed)
        # torch.backends.cudnn.deterministic = True  # It harms a performance.
        # torch.backends.cudnn.benchmark = False  # It harms a performance.

        self.device = torch.device(
            "cuda" if cuda and torch.cuda.is_available() else "cpu")

        self.online_net = None
        self.target_net = None

        # Replay memory which is memory-efficient to store stacked frames.
        if use_per:
            beta_steps = (num_steps - start_steps) / update_interval
            self.memory = LazyPrioritizedMultiStepMemory(
                memory_size, self.env.observation_space.shape,
                self.device, gamma, multi_step, beta_steps=beta_steps)
        else:
            self.memory = LazyMultiStepMemory(
                memory_size, self.env.observation_space.shape,
                self.device, gamma, multi_step)

        self.log_dir = log_dir
        self.model_dir = os.path.join(log_dir, 'model')
        self.summary_dir = os.path.join(log_dir, 'summary')
        if not os.path.exists(self.model_dir):
            os.makedirs(self.model_dir)
        if not os.path.exists(self.summary_dir):
            os.makedirs(self.summary_dir)

        self.writer = SummaryWriter(log_dir=self.summary_dir)
        self.train_return = RunningMeanStats(log_interval)

        self.steps = 0
        self.learning_steps = 0
        self.episodes = 0
        self.best_eval_score = -np.inf
        self.num_actions = self.env.action_space.n
        self.num_steps = num_steps
        self.batch_size = batch_size

        self.double_q_learning = double_q_learning
        self.dueling_net = dueling_net
        self.noisy_net = noisy_net
        self.use_per = use_per

        self.log_interval = log_interval
        self.eval_interval = eval_interval
        self.num_eval_steps = num_eval_steps
        self.gamma_n = gamma ** multi_step
        self.start_steps = start_steps
        self.epsilon_train = LinearAnneaer(
            1.0, epsilon_train, epsilon_decay_steps)
        self.epsilon_eval = epsilon_eval
        self.update_interval = update_interval
        self.target_update_interval = target_update_interval
        self.max_episode_steps = max_episode_steps
        self.grad_cliping = grad_cliping

    def run(self):
        while True:
            self.train_episode()
            if self.steps > self.num_steps:
                break

    def is_update(self):
        return self.steps % self.update_interval == 0 \
               and self.steps >= self.start_steps

    def is_random(self, eval=False):
        # Use e-greedy for evaluation.
        if self.steps < self.start_steps:
            return True
        if eval:
            return np.random.rand() < self.epsilon_eval
        if self.noisy_net:
            return False
        return np.random.rand() < self.epsilon_train.get()

    def update_target(self):
        self.target_net.load_state_dict(
            self.online_net.state_dict())

    def explore(self):
        # Act with randomness.
        action = self.env.action_space.sample()
        return action

    def exploit(self, state):
        # Act without randomness.
        state = torch.ByteTensor(
            state).unsqueeze(0).to(self.device).float() / 255.
        with torch.no_grad():
            action = self.online_net.calculate_q(states=state).argmax().item()
        return action

    @abstractmethod
    def learn(self):
        pass

    def save_models(self, save_dir):
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        torch.save(
            self.online_net.state_dict(),
            os.path.join(save_dir, 'online_net.pth'))
        torch.save(
            self.target_net.state_dict(),
            os.path.join(save_dir, 'target_net.pth'))

    def load_models(self, save_dir):
        self.online_net.load_state_dict(torch.load(
            os.path.join(save_dir, 'online_net.pth')))
        self.target_net.load_state_dict(torch.load(
            os.path.join(save_dir, 'target_net.pth')))

    def train_episode(self):
        self.online_net.train()
        self.target_net.train()

        self.episodes += 1
        episode_return = 0.
        episode_steps = 0

        done = False
        state = self.env.reset()

        while (not done) and episode_steps <= self.max_episode_steps:
            # NOTE: Noises can be sampled only after self.learn(). However, I
            # sample noises before every action, which seems to lead better
            # performances.
            self.online_net.sample_noise()

            if self.is_random(eval=False):
                action = self.explore()
            else:
                action = self.exploit(state)

            next_state, reward, done, _ = self.env.step(action)

            # To calculate efficiently, I just set priority=max_priority here.
            self.memory.append(state, action, reward, next_state, done)

            self.steps += 1
            episode_steps += 1
            episode_return += reward
            state = next_state

            self.train_step_interval()

        # We log running mean of stats.
        self.train_return.append(episode_return)

        # We log evaluation results along with training frames = 4 * steps.
        if self.episodes % self.log_interval == 0:
            self.writer.add_scalar(
                'return/train', self.train_return.get(), 4 * self.steps)

        print(f'Episode: {self.episodes:<4}  '
              f'episode steps: {episode_steps:<4}  '
              f'return: {episode_return:<5.1f}')

    def train_step_interval(self):
        self.epsilon_train.step()

        if self.steps % self.target_update_interval == 0:
            self.update_target()

        if self.is_update():
            self.learn()

        if self.steps % self.eval_interval == 0:
            self.evaluate()
            self.save_models(os.path.join(self.model_dir, 'final'))
            self.online_net.train()

    def evaluate(self):
        self.online_net.eval()
        num_episodes = 0
        num_steps = 0
        total_return = 0.0

        while True:
            state = self.test_env.reset()
            episode_steps = 0
            episode_return = 0.0
            done = False
            while (not done) and episode_steps <= self.max_episode_steps:
                if self.is_random(eval=True):
                    action = self.explore()
                else:
                    action = self.exploit(state)

                next_state, reward, done, _ = self.test_env.step(action)
                num_steps += 1
                episode_steps += 1
                episode_return += reward
                state = next_state

            num_episodes += 1
            total_return += episode_return

            if num_steps > self.num_eval_steps:
                break

        mean_return = total_return / num_episodes

        if mean_return > self.best_eval_score:
            self.best_eval_score = mean_return
            self.save_models(os.path.join(self.model_dir, 'best'))

        # We log evaluation results along with training frames = 4 * steps.
        self.writer.add_scalar(
            'return/test', mean_return, 4 * self.steps)
        print('-' * 60)
        print(f'Num steps: {self.steps:<5}  '
              f'return: {mean_return:<5.1f}')
        print('-' * 60)

    def __del__(self):
        self.env.close()
        self.test_env.close()
        self.writer.close()

[((26, 8, 26, 31), 'torch.manual_seed', 'torch.manual_seed', ({(26, 26, 26, 30): 'seed'}, {}), '(seed)', False, 'import torch\n'), ((27, 8, 27, 28), 'numpy.random.seed', 'np.random.seed', ({(27, 23, 27, 27): 'seed'}, {}), '(seed)', True, 'import numpy as np\n'), ((51, 25, 51, 55), 'os.path.join', 'os.path.join', ({(51, 38, 51, 45): 'log_dir', (51, 47, 51, 54): '"""model"""'}, {}), "(log_dir, 'model')", False, 'import os\n'), ((52, 27, 52, 59), 'os.path.join', 'os.path.join', ({(52, 40, 52, 47): 'log_dir', (52, 49, 52, 58): '"""summary"""'}, {}), "(log_dir, 'summary')", False, 'import os\n'), ((58, 22, 58, 61), 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', (), '', False, 'from torch.utils.tensorboard import SummaryWriter\n'), ((59, 28, 59, 58), 'fqf_iqn_qrdqn.utils.RunningMeanStats', 'RunningMeanStats', ({(59, 45, 59, 57): 'log_interval'}, {}), '(log_interval)', False, 'from fqf_iqn_qrdqn.utils import RunningMeanStats, LinearAnneaer\n'), ((79, 29, 80, 52), 'fqf_iqn_qrdqn.utils.LinearAnneaer', 'LinearAnneaer', ({(80, 12, 80, 15): '1.0', (80, 17, 80, 30): 'epsilon_train', (80, 32, 80, 51): 'epsilon_decay_steps'}, {}), '(1.0, epsilon_train, epsilon_decay_steps)', False, 'from fqf_iqn_qrdqn.utils import RunningMeanStats, LinearAnneaer\n'), ((42, 26, 44, 70), 'fqf_iqn_qrdqn.memory.LazyPrioritizedMultiStepMemory', 'LazyPrioritizedMultiStepMemory', (), '', False, 'from fqf_iqn_qrdqn.memory import LazyMultiStepMemory, LazyPrioritizedMultiStepMemory\n'), ((46, 26, 48, 47), 'fqf_iqn_qrdqn.memory.LazyMultiStepMemory', 'LazyMultiStepMemory', ({(47, 16, 47, 27): 'memory_size', (47, 29, 47, 61): 'self.env.observation_space.shape', (48, 16, 48, 27): 'self.device', (48, 29, 48, 34): 'gamma', (48, 36, 48, 46): 'multi_step'}, {}), '(memory_size, self.env.observation_space.shape, self.\n    device, gamma, multi_step)', False, 'from fqf_iqn_qrdqn.memory import LazyMultiStepMemory, LazyPrioritizedMultiStepMemory\n'), ((53, 15, 53, 45), 'os.path.exists', 'os.path.exists', ({(53, 30, 53, 44): 'self.model_dir'}, {}), '(self.model_dir)', False, 'import os\n'), ((54, 12, 54, 39), 'os.makedirs', 'os.makedirs', ({(54, 24, 54, 38): 'self.model_dir'}, {}), '(self.model_dir)', False, 'import os\n'), ((55, 15, 55, 47), 'os.path.exists', 'os.path.exists', ({(55, 30, 55, 46): 'self.summary_dir'}, {}), '(self.summary_dir)', False, 'import os\n'), ((56, 12, 56, 41), 'os.makedirs', 'os.makedirs', ({(56, 24, 56, 40): 'self.summary_dir'}, {}), '(self.summary_dir)', False, 'import os\n'), ((105, 15, 105, 31), 'numpy.random.rand', 'np.random.rand', ({}, {}), '()', True, 'import numpy as np\n'), ((120, 13, 120, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((129, 15, 129, 39), 'os.path.exists', 'os.path.exists', ({(129, 30, 129, 38): 'save_dir'}, {}), '(save_dir)', False, 'import os\n'), ((130, 12, 130, 33), 'os.makedirs', 'os.makedirs', ({(130, 24, 130, 32): 'save_dir'}, {}), '(save_dir)', False, 'import os\n'), ((133, 12, 133, 52), 'os.path.join', 'os.path.join', ({(133, 25, 133, 33): 'save_dir', (133, 35, 133, 51): '"""online_net.pth"""'}, {}), "(save_dir, 'online_net.pth')", False, 'import os\n'), ((136, 12, 136, 52), 'os.path.join', 'os.path.join', ({(136, 25, 136, 33): 'save_dir', (136, 35, 136, 51): '"""target_net.pth"""'}, {}), "(save_dir, 'target_net.pth')", False, 'import os\n'), ((102, 19, 102, 35), 'numpy.random.rand', 'np.random.rand', ({}, {}), '()', True, 'import numpy as np\n'), ((140, 12, 140, 52), 'os.path.join', 'os.path.join', ({(140, 25, 140, 33): 'save_dir', (140, 35, 140, 51): '"""online_net.pth"""'}, {}), "(save_dir, 'online_net.pth')", False, 'import os\n'), ((142, 12, 142, 52), 'os.path.join', 'os.path.join', ({(142, 25, 142, 33): 'save_dir', (142, 35, 142, 51): '"""target_net.pth"""'}, {}), "(save_dir, 'target_net.pth')", False, 'import os\n'), ((201, 29, 201, 66), 'os.path.join', 'os.path.join', ({(201, 42, 201, 56): 'self.model_dir', (201, 58, 201, 65): '"""final"""'}, {}), "(self.model_dir, 'final')", False, 'import os\n'), ((237, 29, 237, 65), 'os.path.join', 'os.path.join', ({(237, 42, 237, 56): 'self.model_dir', (237, 58, 237, 64): '"""best"""'}, {}), "(self.model_dir, 'best')", False, 'import os\n'), ((34, 31, 34, 56), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((118, 16, 119, 18), 'torch.ByteTensor', 'torch.ByteTensor', ({(119, 12, 119, 17): 'state'}, {}), '(state)', False, 'import torch\n')]