import gym
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple
from itertools import count
from PIL import Image

import torch
import torch.optim as optim
import torchvision.transforms as T

from model import ReplayMemory, DQN
from env_adapter import get_screen
from agent import Agent

TARGET_UPDATE = 10
LR = 0.0001
CMD_LINE_UPDATE_FREQ = 50

env = gym.make('CartPole-v0').unwrapped
env.reset()

# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display

plt.ion()

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

##################################
## Set up the Agent
##################################

# Get screen size so that we can initialize layers correctly based on shape
# returned from AI gym. Typical dimensions at this point are close to 3x40x90
# which is the result of a clamped and down-scaled render buffer in get_screen()
init_screen = get_screen(env, device)
_, _, 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.RMSprop(policy_net.parameters(), lr=LR)
memory = ReplayMemory(10000)

agent = Agent(policy_net, target_net, n_actions, optimizer, memory, device)


##################################
## Main iteration loop
##################################

num_episodes = 5000
for i_episode in range(num_episodes):
    # Initialize the environment and state
    env.reset()
    last_screen = get_screen(env, device)
    current_screen = get_screen(env, device)
    state = current_screen - last_screen
    for t in count():
        # Select and perform an action
        action = agent.select_action(state)
        _, reward, done, _ = env.step(action.item())
        reward = torch.tensor([reward], device=device)

        # Observe new state
        last_screen = current_screen
        current_screen = get_screen(env, device)
        if not done:
            next_state = current_screen - last_screen
        else:
            next_state = None

        # Store the transition in memory
        agent.replay_mem.push(state, action, next_state, reward)

        # Move to the next state
        state = next_state

        # Perform one step of the optimization (on the target network)
        agent.optimize_model()

        if done:
            agent.episode_durations.append(t + 1)
            agent.plot_durations()
            if i_episode % CMD_LINE_UPDATE_FREQ == 0:
                print("episode %d, duration: %d frames" % (i_episode, t+1),end="")
                if i_episode >= CMD_LINE_UPDATE_FREQ:
                    print(", (avg: %.2f frames)" % np.mean(
                        agent.episode_durations[-CMD_LINE_UPDATE_FREQ:]) )
                else:
                    print("")

            break
    # Update the target network, copying all weights and biases in DQN
    if i_episode % TARGET_UPDATE == 0:
        agent.target_net.load_state_dict(agent.policy_net.state_dict())

print('Complete')
env.render()
env.close()
plt.ioff()
plt.show()
