import os
import ffmpeg
import argparse
import skimage
from colorize import colorize
from models.att_unet import AttentionUNet


bw_frames_root = './tmp/bw/'
rgb_frames_root = './tmp/rgb/'


def _remove_images(dir):
    for f in os.listdir(dir):
        if f.lower().endswith('jpg'):
            os.remove(os.path.join(dir, f))


def _extract_raw_frames(source_path):
    bwframe_path_template = bw_frames_root + '%5d.jpg'
    ffmpeg.input(str(source_path)).output(
        str(bwframe_path_template), format='image2', vcodec='mjpeg', qscale=0
    ).run(capture_stdout=True)


def _colorize_raw_frames(source_path, model):
    for img in os.listdir(str(bw_frames_root)):
        img_path = bw_frames_root + '/' + img

        if os.path.isfile(str(img_path)):
            color_image = colorize(img_path, model, (256, 256))
            skimage.io.imsave(str(rgb_frames_root + '/' + img), color_image)


def _get_fps(source_path):
    probe = ffmpeg.probe(str(source_path))
    stream_data = next(
        (stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
    return stream_data['avg_frame_rate']


def _build_video(source_path):
    colorized_path = 'output.mp4'
    colorframes_path_template = str(rgb_frames_root + '%5d.jpg')
    fps = _get_fps(source_path)

    ffmpeg.input(
        str(colorframes_path_template),
        format='image2',
        vcodec='mjpeg',
        framerate=fps,
    ).output(str(colorized_path), crf=17, vcodec='libx264').run(capture_stdout=True)


def colorize_video(path, model):
    _remove_images(bw_frames_root)
    _remove_images(rgb_frames_root)
    _extract_raw_frames(path)
    _colorize_raw_frames(path, model)
    return _build_video(path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Colorize provided images using selected model')
    parser.add_argument('--weights',
                        type=str,
                        help='Saved model weights - required when --model not specified. Uses Attention U-Net.')
    parser.add_argument('--source',
                        type=str,
                        help='Source image path, either grayscale or rgb')
    args = parser.parse_args()

    model = AttentionUNet((256, 256, 3))
    model.load_weights(args.weights)
    colorize_video(args.source, model)
