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):
    """
    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, env, 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):
        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):
    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
