'''
the_man_with_the_personal_computer.py
Robert Luxemburg, 2020, Public Domain

The idea is to split each frame of the source video into 16x16 cells, each of
which has a size of 8x6 pixels. Obviously, higher values give more accurate
results, but take longer to render. In the first pass, each frame is resized
accordingly, and the resulting list of 48-integer vectors is saved to disk.
In the second pass, these vectors are used to populate a k-d tree (see
https://en.wikipedia.org/wiki/K-d_tree) that allows for quick lookup of nearest
neighbors. Then each frame is resized so that its cells have a size of 8x6
pixels, and for each cell, the tree is queried for the best full-frame match.
The above is optimized for speed by using cKDTree instead of KDTree, caching
the cell frames in memory, and allowing multiple instances of the script to run
in parallel. On a dual-core 2015 MacBook Pro, rendering takes around 60 hours. 
'''

import atexit
import json
import os
import sys
import time
import numpy as np
from scipy.spatial import cKDTree
from PIL import Image

video_fn = 'The Man with the Movie Camera.mkv'
video_dn = '.'.join(video_fn.split('.')[:-1])
src_dn = f'{video_dn}/src'

m = 16 # Matrix size
m2 = m ** 2
cx, cy = 8, 6 # Cell size
cs = cx * cy
w, h = 1920, 1440 # Target frame size
cell_w, cell_h = w//m, h//m
c = 1 # Color channels
fps = 24
kf = 1 # Keyframes (every kf-th source frame is a cell frame candidate)
ts, ti = 1, 0 # Number of threads, current thread index

cache = {}
cache_keys = []
cache_size = m2 * fps
cache_hits = 0

def get_image(i):
    global cache, cache_keys, cache_hits
    if i not in cache:
        cache[i] = Image.open(f'{video_dn}/src_{cell_w}x{cell_h}/{i:08d}.jpg')
        cache_keys.append(i)
        if len(cache_keys) > cache_size:
            del cache[cache_keys[0]]
            cache_keys = cache_keys[1:]
    else:
        cache_keys.remove(i)
        cache_keys.append(i)
        cache_hits += 1
    return cache[i]

def load_image(fn):
    image = Image.open(fn)
    if c == 1:
        image = image.convert('L')
    return image

def read_images():
    def save_data():
        np.save(data_fn, data)
    if not os.path.exists(src_dn):
        os.makedirs(src_dn)
        os.system(f'ffmpeg -i "{video_fn}" -q:v 1 "{src_dn}/%08d.jpg"')
    cell_dn = f'{src_dn}_{cell_w}x{cell_h}'
    os.makedirs(cell_dn, exist_ok=True)
    src_fns = [f'{src_dn}/{fn}' for fn in os.listdir(src_dn) if fn[0] != '.']
    src_fns = sorted(src_fns)
    n = len(src_fns)
    data_fn = f'{video_dn}/{cx}x{cy}.npy'
    if not os.path.exists(data_fn):
        data = np.zeros((n, cs) if c == 1 else (n, cs, c), dtype=np.uint8)
    else:
        data = np.load(data_fn)
    atexit.register(save_data)
    for i, src_fn in enumerate(src_fns):
        t = time.time()
        image = None
        cell_fn = f'{cell_dn}/{i:08d}.jpg'
        if not os.path.exists(cell_fn):
            image = load_image(src_fn)
            image.resize((cell_w, cell_h), Image.LANCZOS).save(cell_fn)
        if np.all(data[i]==0):
            if not image:
                image = load_image(src_fn)
            data[i] = np.array(image.resize((cx, cy), Image.LANCZOS).getdata())
        t = time.time() - t
        if t > 0:
            fn = os.path.basename(src_fn)
            print(f'\rreading {fn} {1/t:.1f} fps', ' ' * 16, end='')
    save_data()
    return data.astype(np.int32)

def write_images(data):
    global cache_hits
    dst_dn = f'{video_dn}/dst_{m}x{m}x{cx}x{cy}_{kf}'
    os.makedirs(dst_dn, exist_ok=True)
    src_fns = [f'{src_dn}/{fn}' for fn in os.listdir(src_dn) if fn[0] != '.']
    src_fns = sorted(src_fns)
    n = len(src_fns)
    if c > 1:
        s = data.shape
        data = data.reshape(s[0], s[1] * s[2])
    tree = cKDTree(data if kf == 1 else [data[i] for i in range(0, n, kf)])
    for i, src_fn in enumerate(src_fns):
        if i % ts != ti:
            continue
        dst_fn = f'{dst_dn}/{i:08d}.jpg'
        if os.path.exists(dst_fn):
            continue
        t = time.time()
        cache_hits = 0
        src_image = load_image(src_fn).resize((cx * m, cy * m), Image.LANCZOS)
        dst_image = Image.new('L' if c == 1 else 'RGB', (w, h))
        for y in range(m):
            for x in range(m):
                crop = (x * cx, y * cy, (x + 1) * cx, (y + 1) * cy)
                cell_data = src_image.crop(crop).getdata()
                if c > 1:
                    cell_data = np.array(cell_data)
                    s = cell_data.shape
                    cell_data = cell_data.reshape(s[0] * s[1])
                cell_image = get_image(tree.query(cell_data)[1] * kf)
                dst_image.paste(cell_image, (x * cell_w, y * cell_h))
        dst_image.save(dst_fn)
        t = time.time() - t
        fn = os.path.basename(dst_fn)
        print(f'\rwriting {fn} {cache_hits} {t:.1f} spf', ' ' * 16, end='')
    done = len([f for f in os.listdir(dst_dn) if f[0] != '.']) == n
    dst_fn = f'{dst_dn}/{m}x{m}x{cx}x{cy}_{kf}.mp4'
    if done and not os.path.exists(dst_fn):
        os.system(
            f'ffmpeg -r {fps} -start_number 0 -i "{dst_dn}/%08d.jpg" '
            f'-c:v libx264 -profile:v high -level 4.0 -pix_fmt yuv420p '
            f'-r {fps} -crf 20 "{dst_fn}"'
        )

if __name__ == '__main__':
    if len(sys.argv) == 3:
        ts, ti = int(sys.argv[1]), int(sys.argv[2])
    write_images(read_images())