#!/usr/bin/env python
 
"""
conda create -n ppo python=3.7 numpy ipython matplotlib swig termcolor tqdm scipy tensorboard
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
pip install gym[box2d]
pip install plotille
"""
 
import argparse
import threading
from collections import deque
 
import gym
import matplotlib.pyplot as plt
import numpy as np
import plotille
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import tqdm
from IPython import embed
from torch.utils.tensorboard import SummaryWriter
 
 
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--logdir', type=str)
    parser.add_argument('--max_frames', type=int)
    parser.add_argument('--eval', action='store_true')
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--n_envs', type=int, default=8)
    parser.add_argument('--horizon', type=int, default=2048)
    parser.add_argument('--mini_batch_size', type=int, default=128)
    parser.add_argument('--lr', type=float, default=0.0005)
    parser.add_argument('--weight_decay', type=float, default=0.0)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--lam', type=float, default=0.99)
    parser.add_argument('--clip_range', type=float, default=0.2)
    parser.add_argument('--max_grad', type=float, default=10.0)
    parser.add_argument('--value_coef', type=float, default=0.5)
    parser.add_argument('--ent_coef', type=float, default=0.01)
    return parser.parse_args()
 
 
class Model(nn.Module):
    def __init__(self, n_features, n_actions):
        super().__init__()
        self.fc1 = nn.Linear(n_features, 32)
        self.norm1 = nn.LayerNorm(32)
        self.vf = nn.Linear(32, 1)
        self.policy = nn.Linear(32, n_actions)
 
        self.apply(self._init_weights)
 
    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
 
    def forward(self, x):
        x = torch.relu(self.norm1(self.fc1(x)))
        logit = self.policy(x)
        return self.vf(x), F.log_softmax(logit, -1), F.softmax(logit, -1)
 
 
