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()

[((966, 989), 'torch.manual_seed', 'torch.manual_seed', (['seed'], {}), '(seed)\n', (983, 989), False, 'import torch\n'), ((998, 1018), 'numpy.random.seed', 'np.random.seed', (['seed'], {}), '(seed)\n', (1012, 1018), True, 'import numpy as np\n'), ((2004, 2034), 'os.path.join', 'os.path.join', (['log_dir', '"""model"""'], {}), "(log_dir, 'model')\n", (2016, 2034), False, 'import os\n'), ((2062, 2094), 'os.path.join', 'os.path.join', (['log_dir', '"""summary"""'], {}), "(log_dir, 'summary')\n", (2074, 2094), False, 'import os\n'), ((2296, 2335), 'torch.utils.tensorboard.SummaryWriter', 'SummaryWriter', ([], {'log_dir': 'self.summary_dir'}), '(log_dir=self.summary_dir)\n', (2309, 2335), False, 'from torch.utils.tensorboard import SummaryWriter\n'), ((2364, 2394), 'fqf_iqn_qrdqn.utils.RunningMeanStats', 'RunningMeanStats', (['log_interval'], {}), '(log_interval)\n', (2380, 2394), False, 'from fqf_iqn_qrdqn.utils import RunningMeanStats, LinearAnneaer\n'), ((3037, 3091), 'fqf_iqn_qrdqn.utils.LinearAnneaer', 'LinearAnneaer', (['(1.0)', 'epsilon_train', 'epsilon_decay_steps'], {}), '(1.0, epsilon_train, epsilon_decay_steps)\n', (3050, 3091), False, 'from fqf_iqn_qrdqn.utils import RunningMeanStats, LinearAnneaer\n'), ((1609, 1746), 'fqf_iqn_qrdqn.memory.LazyPrioritizedMultiStepMemory', 'LazyPrioritizedMultiStepMemory', (['memory_size', 'self.env.observation_space.shape', 'self.device', 'gamma', 'multi_step'], {'beta_steps': 'beta_steps'}), '(memory_size, self.env.observation_space.\n    shape, self.device, gamma, multi_step, beta_steps=beta_steps)\n', (1639, 1746), False, 'from fqf_iqn_qrdqn.memory import LazyMultiStepMemory, LazyPrioritizedMultiStepMemory\n'), ((1815, 1918), 'fqf_iqn_qrdqn.memory.LazyMultiStepMemory', 'LazyMultiStepMemory', (['memory_size', 'self.env.observation_space.shape', 'self.device', 'gamma', 'multi_step'], {}), '(memory_size, self.env.observation_space.shape, self.\n    device, gamma, multi_step)\n', (1834, 1918), False, 'from fqf_iqn_qrdqn.memory import LazyMultiStepMemory, LazyPrioritizedMultiStepMemory\n'), ((2110, 2140), 'os.path.exists', 'os.path.exists', (['self.model_dir'], {}), '(self.model_dir)\n', (2124, 2140), False, 'import os\n'), ((2154, 2181), 'os.makedirs', 'os.makedirs', (['self.model_dir'], {}), '(self.model_dir)\n', (2165, 2181), False, 'import os\n'), ((2197, 2229), 'os.path.exists', 'os.path.exists', (['self.summary_dir'], {}), '(self.summary_dir)\n', (2211, 2229), False, 'import os\n'), ((2243, 2272), 'os.makedirs', 'os.makedirs', (['self.summary_dir'], {}), '(self.summary_dir)\n', (2254, 2272), False, 'import os\n'), ((3900, 3916), 'numpy.random.rand', 'np.random.rand', ([], {}), '()\n', (3914, 3916), True, 'import numpy as np\n'), ((4357, 4372), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (4370, 4372), False, 'import torch\n'), ((4583, 4607), 'os.path.exists', 'os.path.exists', (['save_dir'], {}), '(save_dir)\n', (4597, 4607), False, 'import os\n'), ((4621, 4642), 'os.makedirs', 'os.makedirs', (['save_dir'], {}), '(save_dir)\n', (4632, 4642), False, 'import os\n'), ((4717, 4757), 'os.path.join', 'os.path.join', (['save_dir', '"""online_net.pth"""'], {}), "(save_dir, 'online_net.pth')\n", (4729, 4757), False, 'import os\n'), ((4833, 4873), 'os.path.join', 'os.path.join', (['save_dir', '"""target_net.pth"""'], {}), "(save_dir, 'target_net.pth')\n", (4845, 4873), False, 'import os\n'), ((3796, 3812), 'numpy.random.rand', 'np.random.rand', ([], {}), '()\n', (3810, 3812), True, 'import numpy as np\n'), ((4977, 5017), 'os.path.join', 'os.path.join', (['save_dir', '"""online_net.pth"""'], {}), "(save_dir, 'online_net.pth')\n", (4989, 5017), False, 'import os\n'), ((5084, 5124), 'os.path.join', 'os.path.join', (['save_dir', '"""target_net.pth"""'], {}), "(save_dir, 'target_net.pth')\n", (5096, 5124), False, 'import os\n'), ((6987, 7024), 'os.path.join', 'os.path.join', (['self.model_dir', '"""final"""'], {}), "(self.model_dir, 'final')\n", (6999, 7024), False, 'import os\n'), ((8114, 8150), 'os.path.join', 'os.path.join', (['self.model_dir', '"""best"""'], {}), "(self.model_dir, 'best')\n", (8126, 8150), False, 'import os\n'), ((1317, 1342), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1340, 1342), False, 'import torch\n'), ((4263, 4286), 'torch.ByteTensor', 'torch.ByteTensor', (['state'], {}), '(state)\n', (4279, 4286), False, 'import torch\n')]