import argparse
import numpy as np
from itertools import count
from collections import namedtuple
import cv2
import sys
import IPython
import matplotlib.pyplot as plt

import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

from snakeEnv import GridWorld

from collections import deque
import random

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/snakeNet_CNN_DQN', flush_secs=60)


class ReplayBuffer(object):
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, done):
        state      = np.expand_dims(state, 0)
        next_state      = np.expand_dims(next_state, 0)
        self.buffer.append((state, action, reward,next_state, done))
    
    def sample(self, batch_size):
        state, action, reward, next_state, done = zip(*random.sample(self.buffer, batch_size))
        return np.concatenate(state), action, reward, np.concatenate(next_state), done
    
    def __len__(self):
        return len(self.buffer)


# Snake game

env = GridWorld(grid_state=True)
env.seed(42)
torch.manual_seed(42)

class Policy(nn.Module):
    """
    implements both actor and critic in one model
    """
    def __init__(self):
        super(Policy, self).__init__()

        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, stride=2)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=2)
        self.fc1 = nn.Linear(576, 512)

        # Q layer
        self.action_head = nn.Linear(512, 4)


    def forward(self, x):
        x = x/255.0               
        x = self.conv1(x)
        x = F.leaky_relu(x, 0.1)
        x = self.conv2(x)        
        x = F.leaky_relu(x, 0.1)
        x = self.conv3(x)
        x = F.leaky_relu(x, 0.1)

        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = F.leaky_relu(x, 0.1)

        # actor: choses action to take from state s_t 
        # by returning probability of each action
        action_vals = self.action_head(x)

        return action_vals

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Training on: ", device)

temp_model = Policy()
temp_model.load_state_dict(torch.load("SnakeNet_CNN_DQN.pt"))
temp_model.eval()
temp_model.to(device)

target_model = Policy()
target_model.load_state_dict(torch.load("SnakeNet_CNN_DQN.pt"))
target_model.eval()
target_model.to(device)

temp_optimizer = optim.Adam(temp_model.parameters(), lr=7e-4)

eps = np.finfo(np.float32).eps.item()

def select_action(state, ucb):
    with torch.no_grad():
        state = torch.from_numpy(state.copy()).float().unsqueeze(0)
        action_vals = temp_model(state.to(device)) + ucb.to(device)
        action = torch.argmax(action_vals)        
        return action.item()

def finish_episode(batch):
    """
    Training code. Calculates actor and critic loss and performs backprop.
    """

    gamma = 0.99
    state, action, reward, next_state, done = batch

    state      = torch.FloatTensor(np.float32(state))
    next_state = torch.FloatTensor(np.float32(next_state))
    action     = torch.LongTensor(action).reshape(-1,1)
    reward     = torch.FloatTensor(reward).reshape(-1,1)
    done       = torch.FloatTensor(done).reshape(-1,1)

    values = temp_model(state.to(device)).gather(1,action.to(device))
    next_values = target_model(next_state.to(device)).max(1).values.reshape(-1,1).detach()

    expected_values = reward.to(device) + gamma * next_values * (1 - done.to(device))
    loss = (values - expected_values).pow(2).mean()
    
    # perform backprop
    temp_optimizer.zero_grad()
    loss.backward()
    temp_optimizer.step()

    # reset rewards and action buffer
    loss_val = loss.item()

    return loss_val


