import metaworld
import random
import argparse
import gym
import numpy as np
from itertools import count
from collections import namedtuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.distributions import Categorical
import torch_ac

class Single(nn.Module):

    def __init__(self, tasks = len(envs) ):
        super(Single, self).__init__()
        self.actor = torch.nn.ModuleList ( [ nn.Sequential(
            nn.Linear(12, 128), 
            nn.ReLU(), 
            nn.Linear(128, 64), 
            nn.ReLU(), 
            nn.Linear(64, 32), 
            nn.ReLU(), 
            nn.Linear(32, 16), 
            nn.ReLU(),             
            nn.Linear(16, 4)
        ) for i in range(tasks) ] )
        self.value_head = torch.nn.ModuleList ( [nn.Sequential(
            nn.Linear(12, 128), 
            nn.ReLU(), 
            nn.Linear(128, 64), 
            nn.ReLU(), 
            nn.Linear(64, 32), 
            nn.ReLU(), 
            nn.Linear(32, 16), 
            nn.ReLU(),             
            nn.Linear(16, 1)
        ) for i in range(tasks) ] )

        self.saved_actions = [[] for i in range(tasks)] 
        self.rewards = [[] for i in range(tasks)] 
        self.tasks = tasks 

    def forward(self, x):
        tmp=[]
        for i in range(self.tasks) :
            tmp.append( F.softmax(self.actor[i](x) - self.actor[i](x).max()))
        state_values = self.value_head[index](x)
        return tmp, state_values

def select_action(state, tasks, index):
    state = torch.tensor(list(state)).float()
    probs, state_value = model(Variable(state))

    # Obtain the most probable action for each one of the policies
    actions = []
    for i in range(tasks):
        model.saved_actions[i].append(SavedAction(probs[i].log().dot(probs[i]), state_value))
    return probs, state_value

def finish_episode(tasks, alpha, beta, gamma):

    ### Calculate loss function according to Equation 1
    R = 0
    saved_actions = model.saved_actions[index]
    policy_losses = []
    value_losses = []

    ## Obtain the discounted rewards backwards
    rewards = []
    for r in model.rewards[index][::-1]:
        R = r + gamma * R 
        rewards.insert(0, R)

    ## Standardize the rewards to be unit normal (to control the gradient estimator variance)
    rewards = torch.Tensor(rewards)
    rewards = (rewards - rewards.mean()) / (rewards.std() + np.finfo(np.float32).eps)

    for (log_prob, value), r in zip(saved_actions, rewards):
        reward = r - value.data[0]
        policy_losses.append(-log_prob * reward)
        value_losses.append(F.smooth_l1_loss(value, Variable(torch.Tensor([r]))))
        
    optimizer.zero_grad()
    loss = torch.stack(policy_losses).sum() + torch.stack(value_losses).sum()    
    loss.backward()
    optimizer.step()

    # Clean memory
    for i in range(tasks):
        del model.rewards[i][:]
        del model.saved_actions[i][:]

torch.manual_seed(1337)
SavedAction = namedtuple('SavedAction', ['log_prob', 'value'])
ml10 = metaworld.MT10() # Construct the benchmark, sampling tasks
envs = []
for name, env_cls in ml10.train_classes.items():
  env = env_cls()
  task = random.choice([task for task in ml10.train_tasks
                        if task.env_name  ==   name])
  env.set_task(task)
  envs.append(env)
  
file_name = "Single"
batch_size = 128
alpha = 0.5
beta = 0.5
gamma = 0.999
is_plot = False
num_episodes = 500
max_num_steps_per_episode = 10000
learning_rate = 0.001 
tasks = len(envs)
rewardsRec = [[] for _ in range(len(envs))]
model = Single( )
optimizer = optim.Adam(model.parameters(), lr=3e-2)

for rnd in range(10000):
    for index, env in enumerate(envs): 
          total_reward = 0
          state = env.reset()
          for t in range(200):  # Don't infinite loop while learning
              probs, state_value = select_action(state, tasks, index )
              state, reward, done, _ = env.step(probs[index].detach().numpy())
              model.rewards[index].append(reward)
              total_reward += reward
              if done:
                  break
          print(rnd, index, total_reward)
          rewardsRec[index].append(total_reward)
          finish_episode(tasks, alpha, beta, gamma )
    np.save('mt10_single_rewardsRec.npy', rewardsRec)
    torch.save(model.state_dict(), 'mt10_single_params.pkl')
