#!/usr/bin/env python

import aifc
import os
from pathlib import Path
import pickle
import sys

import av
import numpy
from PIL import Image

try:
    archive = sys.argv[1]
except:
    if os.path.exists('shape_dataset.pickle'):
        archive = 'shape_dataset.pickle'
    else:
        print("Please path an archive to sample as command parameter.")
        sys.exit(1)

print(f"Opening {archive}, peeking at the max 10 first entries...")
with open(archive, 'rb') as f:
    dataset, _ = pickle.load(f)[:10]

print("Are audio samples all of the same size?")
audio_shapes = set([dataset[i]['audio'].shape for i in range(len(dataset))])
print(f"{len(audio_shapes) == 1} ({audio_shapes})")

wdir = 'review'
Path(wdir).mkdir(parents=True, exist_ok=True)

print("Save the image data to MP4 files...")
print("Save the sound data to WAV files...")
for i, d in enumerate(dataset):
    with av.open(os.path.join(wdir, f"{d['label'][0]}_{i}.mp4"), 'w') as output:
        out = output.add_stream('mpeg4', rate=30)
        out.width = d['video'].shape[3]
        out.height = d['video'].shape[2]
        out.pix_fmt = 'yuv420p'
        aud = output.add_stream('mp3', rate=22050)
        output.gen_pts = True
        for image, sound in zip(d['video'], d['audio']):
            vf = av.video.frame.VideoFrame.from_ndarray(image.transpose((1,2,0)), format='rgb24')
            output.mux(out.encode(vf))
            af = av.audio.frame.AudioFrame.from_ndarray(sound.astype(dtype=numpy.int16), format='s16p', layout='stereo')
            af.rate = 22050
            output.mux(aud.encode(af))

print(f"Files in {wdir}, named after their supposed class.")
