import numpy as np
from policy import Policy
from value_function import NNValueFunction
import scipy.signal
from utils import Logger, Scaler
from datetime import datetime
import os
import argparse
import signal

class Agent(object):
    """docstring for Agent"""
    def __init__(self, env, gamma = 0.995, lam = 0.98, kl_targ = 0.003, batch_size = 20, hid1_mult = 10, policy_logvar = -1.0,load_model = False):
        self.env = env
        self.obs_dim = self.env.observation_space.shape[0]
        self.obs_dim += 1    # add 1 to obs dimension for time step feature (see run_episode())
        self.act_dim = self.env.action_space.shape[0]

        self.dir_path="./model_walk"
        now = datetime.utcnow().strftime("%b-%d_%H-%M-%S")  # create unique directories
        self.logger = Logger(logname='DarwinController', now=now)
        self.scaler = Scaler(self.obs_dim)

        #self.num_episodes = num_episodes
        self.gamma = gamma
        self.lam = lam
        self.kl_targ = kl_targ
        self.batch_size = batch_size
        self.hid1_mult = hid1_mult
        self.policy_logvar = policy_logvar
        self.load_model = load_model

        self.num_steps = 2048

        self.policy = Policy(self.obs_dim, self.act_dim, self.kl_targ, self.hid1_mult, self.policy_logvar,clipping_range=0.2)
        self.val_func = NNValueFunction(self.obs_dim, self.hid1_mult)

        if self.load_model:
            self.policy.load(self.dir_path+"/policy")
            self.val_func.load(self.dir_path+"/value_function")
            print('load model')



    def run_episode(self, animate=False):
        """ Run single episode with option to animate

        Args:
            env: ai gym environment
            policy: policy object with sample() method
            scaler: scaler object, used to scale/offset each observation dimension
                to a similar range
            animate: boolean, True uses env.render() method to animate episode

        Returns: 4-tuple of NumPy arrays
            observes: shape = (episode len, obs_dim)
            actions: shape = (episode len, act_dim)
            rewards: shape = (episode len,)
            unscaled_obs: useful for training scaler, shape = (episode len, obs_dim)
        """
        obs = self.env.reset()
        observes, actions, rewards, unscaled_obs = [], [], [], []
        done = False
        step = 0.0
        scale, offset = self.scaler.get()
        scale[-1] = 1.0  # don't scale time step feature
        offset[-1] = 0.0  # don't offset time step feature

        for i in range(self.num_steps):
            if animate:
                env.render()
            obs = obs.astype(np.float32).reshape((1, -1))
            obs = np.append(obs, [[step]], axis=1)  # add time step feature
            unscaled_obs.append(obs)
            obs = (obs - offset) * scale  # center and scale observations
            observes.append(obs)
            action = self.policy.sample(obs).reshape((1, -1)).astype(np.float32)
            actions.append(action)
            obs, reward, done, _ = self.env.step(np.squeeze(action, axis=0))
            obs, reward, done, _ = self.env.step(np.squeeze(action, axis=0))
            if not isinstance(reward, float):
                reward = np.asscalar(np.asarray(reward))
            rewards.append(reward)

            step += 1e-3  # increment time step feature
            i += 1
            if done:
                break


        return (np.concatenate(observes), np.concatenate(actions),
                np.array(rewards, dtype=np.float64), np.concatenate(unscaled_obs))


    def run_policy(self, episodes):
        """ Run policy and collect data for a minimum of min_steps and min_episodes

        Args:
            env: ai gym environment
            policy: policy object with sample() method
            scaler: scaler object, used to scale/offset each observation dimension
                to a similar range
            logger: logger object, used to save stats from episodes
            episodes: total episodes to run

        Returns: list of trajectory dictionaries, list length = number of episodes
            'observes' : NumPy array of states from episode
            'actions' : NumPy array of actions from episode
            'rewards' : NumPy array of (un-discounted) rewards from episode
            'unscaled_obs' : NumPy array of (un-discounted) rewards from episode
        """
        total_steps = 0
        trajectories = []
        for e in range(episodes):
            observes, actions, rewards, unscaled_obs= self.run_episode()
            total_steps += observes.shape[0]
            trajectory = {'observes': observes,
                          'actions': actions,
                          'rewards': rewards,
                          'unscaled_obs': unscaled_obs
                          }
            trajectories.append(trajectory)
        unscaled = np.concatenate([t['unscaled_obs'] for t in trajectories])
        self.scaler.update(unscaled)  # update running statistics for scaling observations
        self.logger.log({'_MeanReward': np.mean([t['rewards'].sum() for t in trajectories]),
                    'Steps': total_steps})

        return trajectories


    def discount(self, x, gamma):
        """ Calculate discounted forward sum of a sequence at each point """
        return scipy.signal.lfilter([1.0], [1.0, -gamma], x[::-1])[::-1]


    def add_disc_sum_rew(self, trajectories):
        """ Adds discounted sum of rewards to all time steps of all trajectories

        Args:
            trajectories: as returned by run_policy()
            gamma: discount

        Returns:
            None (mutates trajectories dictionary to add 'disc_sum_rew')
        """
        for trajectory in trajectories:
            if self.gamma < 0.999:  # don't scale for gamma ~= 1
                rewards = trajectory['rewards'] * (1 - self.gamma)
            else:
                rewards = trajectory['rewards']
            disc_sum_rew = self.discount(rewards, self.gamma)
            trajectory['disc_sum_rew'] = disc_sum_rew


    def add_value(self, trajectories):
        """ Adds estimated value to all time steps of all trajectories

        Args:
            trajectories: as returned by run_policy()
            val_func: object with predict() method, takes observations
                and returns predicted state value

        Returns:
            None (mutates trajectories dictionary to add 'values')
        """
        for trajectory in trajectories:
            observes = trajectory['observes']
            values = self.val_func.predict(observes)
            trajectory['values'] = values


    def add_gae(self, trajectories):
        """ Add generalized advantage estimator.
        https://arxiv.org/pdf/1506.02438.pdf

        Args:
            trajectories: as returned by run_policy(), must include 'values'
                key from add_value().
            gamma: reward discount
            lam: lambda (see paper).
                lam=0 : use TD residuals
                lam=1 : A =  Sum Discounted Rewards - V_hat(s)

        Returns:
            None (mutates trajectories dictionary to add 'advantages')
        """
        for trajectory in trajectories:
            if self.gamma < 0.999:  # don't scale for gamma ~= 1
                rewards = trajectory['rewards'] * (1 - self.gamma)
            else:
                rewards = trajectory['rewards']
            values = trajectory['values']
            # temporal differences
            if isinstance(values,np.ndarray):
                tds = rewards - values + np.append(values[1:] * self.gamma, 0)
            else:
                tds = rewards - values
            advantages = self.discount(tds, self.gamma * self.lam)
            trajectory['advantages'] = advantages


    def build_train_set(self, trajectories):
        """

        Args:
            trajectories: trajectories after processing by add_disc_sum_rew(),
                add_value(), and add_gae()

        Returns: 4-tuple of NumPy arrays
            observes: shape = (N, obs_dim)
            actions: shape = (N, act_dim)
            advantages: shape = (N,)
            disc_sum_rew: shape = (N,)
        """
        observes = np.concatenate([t['observes'] for t in trajectories])
        actions = np.concatenate([t['actions'] for t in trajectories])
        disc_sum_rew = np.concatenate([t['disc_sum_rew'] for t in trajectories])
        advantages = np.concatenate([t['advantages'] for t in trajectories])
        # normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-6)

        return observes, actions, advantages, disc_sum_rew


    def log_batch_stats(self, observes, actions, advantages, disc_sum_rew, episode):
        """ Log various batch statistics """
        self.logger.log({'_mean_obs': np.mean(observes),
                    '_min_obs': np.min(observes),
                    '_max_obs': np.max(observes),
                    '_std_obs': np.mean(np.var(observes, axis=0)),
                    '_mean_act': np.mean(actions),
                    '_min_act': np.min(actions),
                    '_max_act': np.max(actions),
                    '_std_act': np.mean(np.var(actions, axis=0)),
                    '_mean_adv': np.mean(advantages),
                    '_min_adv': np.min(advantages),
                    '_max_adv': np.max(advantages),
                    '_std_adv': np.var(advantages),
                    '_mean_discrew': np.mean(disc_sum_rew),
                    '_min_discrew': np.min(disc_sum_rew),
                    '_max_discrew': np.max(disc_sum_rew),
                    '_std_discrew': np.var(disc_sum_rew),
                    '_Episode': episode
                    })

    def save(self):
        self.policy.save(self.dir_path+"/policy")
        self.val_func.save(self.dir_path+"/value_function")

    def load(self):
        self.policy.load(self.dir_path+"/policy")
        self.val_func.load(self.dir_path+"/value_function")