def main():
    running_reward = 0
    render = False
    log_interval = 10
    plotting = False

    # Exploration settings
    N_visits = torch.ones((4,))
    c = 10.0 # exploration coefficient
    min_c = 1.0
    c_decay = 0.999
    number_of_frames = 1

    # setup replay buffer
    batch_size = 32
    max_buffer_size = 100000
    initial_buffer_size = 10000
    replay_buffer = ReplayBuffer(max_buffer_size)

    # used for frame multiplication
    # rot = 1
    act_90 = [3, 2, 0, 1]
    # rot = 2
    act_180 = [1, 0, 3, 2]
    # rot = 3
    act_270 = [2, 3, 1, 0]
    # flip
    act_flip = [1, 0, 2, 3]
    
    rotate_action_list = [act_90, act_180, act_270]

    # updating target model every so many episodes
    update_target_freq = 5

    losses=[]
    avg_rewards=[]

    if plotting:
        plt.ion()    
        fig = plt.figure()
        ax1 = fig.add_subplot(211)
        ax2 = fig.add_subplot(212)
        line0, = ax1.plot([],[])
        line1, = ax2.plot([],[])
        ax1.set_ylabel("Losses")
        ax2.set_ylabel("Average Rewards")
        ax1.grid()
        ax2.grid()

    # run inifinitely many episodes
    for i_episode in count(1):

        # reset environment and episode reward
        state = env.reset()
        ep_reward = 0

        c = max(min_c, c_decay*c)

        # for each episode, only run 9999 steps so that we don't 
        # infinite loop while learning
        for t in range(1, 10000):
            if t % 200 == 0 and ep_reward == 0:
              env.food = np.array(env.randomPosition())
              env.update_state()
              state = env.state

            # select action from policy
            c = 0.0
            ucb = c*torch.sqrt(np.log(number_of_frames)/N_visits)
            action = select_action(state, ucb)

            # update ucb parameters
            N_visits[action] += 1
            number_of_frames += 1

            # take the action
            next_state, reward, done = env.step(action)
            reward = np.sign(reward)

            if render:
                env.render()


            replay_buffer.push(state.copy(), int(action), float(reward), next_state.copy(), float(done))

            ep_reward += reward

            # perform training
            if len(replay_buffer) >= initial_buffer_size:
                loss_val = finish_episode(replay_buffer.sample(batch_size))
            else:
                loss_val = 0.0

            if done:
                break

            # multiply state by 7 through flipping / rotating frame
            for flip in [False,True]:
                for rot in range(1,4):
                    if flip:
                        f_state = np.flip(state.reshape(20,20,3), 1).reshape(3,20,20)
                        f_next_state = np.flip(next_state.reshape(20,20,3),1).reshape(3,20,20)
                    else:
                        f_state = state.copy()
                        f_next_state = next_state.copy()
                    rotated_state = np.rot90(f_state.reshape(20,20,3), rot).reshape(3,20,20) # state gets rotated
                    rotated_next_state = np.rot90(f_next_state.reshape(20,20,3), rot).reshape(3,20,20) # state gets rotated
                    if flip:
                        f_action = act_flip[action]
                    else:
                        f_action = int(action)
                    rotated_action = rotate_action_list[rot-1][f_action] # action gets rotated
                    replay_buffer.push(rotated_state.copy(), int(rotated_action), float(reward), rotated_next_state.copy(), float(done))

            # update state for next iteration
            state = next_state.copy()


        # update target model
        if len(replay_buffer) >= initial_buffer_size:
            if i_episode % update_target_freq == 0:
                # copy temp model parameters into target model
                target_model.load_state_dict(temp_model.state_dict())
                target_model.eval()

        # update cumulative reward
        running_reward = 0.05 * ep_reward + (1 - 0.05) * running_reward

        # record average reward and loss
        avg_rewards.append(running_reward)
        if loss_val == 0 and len(losses)>0:
            losses.append(losses[-1])
        else:
            losses.append(loss_val)
        writer.add_scalar('Training loss', loss_val, i_episode)
        writer.add_scalar('Average reward', running_reward, i_episode)

        # log results
        if i_episode % log_interval == 0:
            print('Episode {}\tLast reward: {:.2f}\tAverage reward: {:.2f}'.format(
                  i_episode, ep_reward, running_reward))
            print(N_visits)
            if plotting:
                line0.set_ydata(losses)
                line1.set_ydata(avg_rewards)
                line0.set_xdata(np.arange(len(losses)))
                line1.set_xdata(np.arange(len(avg_rewards)))
                if len(losses) > 0:
                    ax1.set_xlim(0,len(losses))
                    ax2.set_xlim(0,len(losses))
                    ax1.set_ylim(np.min(losses), np.max(losses))
                    ax2.set_ylim(np.min(avg_rewards), np.max(avg_rewards))

                fig.canvas.draw()
                plt.pause(0.1)

        # save model every so often
        if i_episode % 1000 == 0:
            print("Saving model")
            torch.save(target_model.state_dict(), "SnakeNet_CNN_DQN.pt")

        # check if we have "solved" the problem
        if running_reward > 100:
            print("Solved! Running reward is now {} and "
                  "the last episode runs to {} time steps!".format(running_reward, t))
            break

# save final model
print("Saving model")
torch.save(target_model.state_dict(), "SnakeNet_CNN_DQN.pt")

if __name__ == '__main__':
    try:
        main()
    finally:
        # save final model
        print("Saving model")
        torch.save(target_model.state_dict(), "SnakeNet_CNN_DQN.pt")

import argparse
import numpy as np
from itertools import count
from collections import namedtuple
import cv2
import sys
import IPython
import matplotlib.pyplot as plt

import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical

from snakeEnv import GridWorld

from collections import deque
import random

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/snakeNet_CNN_DQN', flush_secs=60)


