import gym
import numpy as np

class VideoWrapper(gym.Wrapper):
    def __init__(self, 
            env, 
            mode='rgb_array',
            enabled=True,
            steps_per_render=1,
            **kwargs
        ):
        super().__init__(env)
        
        self.mode = mode
        self.enabled = enabled
        self.render_kwargs = kwargs
        self.steps_per_render = steps_per_render

        self.frames = list()
        self.step_count = 0

    def reset(self, **kwargs):
        obs = super().reset(**kwargs)
        self.frames = list()
        self.step_count = 1
        if self.enabled:
            frame = self.env.render(
                mode=self.mode, **self.render_kwargs)
            assert frame.dtype == np.uint8
            self.frames.append(frame)
        return obs
    
    def step(self, action):
        result = super().step(action)
        self.step_count += 1
        if self.enabled and ((self.step_count % self.steps_per_render) == 0):
            frame = self.env.render(
                mode=self.mode, **self.render_kwargs)
            assert frame.dtype == np.uint8
            self.frames.append(frame)
        return result
    
    def render(self, mode='rgb_array', **kwargs):
        return self.frames