#!/usr/bin/env python3

import random
import timeit
import numpy as np
import os.path as osp
from itertools import count
from datetime import timedelta
import matplotlib
import matplotlib.pyplot as plt
import collections
from collections import namedtuple
from tensorboardX import SummaryWriter

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as T

import pybullet as p
from pybullet_envs.bullet.kuka_diverse_object_gym_env \
    import KukaDiverseObjectEnv

# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display
plt.ion()  # turn interactive mode on

# if gpu is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward'))


class ReplayMemory(object):
    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        """Saves a transition."""
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = Transition(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)


class DQN(nn.Module):
    def __init__(self, h, w, outputs):
        super(DQN, self).__init__()
        self.conv1 = nn.Conv2d(STACK_SIZE, 32, kernel_size=8, stride=4)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)

        def conv2d_size_out(size, kernel_size=5, stride=2):
            return (size - (kernel_size - 1) - 1) // stride + 1
        convw = conv2d_size_out(
            conv2d_size_out(conv2d_size_out(w, 8, 4), 4, 2), 3, 1)
        convh = conv2d_size_out(
            conv2d_size_out(conv2d_size_out(h, 8, 4), 4, 2), 3, 1)
        linear_input_size = convw * convh * 64
        self.linear = nn.Linear(linear_input_size, 512)
        self.head = nn.Linear(512, outputs)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.conv3(x))
        x = F.relu(self.linear(x.view(x.size(0), -1)))
        return self.head(x)


def get_screen():
    global stacked_screens
    # returned screen requested by gym is 400x600x3, but is sometimes larger
    # such as 800x1200x3. Transpose it into torch order (CHW).
    screen = env._get_observation().transpose((2, 0, 1))
    # convert to float, rescale, convert to torch tensor
    # (this doesn't require a copy)

    screen = np.ascontiguousarray(screen, dtype=np.float32) / 255
    screen = torch.from_numpy(screen)
    # resize, and add a batch dimension (BCHW)
    return preprocess(screen).unsqueeze(0).to(device)


def select_action(state, i_episode):
    global steps_done
    global eps_threshold
    sample = random.random()
    eps_threshold = max(EPS_END, EPS_START - i_episode / EPS_DECAY_LAST_FRAME)
    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1) will return largest column value of each row.
            # second column on max result is index of where max element was
            # found, so we pick action with the larger expected reward.
            return policy_net(state).max(1)[1].view(1, 1)
    else:
        return torch.tensor(
            [[random.randrange(n_actions)]], device=device, dtype=torch.long)


def optimize_model():
    if len(memory) < BATCH_SIZE:
        return
    transitions = memory.sample(BATCH_SIZE)
    # transpose the batch (see https://stackoverflow.com/a/19343/3343043 for
    # detailed explanation). This converts batch-array of Transitions
    # to Transition of batch-arrays.
    batch = Transition(*zip(*transitions))

    # compute a mask of non-final states and concatenate the batch elements
    # (a final state would've been the one after which simulation ended)
    non_final_mask = torch.tensor(
        tuple(map(lambda s: s is not None, batch.next_state)),
        device=device, dtype=torch.bool)
    non_final_next_states = torch.cat(
        [s for s in batch.next_state if s is not None])
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)

    # compute Q(s_t, a) - the model computes Q(s_t), then we select the
    # columns of actions taken. These are the actions which would've been taken
    # for each batch state according to policy_net
    state_action_values = policy_net(state_batch).gather(1, action_batch)

    # compute V(s_{t+1}) for all next states.
    # expected values of actions for non_final_next_states are computed based
    # on the "older" target_net; selecting their best reward with max(1)[0].
    # this is merged based on the mask,
    # such that we'll have either the expected
    # state value or 0 in case the state was final.
    next_state_values = torch.zeros(BATCH_SIZE, device=device)
    next_state_values[non_final_mask] = \
        target_net(non_final_next_states).max(1)[0].detach()
    # compute the expected Q values
    expected_state_action_values = (next_state_values * GAMMA) + reward_batch

    # compute Huber loss
    loss = F.smooth_l1_loss(
        state_action_values, expected_state_action_values.unsqueeze(1))

    # optimize the model
    optimizer.zero_grad()
    loss.backward()
    for param in policy_net.parameters():
        param.grad.data.clamp_(-1, 1)
    optimizer.step()


