import numpy as np
import gym
import torch
import torch.nn as nn
from collections import namedtuple
from tensorboardX import SummaryWriter
import shutil


# Constants (parameters) of the algo
HIDDEN_SIZE = 128
BATCH_SIZE = 16
PERCENTILE = 70

# Neural network
class Net(nn.Module):
    def __init__(self, obs_size, hidden_size, n_actions):
        super(Net, self).__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, n_actions)
        )

    def forward(self, x):
        return self.net(x)

Episode = namedtuple('Episode', field_names=['reward','steps']) # 1 Episode = total reward & collection of EpisodeStep
EpisodeStep = namedtuple('EpisodeStep', field_names=['observation','action']) # 1 EpisodeStep = observation (state) & action

# function to generate batches of episodes
def iterate_batches(env, net, batch_size):
    batch = []
    episode_reward = 0.
    episode_steps = []
    obs = env.reset()
    sm = nn.Softmax(dim=1)
    while True:
        obs_v = torch.FloatTensor([obs])
        act_probs_v = sm(net(obs_v))
        act_probs = act_probs_v.data.numpy()[0]
        action = np.random.choice(len(act_probs), p=act_probs)
        next_obs, reward, is_done, _ = env.step(action)
        episode_reward += reward
        episode_steps.append(EpisodeStep(observation=obs, action=action))
        if is_done:
            batch.append(Episode(reward=episode_reward, steps=episode_steps))
            episode_reward = 0.
            episode_steps = []
            next_obs = env.reset()
            if len(batch) == batch_size:
                yield batch
                batch = []
        obs = next_obs

# Filter PERCENTILE% of the batch
def filter_batch(batch, percentile):
    rewards = list(map(lambda s: s.reward, batch))
    reward_bound = np.percentile(rewards, percentile)
    reward_mean = float(np.mean(rewards))
    train_obs = []
    train_act = []
    for sample in batch:
        if sample.reward < reward_bound:
            continue
        train_obs.extend(map(lambda step: step.observation, sample.steps))
        train_act.extend(map(lambda step: step.action, sample.steps))
    train_obs_v = torch.FloatTensor(train_obs)
    train_act_v = torch.LongTensor(train_act)
    return train_obs_v, train_act_v, reward_bound, reward_mean

if __name__ == '__main__':
    env = gym.make('CartPole-v0')
    net = Net(env.observation_space.shape[0], HIDDEN_SIZE, env.action_space.n)
    loss = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(params = net.parameters(), lr = 0.01)
    try:
        shutil.rmtree('runs')
        writer = SummaryWriter('runs')
    except:
        writer = SummaryWriter('runs')
    for iter_nb, batch in enumerate(iterate_batches(env, net, BATCH_SIZE)):
        train_obs_v, train_act_v, reward_bound, reward_mean = filter_batch(batch, PERCENTILE)
        optimizer.zero_grad()
        action_scores_v = net(train_obs_v)
        loss_v = loss(action_scores_v, train_act_v)
        loss_v.backward()
        optimizer.step()
        print('%d: loss=%.3f, reward_mean=%.1f, reward_bound=%.1f' %(iter_nb, loss_v.item(), reward_bound, reward_mean))
        writer.add_scalar('loss', loss_v.item(), iter_nb)
        writer.add_scalar('reward mean', reward_mean, iter_nb)
        writer.add_scalar('reward bound', reward_bound, iter_nb)
        if reward_mean>199:
            print('Solved')
            break
    writer.close()