class ReplayBuffer(object):
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, done):
        state      = np.expand_dims(state, 0)
        next_state      = np.expand_dims(next_state, 0)
        self.buffer.append((state, action, reward,next_state, done))
    
    def sample(self, batch_size):
        state, action, reward, next_state, done = zip(*random.sample(self.buffer, batch_size))
        return np.concatenate(state), action, reward, np.concatenate(next_state), done
    
    def __len__(self):
        return len(self.buffer)


# Snake game

env = GridWorld(grid_state=True)
env.seed(42)
torch.manual_seed(42)

class Policy(nn.Module):
    """
    implements both actor and critic in one model
    """
    def __init__(self):
        super(Policy, self).__init__()

        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, stride=2)
        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=2)
        self.fc1 = nn.Linear(576, 512)

        # Q layer
        self.action_head = nn.Linear(512, 4)


    def forward(self, x):
        x = x/255.0               
        x = self.conv1(x)
        x = F.leaky_relu(x, 0.1)
        x = self.conv2(x)        
        x = F.leaky_relu(x, 0.1)
        x = self.conv3(x)
        x = F.leaky_relu(x, 0.1)

        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = F.leaky_relu(x, 0.1)

        # actor: choses action to take from state s_t 
        # by returning probability of each action
        action_vals = self.action_head(x)

        return action_vals

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Training on: ", device)

temp_model = Policy()
temp_model.load_state_dict(torch.load("SnakeNet_CNN_DQN.pt"))
temp_model.eval()
temp_model.to(device)

target_model = Policy()
target_model.load_state_dict(torch.load("SnakeNet_CNN_DQN.pt"))
target_model.eval()
target_model.to(device)

temp_optimizer = optim.Adam(temp_model.parameters(), lr=7e-4)

eps = np.finfo(np.float32).eps.item()

def select_action(state, ucb):
    with torch.no_grad():
        state = torch.from_numpy(state.copy()).float().unsqueeze(0)
        action_vals = temp_model(state.to(device)) + ucb.to(device)
        action = torch.argmax(action_vals)        
        return action.item()

def finish_episode(batch):
    """
    Training code. Calculates actor and critic loss and performs backprop.
    """

    gamma = 0.99
    state, action, reward, next_state, done = batch

    state      = torch.FloatTensor(np.float32(state))
    next_state = torch.FloatTensor(np.float32(next_state))
    action     = torch.LongTensor(action).reshape(-1,1)
    reward     = torch.FloatTensor(reward).reshape(-1,1)
    done       = torch.FloatTensor(done).reshape(-1,1)

    values = temp_model(state.to(device)).gather(1,action.to(device))
    next_values = target_model(next_state.to(device)).max(1).values.reshape(-1,1).detach()

    expected_values = reward.to(device) + gamma * next_values * (1 - done.to(device))
    loss = (values - expected_values).pow(2).mean()
    
    # perform backprop
    temp_optimizer.zero_grad()
    loss.backward()
    temp_optimizer.step()

    # reset rewards and action buffer
    loss_val = loss.item()

    return loss_val


