# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import cv2
import imageio
import numpy as np


class VideoRecorder:
    def __init__(self, root_dir, render_size=256, fps=20):
        if root_dir is not None:
            self.save_dir = root_dir / 'eval_video'
            self.save_dir.mkdir(exist_ok=True)
        else:
            self.save_dir = None

        self.render_size = render_size
        self.fps = fps
        self.frames = []

    def init(self, env, enabled=True, metaworld=False):
        self.frames = []
        self.enabled = self.save_dir is not None and enabled
        self.record(env, metaworld)

    def record(self, env, metaworld=False):
        if self.enabled:
            if metaworld:
                render_kwargs = dict(height=self.render_size, width=self.render_size, mode='offscreen', camera_name='corner2')
                frame=env._env._env._env._env.sim.render(**render_kwargs)
                self.frames.append(frame)
            else:
                if hasattr(env, 'physics'):
                    frame = env.physics.render(height=self.render_size,
                                                width=self.render_size,
                                                camera_id=0)
                else:
                    frame = env.render()
                self.frames.append(frame)
        # if self.enabled:
        #     if hasattr(env, 'physics'):
        #         frame = env.physics.render(height=self.render_size,
        #                                     width=self.render_size,
        #                                     camera_id=0)
        #     else:
        #         frame = env.get_pixels_with_width_height(self.render_size, self.render_size)
        #     self.frames.append(frame)
    def save(self, file_name):
        if self.enabled:
            path = self.save_dir / file_name
            imageio.mimsave(str(path), self.frames, fps=self.fps)


class TrainVideoRecorder:
    def __init__(self, root_dir, render_size=256, fps=20):
        if root_dir is not None:
            self.save_dir = root_dir / 'train_video'
            self.save_dir.mkdir(exist_ok=True)
        else:
            self.save_dir = None

        self.render_size = render_size
        self.fps = fps
        self.frames = []

    def init(self, obs, enabled=True):
        self.frames = []
        self.enabled = self.save_dir is not None and enabled
        self.record(obs)

    def record(self, obs):
        if self.enabled:
            frame = cv2.resize(obs[-3:].transpose(1, 2, 0),
                               dsize=(self.render_size, self.render_size),
                               interpolation=cv2.INTER_CUBIC)
            self.frames.append(frame)

    def save(self, file_name):
        if self.enabled:
            path = self.save_dir / file_name
            imageio.mimsave(str(path), self.frames, fps=self.fps)