__all__ = ['MonitorEval', 'get_monitor_files', 'load_results']

import gym
from gym.core import Wrapper
from os import path
import time
from glob import glob
import numpy as np
from PIL import Image
import os
import shutil

try:
    import ujson as json # Not necessary for monitor writing, but very useful for monitor loading
except ImportError:
    import json

class MonitorEval(Wrapper):
    EXT = "evalMonitor.json"
    f = None

    def __init__(self, env, filename, base_dir, allow_early_resets=False, use_monitor=True, create_video=False, curiosity_search=False, no_grid=False):
        Wrapper.__init__(self, env=env)
        self.tstart = time.time()
        if filename is None:
            self.f = None
            self.logger = None
        else:
            if not filename.endswith(MonitorEval.EXT):
                filename = filename + "." + MonitorEval.EXT
            self.f = open(filename, "wt")
            self.logger = JSONLogger(self.f)
            self.logger.writekvs({"t_start": self.tstart, "gym_version": gym.__version__,
                "env_id": env.spec.id if env.spec else 'Unknown'})
        self.allow_early_resets = allow_early_resets
        self.rewards = None
        self.intrinsics = None # ADDED BY CMS 12/10/17
        self.roomCounts = None # ADDED BY CMS 16/12/17
        self.videoFrames = None # ADDED BY CMS 16/12/17
        self.base_dir = base_dir # ADDED BY CMS 16/12/17
        self.useMonitor = use_monitor # ADDED BY CMS 16/12/17
        self.needs_reset = True
        self.episode_rewards = []
        self.episode_lengths = []
        self.total_steps = 0
        self.current_metadata = {} # extra info that gets injected into each log entry
        # Useful for metalearning where we're modifying the environment externally
        # But want our logs to know about these modifications
        ### BEGIN CMS ADDITIONS 13/12/2017
        self.curiositySearch = curiosity_search
        self.createVideo = create_video
        self.noGrid = no_grid
        ### END CMS ADDITIONS

    def __getstate__(self): # XXX
        d = self.__dict__.copy()
        if self.f:
            del d['f'], d['logger']
            d['_filename'] = self.f.name
            d['_num_episodes'] = len(self.episode_rewards)
        else:
            d['_filename'] = None
        return d
    def __setstate__(self, d):
        filename = d.pop('_filename')
        self.__dict__ = d
        if filename is not None:
            nlines = d.pop('_num_episodes') + 1
            self.f = open(filename, "r+t")
            for _ in range(nlines):
                self.f.readline()
            self.f.truncate()        
            self.logger = JSONLogger(self.f)


    def reset(self):
        if not self.allow_early_resets and not self.needs_reset:
            raise RuntimeError("Tried to reset an environment before done. If you want to allow early resets, wrap your env with MonitorEval(env, path, allow_early_resets=True)")
        self.rewards = []
        self.intrinsics = [] # ADDED BY CMS 12/10/17
        self.roomCounts = [] # ADDED BY CMS 16/12/17
        self.videoFrames = [] # ADDED BY CMS 16/12/17
        self.needs_reset = False
        return self.env.reset()

    def step(self, action):
        if self.needs_reset:
            raise RuntimeError("Tried to step environment that needs reset")
        ob, rew, done, info = self.env.step(action)
        self.rewards.append(rew)
        self.intrinsics.append(info["intrinsic"]) # ADDED BY CMS 12/10/17
        self.roomCounts.append(info["totalRooms"]) # ADDED BY CMS 16/12/17
        if done:
            self.needs_reset = True
            eprew = sum(self.rewards)
            eplen = len(self.rewards)
            epint = sum(self.intrinsics) # ADDED BY CMS 12/10/17
            eproom = np.max(self.roomCounts)
            epinfo = {"rew": eprew, "i": epint, "room": eproom, "steps": self.total_steps}
            epinfo.update(self.current_metadata)
            if self.logger and self.useMonitor:
                self.logger.writekvs(epinfo)
            if self.useMonitor:
                print("FINAL SCORE: " + str(eprew) + ", INTRINSIC: " + str(epint) + ", ROOMS: " + str(eproom))
            self.episode_rewards.append(eprew)
            self.episode_lengths.append(eplen)
            info['episode'] = epinfo
            if self.createVideo and self.useMonitor:
                baseFrame = self.env.render(mode='rgb_array')
                if (self.curiositySearch and self.env.unwrapped.curiositySearch):
                    csFrame = np.array(Image.fromarray(self.env.unwrapped.getCurrCSGrid()[:,:,0]).rotate(90).resize((baseFrame.shape[1], baseFrame.shape[0]), resample=Image.NEAREST).convert('RGB'), dtype=np.uint8)
                    self.videoFrames.append(np.concatenate((baseFrame, csFrame), 1))
                else:
                    self.videoFrames.append(baseFrame)
                outputFolder = path.join(self.base_dir, "video", "frame_")
                if path.exists(path.join(self.base_dir, "video")):
                    shutil.rmtree(path.join(self.base_dir, "video"))
                os.makedirs(path.join(self.base_dir, "video"))
                for i in range(len(self.videoFrames)):
                    Image.fromarray(self.videoFrames[i]).save(outputFolder + str(i) + ".png")
        self.total_steps += 1
        if (self.curiositySearch and self.env.unwrapped.curiositySearch and not self.noGrid):
            info["csGrid"] = self.env.unwrapped.getCurrCSGrid()
        if self.createVideo and self.useMonitor:
            baseFrame = self.env.render(mode='rgb_array')
            if (self.curiositySearch and self.env.unwrapped.curiositySearch and not self.noGrid):
                csFrame = np.array(Image.fromarray(info["csGrid"][:,:,0]).rotate(90).resize((baseFrame.shape[1], baseFrame.shape[0]), resample=Image.NEAREST).convert('RGB'), dtype=np.uint8)
                self.videoFrames.append(np.concatenate((baseFrame, csFrame), 1))
            else:
                self.videoFrames.append(baseFrame)
        return (ob, rew, done, info)

    def close(self):
        if self.f is not None:
            self.f.close()

    def get_total_steps(self):
        return self.total_steps

    def get_episode_rewards(self):
        return self.episode_rewards

    def get_episode_lengths(self):
        return self.episode_lengths