def main():
    running_reward = 0
    render = False
    log_interval = 10
    plotting = False

    # Exploration settings
    N_visits = torch.ones((4,))
    c = 10.0 # exploration coefficient
    min_c = 1.0
    c_decay = 0.999
    number_of_frames = 1

    # setup replay buffer
    batch_size = 32
    max_buffer_size = 100000
    initial_buffer_size = 10000
    replay_buffer = ReplayBuffer(max_buffer_size)

    # used for frame multiplication
    # rot = 1
    act_90 = [3, 2, 0, 1]
    # rot = 2
    act_180 = [1, 0, 3, 2]
    # rot = 3
    act_270 = [2, 3, 1, 0]
    # flip
    act_flip = [1, 0, 2, 3]
    
    rotate_action_list = [act_90, act_180, act_270]

    # updating target model every so many episodes
    update_target_freq = 5

    losses=[]
    avg_rewards=[]

    if plotting:
        plt.ion()    
        fig = plt.figure()
        ax1 = fig.add_subplot(211)
        ax2 = fig.add_subplot(212)
        line0, = ax1.plot([],[])
        line1, = ax2.plot([],[])
        ax1.set_ylabel("Losses")
        ax2.set_ylabel("Average Rewards")
        ax1.grid()
        ax2.grid()

    # run inifinitely many episodes
    for i_episode in count(1):

        # reset environment and episode reward
        state = env.reset()
        ep_reward = 0

        c = max(min_c, c_decay*c)

        # for each episode, only run 9999 steps so that we don't 
        # infinite loop while learning
        for t in range(1, 10000):
            if t % 200 == 0 and ep_reward == 0:
              env.food = np.array(env.randomPosition())
              env.update_state()
              state = env.state

            # select action from policy
            c = 0.0
            ucb = c*torch.sqrt(np.log(number_of_frames)/N_visits)
            action = select_action(state, ucb)

            # update ucb parameters
            N_visits[action] += 1
            number_of_frames += 1

            # take the action
            next_state, reward, done = env.step(action)
            reward = np.sign(reward)

            if render:
                env.render()


            replay_buffer.push(state.copy(), int(action), float(reward), next_state.copy(), float(done))

            ep_reward += reward

            # perform training
            if len(replay_buffer) >= initial_buffer_size:
                loss_val = finish_episode(replay_buffer.sample(batch_size))
            else:
                loss_val = 0.0

            if done:
                break

            # multiply state by 7 through flipping / rotating frame
            for flip in [False,True]:
                for rot in range(1,4):
                    if flip:
                        f_state = np.flip(state.reshape(20,20,3), 1).reshape(3,20,20)
                        f_next_state = np.flip(next_state.reshape(20,20,3),1).reshape(3,20,20)
                    else:
                        f_state = state.copy()
                        f_next_state = next_state.copy()
                    rotated_state = np.rot90(f_state.reshape(20,20,3), rot).reshape(3,20,20) # state gets rotated
                    rotated_next_state = np.rot90(f_next_state.reshape(20,20,3), rot).reshape(3,20,20) # state gets rotated
                    if flip:
                        f_action = act_flip[action]
                    else:
                        f_action = int(action)
                    rotated_action = rotate_action_list[rot-1][f_action] # action gets rotated
                    replay_buffer.push(rotated_state.copy(), int(rotated_action), float(reward), rotated_next_state.copy(), float(done))

            # update state for next iteration
            state = next_state.copy()


        # update target model
        if len(replay_buffer) >= initial_buffer_size:
            if i_episode % update_target_freq == 0:
                # copy temp model parameters into target model
                target_model.load_state_dict(temp_model.state_dict())
                target_model.eval()

        # update cumulative reward
        running_reward = 0.05 * ep_reward + (1 - 0.05) * running_reward

        # record average reward and loss
        avg_rewards.append(running_reward)
        if loss_val == 0 and len(losses)>0:
            losses.append(losses[-1])
        else:
            losses.append(loss_val)
        writer.add_scalar('Training loss', loss_val, i_episode)
        writer.add_scalar('Average reward', running_reward, i_episode)

        # log results
        if i_episode % log_interval == 0:
            print('Episode {}\tLast reward: {:.2f}\tAverage reward: {:.2f}'.format(
                  i_episode, ep_reward, running_reward))
            print(N_visits)
            if plotting:
                line0.set_ydata(losses)
                line1.set_ydata(avg_rewards)
                line0.set_xdata(np.arange(len(losses)))
                line1.set_xdata(np.arange(len(avg_rewards)))
                if len(losses) > 0:
                    ax1.set_xlim(0,len(losses))
                    ax2.set_xlim(0,len(losses))
                    ax1.set_ylim(np.min(losses), np.max(losses))
                    ax2.set_ylim(np.min(avg_rewards), np.max(avg_rewards))

                fig.canvas.draw()
                plt.pause(0.1)

        # save model every so often
        if i_episode % 1000 == 0:
            print("Saving model")
            torch.save(target_model.state_dict(), "SnakeNet_CNN_DQN.pt")

        # check if we have "solved" the problem
        if running_reward > 100:
            print("Solved! Running reward is now {} and "
                  "the last episode runs to {} time steps!".format(running_reward, t))
            break

# save final model
print("Saving model")
torch.save(target_model.state_dict(), "SnakeNet_CNN_DQN.pt")

if __name__ == '__main__':
    try:
        main()
    finally:
        # save final model
        print("Saving model")
        torch.save(target_model.state_dict(), "SnakeNet_CNN_DQN.pt")

# Snake game implements both actor and critic in one model # Q layer # actor: choses action to take from state s_t # by returning probability of each action Training code. Calculates actor and critic loss and performs backprop. # perform backprop # reset rewards and action buffer # Exploration settings # exploration coefficient # setup replay buffer # used for frame multiplication # rot = 1 # rot = 2 # rot = 3 # flip # updating target model every so many episodes # run inifinitely many episodes # reset environment and episode reward # for each episode, only run 9999 steps so that we don't # infinite loop while learning # select action from policy # update ucb parameters # take the action # perform training # multiply state by 7 through flipping / rotating frame # state gets rotated # state gets rotated # action gets rotated # update state for next iteration # update target model # copy temp model parameters into target model # update cumulative reward # record average reward and loss # log results # save model every so often # check if we have "solved" the problem # save final model # save final model