def train():
    writer = SummaryWriter()
    total_rewards = []
    ten_rewards = 0
    best_mean_reward = None
    start_time = timeit.default_timer()
    for i_episode in range(num_episodes):
        # initialize the environment and state
        env.reset()
        state = get_screen()
        stacked_states = collections.deque(
            STACK_SIZE*[state], maxlen=STACK_SIZE)
        for t in count():
            stacked_states_t = torch.cat(tuple(stacked_states), dim=1)
            # select and perform an action
            action = select_action(stacked_states_t, i_episode)
            _, reward, done, _ = env.step(action.item())
            reward = torch.tensor([reward], device=device)

            # observe new state
            next_state = get_screen()
            if not done:
                next_stacked_states = stacked_states
                next_stacked_states.append(next_state)
                next_stacked_states_t = torch.cat(
                    tuple(next_stacked_states), dim=1)
            else:
                next_stacked_states_t = None

            # store the transition in memory
            memory.push(stacked_states_t, action,
                        next_stacked_states_t,
                        reward)

            # move to the next state
            stacked_states = next_stacked_states

            # perform one step of the optimization (on the target network)
            optimize_model()
            if done:
                reward = reward.cpu().numpy().item()
                ten_rewards += reward
                total_rewards.append(reward)
                mean_reward = np.mean(total_rewards[-100:])*100
                writer.add_scalar("epsilon", eps_threshold, i_episode)
                if (best_mean_reward is None or
                    best_mean_reward < mean_reward) and \
                   i_episode > 100:
                    # for saving the model and possibly resuming training
                    torch.save({
                        'policy_net_state_dict':
                            policy_net.state_dict(),
                        'target_net_state_dict':
                            target_net.state_dict(),
                        'optimizer_policy_net_state_dict':
                            optimizer.state_dict()
                        }, save_model_path)
                    if best_mean_reward is not None:
                        print("Best mean reward updated %.1f -> %.1f, \
                               model saved"
                              % (best_mean_reward, mean_reward))
                    best_mean_reward = mean_reward
                break

        if i_episode % 10 == 0:
            writer.add_scalar('ten episodes average rewards',
                              ten_rewards/10.0, i_episode)
            ten_rewards = 0

        # update the target network, copying all weights and biases in DQN
        if i_episode % TARGET_UPDATE == 0:
            target_net.load_state_dict(policy_net.state_dict())
        if i_episode >= 200 and mean_reward > 50:
            print('Environment solved in {:d} \
                   episodes!\tAverage Score: {:.2f}'.format(
                       i_episode+1, mean_reward))
            break

    print('Average Score: {:.2f}'.format(mean_reward))
    elapsed = timeit.default_timer() - start_time
    print("Elapsed time: {}".format(timedelta(seconds=elapsed)))
    writer.close()
    env.close()


if __name__ == '__main__':
    STACK_SIZE = 5
    preprocess = T.Compose(
        [T.ToPILImage(),
         T.Grayscale(num_output_channels=1),
         T.Resize(40, interpolation=T.InterpolationMode.BICUBIC),
         T.ToTensor()])

    isRendersTrain = True
    env = KukaDiverseObjectEnv(
        renders=isRendersTrain,
        isDiscrete=True,
        removeHeightHack=False,
        maxSteps=20)
    env.cid = p.connect(p.DIRECT)
    env.reset()

    BATCH_SIZE = 32
    GAMMA = 0.99
    EPS_START = 0.9
    EPS_END = 0.1
    EPS_DECAY = 200
    EPS_DECAY_LAST_FRAME = 10**4
    TARGET_UPDATE = 1000
    LEARNING_RATE = 1e-4

    # get screen size so that we can initialize layers correctly based on shape
    # returned from pybullet (48, 48, 3).
    init_screen = get_screen()
    _, _, screen_height, screen_width = init_screen.shape

    # get number of actions from gym action space
    n_actions = env.action_space.n

    policy_net = DQN(screen_height, screen_width, n_actions).to(device)
    target_net = DQN(screen_height, screen_width, n_actions).to(device)
    target_net.load_state_dict(policy_net.state_dict())
    target_net.eval()

    optimizer = optim.Adam(policy_net.parameters(), lr=LEARNING_RATE)
    memory = ReplayMemory(10000)
    eps_threshold = 0

    save_model_path = osp.join(
        osp.dirname(__file__),
        '..',
        'data',
        'result',
        'robot',
        'policy_dqn.pt')
    num_episodes = 3
    train()
