import os
import time
import json
import glob
import logging
try:
    import cPickle as pickle
except ImportError:
    import pickle
import numpy as np

from gym import error
from gym.utils import atomic_write, closer
import gym.wrappers
logger = logging.getLogger(__name__)

__all__ = ['TraceRecordingWrapper', 'scan_recorded_traces']

trace_record_closer = closer.Closer()

class TraceRecordingWrapper(gym.Wrapper):
    """

    A Wrapper that records a trace of every action, observation, and reward generated by an environment.
    For an episode of length N, this will consist of:
      actions [0..N]
      observations [0..N+1]. Including the initial observation from `env.reset()`
      rewards [0..N]

    Usage:

      from gym.wrappers import TraceRecordingWrapper
      if args.record_trace:
        env = TraceRecordingWrapper(env, '/tmp/mytraces')

    It'll save a numbered series of pickled files with a manifest in /tmp/mytraces/openaigym.traces.*

    Later you can load the recorded traces:

      import gym.wrappers.trace_recording

      def episode_cb(observations, actions, rewards):
          ... do something the episode ...

      gym.wrappers.trace_recording.scan_recorded_traces('/tmp/mytraces', episode_cb)

    For an episode of length N, episode_cb receives 3 numpy arrays:
      observations.shape = [N + 1, observation_dim]
      actions.shape = [N, action_dim]
      rewards.shape = [N]


    """
    def __init__(self, env, directory):
        """
        Create a TraceRecordingWrapper around env, writing into directory
        """
        super(TraceRecordingWrapper, self).__init__(env)

        closer_id = trace_record_closer.register(self)

        self.directory = directory
        self.file_prefix = 'openaigym.trace.{}.{}'.format(closer_id, os.getpid())

        self.done = None
        self.closed = False

        self.actions = []
        self.observations = []
        self.rewards = []
        self.episode_id = 0

        self.buffered_step_count = 0
        self.buffer_batch_size = 1000

        self.episodes_first = 0
        self.episodes = []
        self.batches = []

    def _step(self, action):
        self.actions.append(action)
        observation, reward, done, info = self.env.step(action)
        if done:
            self.done = True
        self.observations.append(observation)
        self.rewards.append(reward)
        self.buffered_step_count += 1

        return observation, reward, done, info

    def _reset(self):
        assert not self.closed
        self.done = False
        observation = self.env.reset()
        self._end_episode()
        self.observations.append(observation)
        return observation

    def _end_episode(self):
        """
        if len(observations) == 0, nothing has happened yet.
        If len(observations) == 1, then len(actions) == 0, and we have only called reset and done a null episode.
        """
        if len(self.observations) > 0:
            if len(self.episodes)==0:
                self.episodes_first = self.episode_id

            self.episodes.append({
                'actions': optimize_list_of_ndarrays(self.actions),
                'observations': optimize_list_of_ndarrays(self.observations),
                'rewards': optimize_list_of_ndarrays(self.rewards),
            })
            self.actions = []
            self.observations = []
            self.rewards = []
            self.episode_id += 1

            if self.buffered_step_count >= self.buffer_batch_size:
                self._save_complete()

    def _save_complete(self):
        """
        Save the latest batch and write a manifest listing all the batches
        """

        batch_fn = '{}.ep{:06}.pickle'.format(self.file_prefix, self.episodes_first)
        p = {
            'episodes_first': self.episodes_first,
            'episodes': self.episodes,
        }
        with atomic_write.atomic_write(os.path.join(self.directory, batch_fn), True) as f:
            pickle.dump(p, f)
            bytes_per_step = float(f.tell()) / float(self.buffered_step_count)
        self.batches.append({
            'first': self.episodes_first,
            'len': len(self.episodes),
            'fn': batch_fn})

        manifest = {'batches': self.batches}
        manifest_fn = os.path.join(self.directory, '{}.manifest.json'.format(self.file_prefix))
        with atomic_write.atomic_write(os.path.join(self.directory, manifest_fn), True) as f:
            json.dump(manifest, f)

        # Adjust batch size, aiming for 5 MB per file.
        # This seems like a reasonable tradeoff between:
        #   writing speed (not too much overhead creating small files)
        #   local memory usage (buffering an entire batch before writing)
        #   random read access (loading the whole file isn't too much work when just grabbing one episode)
        self.buffer_batch_size = max(1, min(50000, int(5000000 / bytes_per_step + 1)))

        self.episodes = []
        self.episodes_first = None
        self.buffered_step_count = 0

    def close(self):
        """
        Flush any buffered data to disk and close. It should get called automatically at program exit time, but
        you can free up memory by calling it explicitly when you're done
        """
        self._end_episode()
        if len(self.episodes) > 0:
            self._save_complete()
        self.closed = True


def optimize_list_of_ndarrays(x):
    """
    Replace a list of ndarrays with a single ndarray with an extra dimension.
    Should return unchanged a list of other kinds of observations or actions, like Discrete or Tuple
    """
    if type(x) == np.ndarray:
        return x
    if len(x) == 0:
        return np.array([[]])
    if type(x[0]) == float or type(x[0]) == np.ndarray:
        return np.array(x)
    return x


def scan_recorded_traces(directory, episode_cb=None, max_episodes=None):
    """
    Go through all the traces recorded to directory, and call episode_cb for every episode.
    Set max_episodes to end after a certain number (or you can just throw an exception from episode_cb
    if you want to end the iteration early)
    """
    added_episode_count = 0
    manifest_ptn = os.path.join(directory, 'openaigym.trace.*.manifest.json')
    trace_manifest_fns = glob.glob(manifest_ptn)
    logger.info('Trace manifests %s %s', manifest_ptn, trace_manifest_fns)
    for trace_manifest_fn in trace_manifest_fns:
        trace_manifest_f = open(trace_manifest_fn, 'rb')
        trace_manifest = json.load(trace_manifest_f)
        trace_manifest_f.close()
        for batch in trace_manifest['batches']:
            batch_fn = os.path.join(directory, batch['fn'])
            batch_f = open(batch_fn, 'rb')
            batch_d = pickle.load(batch_f)
            batch_f.close()
            for ep in batch_d['episodes']:
                episode_cb(ep['observations'], ep['actions'], ep['rewards'])
                added_episode_count += 1
                if max_episodes is not None and added_episode_count >= max_episodes: return
