import subprocess
import os
import shutil
import datetime
import time
import calendar

import numpy as np
from sortedcontainers import SortedDict
import tqdm
from PIL import Image

years = [2024]
sources = ["goes16", "goes18"]
ffmpeg_path = r"..\ffmpeg.exe"

for year in years:
    for source in sources:
        path_to_images = f"..\\composite\\{source}"
        output_file = f"..\\{source}_{year}_nofilt.mp4"
        interp_file = f"..\\{source}_{year}_interp_nofilt.mp4"

        starttime = calendar.timegm(datetime.datetime(year, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
        stoptime = calendar.timegm(datetime.datetime(year+1, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
        min_file_size = 350000 # Detect and remove corrupted files by filtering by file size
        max_file_size = 450000
        encoding_crf = 16
        max_frame_interp = 120

        # Set up ffmpeg to stream images from an input pipe
        command_line = f'{ffmpeg_path} -y -f image2pipe -framerate 60 -i - -c:v libx264 -crf {encoding_crf} -preset veryfast {output_file}'
        print(command_line)
        p = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE)

        black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8'))

        files_by_time = SortedDict()
        for root, dirs, files in tqdm.tqdm(os.walk(path_to_images), desc="Finding and filtering files."):
            for f in files:
                if f.endswith('.jpg'):
                    fpath = os.path.join(root, f)
                    fsize = os.path.getsize(fpath)
                    time_chunk = f.replace('-', '.').replace('_', '.').split('.')[1]
                    ftime = int(time_chunk)
                    if (fsize > min_file_size) and (fsize < max_file_size) and (ftime > starttime) and (ftime < stoptime):
                        files_by_time[ftime] = fpath

        difftimes = np.diff(files_by_time.keys())
        unique, counts = np.unique(difftimes, return_counts=True)
        interval = unique[0]
        assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval

        prevtime = files_by_time.peekitem(0)[0] - interval
        for t, f in tqdm.tqdm(files_by_time.items(), desc="FFMPEG encoding images", total = len(files_by_time)):
            framejump = (t - prevtime) // interval
            if framejump < max_frame_interp:
                for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg
                    black_image.save(p.stdin, 'jpeg', quality = 95)
            else: # Skip past intervals that are too large to fill reasonably
                print(f"Detected a frame gap of: {framejump}!")
            with open(f, 'rb') as fh:
                p.stdin.write(fh.read())
            prevtime = t

        p.stdin.close() # Close the ffmpeg input pipe
        p.wait() # Wait for ffmpeg to finish encoding

        command_line = f'{ffmpeg_path} -y -i {output_file} -vf blackframe=0,metadata=select:key=lavfi.blackframe.pblack:value=95:function=less,minterpolate=mi_mode=mci:mc_mode=aobmc:me_mode=bidir:vsbmc=1:me=fss -c:a copy -c:v libx264 -crf {encoding_crf} -preset veryfast {interp_file}'
        print(command_line)

        pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout
        output = pipe.read().decode()
        pipe.close()