# Background removal worker for a remote machine with an NVIDIA GPU. matte.ts copies this script to
# the remote machine together with the camera track and runs it in a temporary directory that holds
# in.webm. It runs the same Robust Video Matting model as the local path, through onnxruntime with
# the CUDA execution provider, and writes out.webm with the ffmpeg arguments given in the spec, so
# the result matches a local run. Stage changes and progress are reported on stderr as
# "takes:stage <label>" and "takes:progress <frames>" lines.

import json
import os
import subprocess
import sys
import urllib.request

import numpy as np
import onnxruntime as ort


def report(line):
    print(f'takes:{line}', file=sys.stderr, flush=True)


def ensure_model(url):
    path = os.path.expanduser(os.path.join('~', '.cache', 'takes', 'models', os.path.basename(url)))
    if not os.path.exists(path):
        report('stage Downloading model')
        os.makedirs(os.path.dirname(path), exist_ok=True)
        tmp = f'{path}.{os.getpid()}.tmp'
        urllib.request.urlretrieve(url, tmp)
        os.rename(tmp, path)
    return path


def to_u8(x):
    # Round half to even and clamp, as a Uint8ClampedArray does in the local path.
    return np.clip(np.rint(x * 255), 0, 255).astype(np.uint8)


def main():
    spec = json.loads(sys.argv[1])
    W, H = spec['width'], spec['height']
    model = ensure_model(spec['modelUrl'])
    report('stage Loading model on the GPU')
    if hasattr(ort, 'preload_dlls'):
        ort.preload_dlls()  # the CUDA and cuDNN libraries installed by pip
    opts = ort.SessionOptions()
    opts.log_severity_level = 3  # errors only; a warning about memcpy nodes is expected for this model
    sess = ort.InferenceSession(model, opts, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
    if 'CUDAExecutionProvider' not in sess.get_providers():
        sys.exit('onnxruntime could not use the GPU (CUDAExecutionProvider is unavailable); see the messages above')

    report('stage Removing background')
    decoder = subprocess.Popen(['ffmpeg', *spec['decode']], stdout=subprocess.PIPE)
    encoder = subprocess.Popen(['ffmpeg', *spec['encode']], stdin=subprocess.PIPE)
    rec = [np.zeros((1, 1, 1, 1), np.float32)] * 4
    dsr = np.array([spec['downsampleRatio']], np.float32)
    rgba = np.empty((H, W, 4), np.uint8)
    size = W * H * 3
    n = 0
    while True:
        buf = decoder.stdout.read(size)
        if len(buf) < size:
            break
        frame = np.frombuffer(buf, np.uint8).reshape(H, W, 3)
        src = (frame.transpose(2, 0, 1)[None] / np.float32(255)).astype(np.float32)
        fgr, pha, *rec = sess.run(
            ['fgr', 'pha', 'r1o', 'r2o', 'r3o', 'r4o'],
            {'src': src, 'r1i': rec[0], 'r2i': rec[1], 'r3i': rec[2], 'r4i': rec[3], 'downsample_ratio': dsr},
        )
        rgba[..., :3] = to_u8(fgr[0].transpose(1, 2, 0))
        rgba[..., 3] = to_u8(pha[0, 0])
        encoder.stdin.write(rgba.tobytes())
        n += 1
        if n % 10 == 0:
            report(f'progress {n}')
    encoder.stdin.close()
    if decoder.wait() != 0 or encoder.wait() != 0:
        sys.exit('ffmpeg failed')
    report(f'progress {n}')


if __name__ == '__main__':
    main()