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

from gym import error
from gym.utils import atomic_write

class TraceRecorder(object):
    """
    Record 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]

    """
    def __init__(self, directory, file_prefix):
        self.directory = directory
        self.file_prefix = file_prefix

        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 before_step(self, action):
        assert not self.closed

        if self.done:
            raise error.ResetNeeded("Trying to step environment which is currently done. While the monitor is active, you cannot step beyond the end of an episode. Call 'env.reset()' to start the next episode.")

        self.actions.append(action)

    def after_step(self, observation, reward, done, info):
        if done:
            self.done = True
        self.observations.append(observation)
        self.rewards.append(reward)
        self.buffered_step_count += 1

    def before_reset(self):
        assert not self.closed
        self.done = False

    def after_reset(self, observation):
        self.end_episode()
        self.observations.append(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):
        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