if __name__ == '__main__':
 
    args = parse_args()
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    writer = SummaryWriter(args.logdir)
 
    # 환경 생성
    env = gym.vector.make('LunarLander-v2', num_envs=args.n_envs)
    env.seed(args.seed)
    n_features = env.observation_space.shape[1]
    n_actions = env.action_space[0].n
 
    # 모델 & 옵티마이저 생성
    model = Model(n_features, n_actions)
    model_old = Model(n_features, n_actions)
    model_old.load_state_dict(model.state_dict())
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
 
    # 테스트 게임 시작
    def run_test_game():
        test_env = gym.make('LunarLander-v2')
        test_env.seed(args.seed)
 
        while True:
            test_state = test_env.reset()
            score = 0
            while True:
                with torch.no_grad():
                    _, _, prob = model_old(
                        torch.from_numpy(test_state).float().view(1, -1)
                    )
                action = prob.multinomial(1).item()
                test_state, r, done, _ = test_env.step(action)
                score += r
                test_env.render()
                if done:
                    score = 0
                    break
 
    if args.eval:
        threading.Thread(target=run_test_game, daemon=True).start()
 
    # 버퍼 생성
    D_obs = np.zeros((args.horizon, args.n_envs, n_features))
    D_action = np.zeros((args.horizon, args.n_envs))
    D_reward = np.zeros((args.horizon, args.n_envs))
    D_done = np.zeros((args.horizon, args.n_envs))
    D_value = np.zeros((args.horizon, args.n_envs))
    D_logp = np.zeros((args.horizon, args.n_envs, n_actions))
 
    # 학습 시작
    frames = 0
    score = np.zeros(args.n_envs)
    n_scores_in_epoch = 0
    scores = deque(maxlen=10000)
 
    obs_prime = env.reset()
 
    while True:
 
        obs = obs_prime
        for D_i in tqdm.trange(args.horizon, desc='Rollout'):
            # 게임 플레이 & 데이터 수집
            with torch.no_grad():
                value, logp, prob = model_old(torch.from_numpy(obs).float())
                action = prob.multinomial(num_samples=1).numpy().reshape(-1)
            obs_prime, reward, done, info = env.step(action)
 
            # 점수 기록
            score += reward
            scores.extend([(frames, s) for s in score[done]])
            score[done] = 0
            n_scores_in_epoch += done.sum()
 
            # 데이터 저장
            D_obs[D_i] = obs
            D_action[D_i] = action
            D_reward[D_i] = reward / 100.
            D_done[D_i] = done
            D_value[D_i] = value.view(-1).numpy()
            D_logp[D_i] = logp.numpy()
 
            obs = obs_prime
            frames += args.n_envs
 
        # 데이터 수집 완료
        D_i = 0
        # gamma 
        gamma = args.gamma * (1 - D_done)
        # return 계산
        D_ret = np.zeros((args.horizon + 1, args.n_envs))
        with torch.no_grad():
            value, _, _ = model_old(torch.from_numpy(D_obs[-1]).float())
        D_ret[-1] = value.view(-1).numpy()
        for t in reversed(range(args.horizon)):
            D_ret[t] = D_reward[t] + gamma[t] * D_ret[t+1]
        D_ret = D_ret[:-1]
        # adv 계산
        value_ = np.vstack([D_value, value.numpy().transpose(1, 0)])
        delta = D_reward + gamma * value_[1:] - value_[:-1]
        D_adv = np.zeros((args.horizon, args.n_envs))
        gae = 0
        for t in reversed(range(args.horizon)):
            gae = gae * gamma[t] * args.lam + delta[t]
            D_adv[t] = gae
 
        # batch 차원 제거
        FD_obs = D_obs.reshape(-1, n_features)
        FD_action = D_action.reshape(-1)
        FD_logp = D_logp.reshape(-1, n_actions)
        FD_ret = D_ret.reshape(-1)
        FD_adv = D_adv.reshape(-1)
        # adv 정규화
        FD_adv = (FD_adv - FD_adv.mean()) / (FD_adv.std() + 1e-8)
 
        # 미니배치 index 준비
        idx = np.arange(args.horizon * args.n_envs)
        np.random.shuffle(idx)
        n_mini_batchs = args.horizon * args.n_envs // args.mini_batch_size
        for mb_i in tqdm.trange(n_mini_batchs, desc='Fit'):
            # 미니배치 준비
            sel = idx[mb_i * args.mini_batch_size: (mb_i+1) * args.mini_batch_size]
            obs = torch.from_numpy(FD_obs[sel]).float()
            action = torch.from_numpy(FD_action[sel]).long()
            ret = torch.from_numpy(FD_ret[sel]).float()
            adv = torch.from_numpy(FD_adv[sel]).float()
            logp_old = torch.from_numpy(FD_logp[sel]).float()
 
            # 그래프 생성
            value, logp, prob = model(obs)
            logp_a = logp.gather(1, action.view(-1, 1)).view(-1)
            logp_old_a = logp_old.gather(1, action.view(-1, 1)).view(-1)            
 
            # loss_v
            ## loss_v = 0.5 * (ret - value.view(ret.shape)).pow(2).mean()
            loss_v = F.smooth_l1_loss(value, ret.view(value.shape))
            # loss_pi
            ratios = torch.exp(logp_a - logp_old_a)
            policy_loss1 = -adv * ratios
            clip_range = args.clip_range
            policy_loss2 = -adv * torch.clamp(ratios, min=1.0 - clip_range, max=1.0 + clip_range)
            loss_pi = torch.mean(torch.max(policy_loss1, policy_loss2))
            # entropy
            entropy = -(prob * logp).sum(-1).mean()
            loss = loss_v * args.value_coef + loss_pi - args.ent_coef * entropy
 
            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad)
            optimizer.step()
 
        # 학습결과 출력
        scores_x, scores_y = zip(*scores)
        mean_score = np.mean(scores_y[-n_scores_in_epoch:])
        print(plotille.scatter(scores_x, scores_y, height=25, color_mode='byte'))
        print(plotille.histogram(scores_y[-n_scores_in_epoch:], bins=n_scores_in_epoch//5, height=25, color_mode='byte'))
        n_scores_in_epoch = 0
        print(
            f'{frames:,} => {mean_score:.3f}, '
            f'ret: {D_ret.mean():.3f}, '
            f'v: {value.mean():.3f}, '
            f'ent.: {entropy:.3f}, '
            f'L_v: {loss_v.item():.3f}, '
            f'L_pi: {loss_pi.item():.3f}, '
        )
        writer.add_scalar('metric/score', mean_score, global_step=frames)
        writer.add_scalar('metric/ret', D_ret.mean(), global_step=frames)
        writer.add_scalar('metric/v', value.mean(), global_step=frames)
        writer.add_scalar('metric/ent', entropy, global_step=frames)
        writer.add_scalar('loss/v', loss_v.item(), global_step=frames)
        writer.add_scalar('loss/pi', loss_pi.item(), global_step=frames)
 
        # target 모델 교체
        model_old.load_state_dict(model.state_dict())
 
        if args.max_frames is not None and frames > args.max_frames:
            writer.add_hparams(dict(algo='ppo'), dict(final_score=mean_score))
            break