#!/usr/bin/env python3.8
import argparse
import gizeh as gz
import numpy as np
import csv
from collections import OrderedDict
from scipy import interpolate
from moviepy.video.VideoClip import VideoClip
from moviepy.video.compositing.CompositeVideoClip import CompositeVideoClip
from moviepy.video.io.VideoFileClip import VideoFileClip
from moviepy.video.fx.resize import resize
from pathlib import Path

argp = argparse.ArgumentParser(description="Generate video from keylog and overlay on background video")
argp.add_argument("keylog_path", help="Path to the keylog txt generated by record-moves.py")
argp.add_argument("--keyvid-out", help="Destination of the small key visualization video")
argp.add_argument("--bgvid-path", help="Path to the background video file")
argp.add_argument("--keys-appear-at", type=float,
  help="Time, in seconds, from which to show the keys in the composite video")
argp.add_argument("--composite-out", help="Destination of the composite of keyvid over background")
args = argp.parse_args()

if args.bgvid_path and not args.keys_appear_at:
    argp.error("Can't generate video (from --bgvid-path) without supplying a --keys-appear-at value")

v = np.array

BORDER = 4
ALPHA = 0.5
ALPHA_PRESSED = 0.5
BORDER_RGBA = (0.1,0.1,0.2,ALPHA)
RES = (300, 200)

# Keep synced with record-moves.py
names_layout = [
    ['crouch', 'forward', 'jump'],
    ['left', 'backward', 'right'],
]

chars_layout = [
    ['C', '↑', 'J'],
    ['←', '↓', '→'],
]

def set2layout(keys_down):
    return [
      [col if col == '' else (col in keys_down) for col in row]
      for row in names_layout
    ]

def render_keys(keys_down):
    isdown_layout = set2layout(keys_down)

    s = gz.Surface(RES[0]+2*BORDER, RES[1]+2*BORDER)
    #gz.rectangle(xy=(RES[0]/2, RES[1]/2), lx=RES[0], ly=RES[1], fill=(0,1,0)).draw(s) # greenscreen?
    gz.rectangle(xy=(RES[0]/2+BORDER, RES[1]/2+BORDER), lx=RES[0], ly=RES[1], fill=BORDER_RGBA).draw(s)

    row_height = RES[1]/len(chars_layout)

    xy = v([BORDER, -row_height/2 + BORDER])

    for row, downs in zip(chars_layout, isdown_layout):
        xy += v([0, row_height]) # place pen at start pos for row

        tot_chars = sum([len(col) for col in row])

        for col, is_down in zip(row, downs):
            col_width = RES[0] * len(col) / tot_chars
            xy += v([col_width/2, 0]) # advance pen

            if col != ' ':
                fill = (0.5,0.5,1,ALPHA_PRESSED) if is_down else (0.2,0.2,0.4,ALPHA)
                gz.rectangle(xy=xy, lx=col_width, ly=row_height,
                             fill=fill, stroke=BORDER_RGBA, stroke_width=BORDER).draw(s)
                gz.text(col, 'Arial', 60, fill=(1,1,1,ALPHA), xy=xy).draw(s)

            xy += v([col_width/2, 0]) # advance pen

        xy[0] = BORDER

    return s

keys_down_timeline = OrderedDict()
get_most_recent_time = None
last_t_ms = None

print('Building timeline...')
with open(args.keylog_path, newline='') as infile:
    reader = csv.reader(infile, delimiter='\t')
    next(reader) # skip headings

    last_keys_down = None

    for entry in reader:
        time_ms = int(entry[0])
        keyname = entry[1]
        is_down = True if entry[2] == 'on' else False

        current_keys_down = set() if last_keys_down is None else set(last_keys_down)

        if is_down:
            current_keys_down.add(keyname)
        else:
            current_keys_down.remove(keyname)

        keys_down_timeline[time_ms] = current_keys_down

        last_keys_down = current_keys_down
        last_t_ms = time_ms

    ts = v(list(keys_down_timeline.keys()))
    get_most_recent_time = interpolate.interp1d(ts, ts, kind='previous')

def get_keys_down_at(time_ms):
    if time_ms > last_t_ms:
        return set()
    time = int(get_most_recent_time(time_ms))
    return keys_down_timeline[time]

def frame_at(t_sec):
    return render_keys(get_keys_down_at(t_sec*1000)).get_npimage(transparent=True)

def rgb_at(t_sec):
    return frame_at(t_sec)[:,:,:3]

def alpha_at(t_sec):
    return frame_at(t_sec)[:,:,3] / 255.0

print('Generating video(s)...')

keyvid = None

if args.keyvid_out is not None:
    keyvid = VideoClip(rgb_at, duration=last_t_ms / 1000.0 + 3)
    keyvid.write_videofile(args.keyvid_out, fps=24)

if args.bgvid_path is not None:
    # fps_source to fix 2x speed bug
    # https://github.com/Zulko/moviepy/issues/1263#issuecomment-659470362
    bgvid = VideoFileClip(args.bgvid_path, fps_source='fps')

    # Transparency in moviepy and gizeh
    # https://github.com/Zulko/moviepy/issues/898#issuecomment-452361230
    alpha_mask = VideoClip(alpha_at, duration=bgvid.duration - args.keys_appear_at, ismask=True)

    if keyvid is None:
        keyvid = VideoClip(rgb_at)
    keyvid.duration = alpha_mask.duration
    keyvid = keyvid.set_mask(alpha_mask)

    composite = CompositeVideoClip([
        bgvid,
        keyvid
            .set_position(('center', 0.6), relative=True)
            .fx(resize, 0.5)
            .set_start(args.keys_appear_at)
    ])

    out_path = args.composite_out
    if out_path is None:
        out_path = Path(args.bgvid_path)
        # Awaiting https://github.com/python/cpython/pull/19295#issue-397309516
        out_path = out_path.with_name(out_path.stem + '-with-keys' + out_path.suffix)
        out_path = str(out_path)
    composite.write_videofile(out_path, fps=24)
