import mrl
from mrl.utils.misc import AttrDict
import numpy as np

class EpisodicEval(mrl.Module):
  def __init__(self):
    super().__init__('eval', required_agent_modules = ['eval_env', 'policy'], locals=locals())
  
  def __call__(self, num_episodes : int, *unused_args):
    """
    Runs num_steps steps in the environment and returns results.
    Results tracking is done here instead of in process_experience, since 
    experiences aren't "real" experiences; e.g. agent cannot learn from them.  
    """
    self.eval_mode()
    env = self.eval_env
    num_envs = env.num_envs
    
    episode_rewards, episode_steps = [], []
    discounted_episode_rewards = []
    is_successes = []
    record_success = False

    while len(episode_rewards) < num_episodes:
      state = env.reset()
      dones = np.zeros((num_envs,))
      steps = np.zeros((num_envs,))
      is_success = np.zeros((num_envs,))
      ep_rewards = [[] for _ in range(num_envs)]

      while not np.all(dones):
        action, _ = self.policy(state)
        state, reward, dones_, infos = env.step(action)
        for i, (rew, done, info) in enumerate(zip(reward, dones_, infos)):
          if dones[i]:
            continue
          ep_rewards[i].append(rew)
          steps[i] += 1
          if done:
            dones[i] = 1. 
            if 'is_success' in info:
              record_success = True
              is_success[i] = info['is_success']

      for ep_reward, step, is_succ in zip(ep_rewards, steps, is_success):
        if record_success:
          is_successes.append(is_succ)
        episode_rewards.append(sum(ep_reward))
        discounted_episode_rewards.append(discounted_sum(ep_reward, self.config.gamma))
        episode_steps.append(step)
    
    if hasattr(self, 'logger'):
      if len(is_successes):
        self.logger.add_scalar('Test/Success', np.mean(is_successes))
      self.logger.add_scalar('Test/Episode_rewards', np.mean(episode_rewards))
      self.logger.add_scalar('Test/Discounted_episode_rewards', np.mean(discounted_episode_rewards))
      self.logger.add_scalar('Test/Episode_steps', np.mean(episode_steps))

    return AttrDict({
      'rewards': episode_rewards,
      'steps': episode_steps
    })

def discounted_sum(lst, discount):
  sum = 0
  gamma = 1
  for i in lst:
    sum += gamma*i
    gamma *= discount
  return sum