import numpy as np
import imageio


def load_video(filename):
    vid = imageio.get_reader(filename, 'ffmpeg')
    fps = vid.get_meta_data()['fps']
    nframes = vid.count_frames()
    return vid, fps, nframes


class SaveVideo:
    def __init__(self, outname, fps):
        self.outname = outname
        self.fps = fps

    def __enter__(self):
        self.writter = imageio.get_writer(self.outname,
                                          format='FFMPEG',
                                          fps=self.fps)
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.writter.close()

    def __iadd__(self, data):
        if np.max(data) <= 1:
            data = np.array(255*data, dtype=np.uint8)
        else:
            data = np.array(data, dtype=np.uint8)
        self.writter.append_data(data)
        return self