import copy
import glob
import os
import time
import types
from collections import deque

import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import algo
from arguments import get_args
from envs import make_vec_envs
from model import Policy
from storage import RolloutStorage
from visualize import visdom_plot, get_reward_log

import sys 
sys.path.append('/home/mehdi/PythonBuilds/reacher/reacher_4/')
import reacher_4

import matplotlib.pyplot as plt 
from tensorboardX import SummaryWriter

args = get_args()

num_updates = int(args.num_frames) // args.num_steps // args.num_processes


try:
    os.makedirs(args.log_dir)
except OSError:
    files = glob.glob(os.path.join(args.log_dir, '*.monitor.csv'))
    for f in files:
        os.remove(f)

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)



def main():

    print('Preparing parameters')

    torch.set_num_threads(1)
    device = torch.device("cuda:0" if args.cuda else "cpu")

    # print('Initializing visdom')
    # if args.vis:
    #     from visdom import Visdom
    #     viz = Visdom(port=args.port)
    #     win = None


    print('Creating envs')
    envs = make_vec_envs(args.env_name, args.seed, args.num_processes,
                        args.gamma, args.log_dir, args.add_timestep, device, False)

    print('Creating network')
    actor_critic = Policy(envs.observation_space.shape, envs.action_space,
        base_kwargs={'recurrent': args.recurrent_policy})
    actor_critic.to(device)


    print('Initializing PPO')
    agent = algo.PPO(actor_critic, args.clip_param, args.ppo_epoch, args.num_mini_batch,
                         args.value_loss_coef, args.entropy_coef, lr=args.lr,
                               eps=args.eps,
                               max_grad_norm=args.max_grad_norm)
    print('Memory')
    rollouts = RolloutStorage(args.num_steps, args.num_processes,
                        envs.observation_space.shape, envs.action_space)

    obs = envs.reset()
    rollouts.obs[0].copy_(obs)
    rollouts.to(device)

    episode_rewards = deque(maxlen=10)



    # ===================== TB visualisation =================

    writer = SummaryWriter()
    last_index = 0
    
    print('Starting ! ')

    start = time.time()
    for j in range(num_updates):
        for step in range(args.num_steps):
            # Sample actions
            with torch.no_grad():
                value, action, action_log_prob = actor_critic.act(
                        rollouts.obs[step], rollouts.masks[step])

            # Obser reward and next obs
            obs, reward, done, infos = envs.step(action)

            for info in infos:
                if 'episode' in info.keys():
                    episode_rewards.append(info['episode']['r'])

            # If done then clean the history of observations.
            masks = torch.FloatTensor([[0.0] if done_ else [1.0] for done_ in done])
            rollouts.insert(obs, action, action_log_prob, value, reward, masks)

        with torch.no_grad():
            next_value = actor_critic.get_value(rollouts.obs[-1], rollouts.masks[-1]).detach()

        rollouts.compute_returns(next_value, args.use_gae, args.gamma, args.tau)

        value_loss, action_loss, dist_entropy = agent.update(rollouts)
        writer.add_scalar('Agents metrics/Policy loss',action_loss,j)
        writer.add_scalar('Agents metrics/Value loss',value_loss,j)
        writer.add_scalar('Agents metrics/Entropy loss',dist_entropy,j)

        rollouts.after_update()

        if j % args.save_interval == 0 and args.save_dir != "":
            save_path = os.path.join(args.save_dir, args.algo)
            try:
                os.makedirs(save_path)
            except OSError:
                pass

            # A really ugly way to save a model to CPU
            save_model = actor_critic
            if args.cuda:
                save_model = copy.deepcopy(actor_critic).cpu()

            save_model = [save_model,
                            hasattr(envs.venv, 'ob_rms') and envs.venv.ob_rms or None]

            torch.save(save_model, os.path.join(save_path, args.env_name + ".pt"))

        total_num_steps = (j + 1) * args.num_processes * args.num_steps

        if j % args.log_interval == 0 and len(episode_rewards) > 1:
            end = time.time()
            print("Updates {}, num timesteps {}, FPS {} \n Last {} training episodes: mean/median reward {:.1f}/{:.1f}, min/max reward {:.1f}/{:.1f}\n".
                format(j, total_num_steps,
                       int(total_num_steps / (end - start)),
                       len(episode_rewards),
                       np.mean(episode_rewards),
                       np.median(episode_rewards),
                       np.min(episode_rewards),
                       np.max(episode_rewards), dist_entropy,
                       value_loss, action_loss))

        
        if j % args.vis_interval == 0:
            try:

                # Sometimes monitor doesn't properly flush the outputs
                # win, tx, ty = visdom_plot(viz, win, args.log_dir, args.env_name, args.algo, args.num_frames)
                tx, ty = get_reward_log(args.log_dir)
                if tx != None and ty != None: 
                    max_index = len(tx)
                    for ind_iter in range(last_index, max_index): 
                        writer.add_scalar('Reward',ty[ind_iter], tx[ind_iter])
                    last_index = max_index



                # tx, ty = get_reward_log(viz, win, args.log_dir, args.env_name,
                #                   args.algo, args.num_frames)

                # if tx != None and ty != None: 
                #     plt.cla()
                #     plt.plot(tx,ty)
                #     plt.pause(0.1)

                #     plt.show()

                # if(ty != None and tx != None): 

                #     input(ty)
                #     writer.add_scalar('Reward', ty[-1], tx[-1])
                # if(tx != None and ty != None): 
                #     plt.cla()
                #     plt.plot(tx, ty)
                #     plt.pause(0.1)
            except IOError:
                pass



if __name__ == "__main__":
    main()