class JSONLogger(object):
    def __init__(self, file):
        self.file = file

    def writekvs(self, kvs):
        for k,v in kvs.items():
            if hasattr(v, 'dtype'):
                v = v.tolist()
                kvs[k] = float(v)
        self.file.write(json.dumps(kvs) + '\n')
        self.file.flush()


class LoadMonitorResultsError(Exception):
    pass

def get_monitor_files(dir):
    return glob(path.join(dir, "*" + MonitorEval.EXT))

def load_results(dir, raw_episodes=False):
    fnames = get_monitor_files(dir)
    if not fnames:
        raise LoadMonitorResultsError("no monitor files of the form *%s found in %s" % (MonitorEval.EXT, dir))
    episodes = []
    headers = []
    for fname in fnames:
        with open(fname, 'rt') as fh:
            lines = fh.readlines()
        header = json.loads(lines[0])
        headers.append(header)
        for line in lines[1:]:
            episode = json.loads(line)
            episode['abstime'] = header['t_start'] + episode['t']
            del episode['t']
            episodes.append(episode)
    header0 = headers[0]
    for header in headers[1:]:
        assert header['env_id'] == header0['env_id'], "mixing data from two envs"
    episodes = sorted(episodes, key=lambda e: e['abstime'])
    if raw_episodes: 
        return episodes
    else:
        return {
            'env_info': {'env_id': header0['env_id'], 'gym_version': header0['gym_version']},
            'episode_end_times': [e['abstime'] for e in episodes],
            'episode_lengths': [e['l'] for e in episodes],
            'episode_rewards': [e['r'] for e in episodes],
            'initial_reset_time': min([min(header['t_start'] for header in headers)])
        }
