import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import matplotlib.pyplot as plt
from maze_env import Maze
from utility import decode

# This is a modified version of https://github.com/pytorch/examples/blob/master/reinforcement_learning/reinforce.py

class Policy(nn.Module):
    def __init__(self):
        super(Policy, self).__init__()
        self.affine1 = nn.Linear(100, 128)
        self.dropout = nn.Dropout(p=0.6)
        self.affine2 = nn.Linear(128, 4)
        self.gamma = 0.999 # reward_decay
        self.saved_log_probs = []
        self.rewards = []
        self.eps = []
        self.display_name="Policy Gradient"
        self.optimizer = optim.Adam(self.parameters(), lr=5e-5) # lr = 1e-3 for task 3, 1e-2 for task 2, 5e-5 task 1

    def forward(self, x):
        x = self.affine1(x)
        x = self.dropout(x)
        x = F.relu(x)
        action_scores = self.affine2(x)
        return F.softmax(action_scores, dim=1)

    def choose_action(self, state):
        temp_s = state.strip('][').split(', ')
        state = [float(coord) for coord in temp_s]
        state = torch.from_numpy(np.array(decode(state))).float().unsqueeze(0)
        probs = self(state)
        m = Categorical(probs)
        action = m.sample()
        self.saved_log_probs.append(m.log_prob(action))
        return action.item()

    def learn(self):
        R = 0
        policy_loss = []
        returns = []
        for r in self.rewards[::-1]:
            R = r + self.gamma * R
            returns.insert(0, R)
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) #/ (returns.std() + torch.tensor(self.eps))
        for log_prob, R in zip(self.saved_log_probs, returns):
            policy_loss.append(-log_prob * R)
        self.optimizer.zero_grad()
        policy_loss = torch.cat(policy_loss).sum()
        policy_loss.backward()
        self.optimizer.step()
        del self.rewards[:]
        del self.saved_log_probs[:]