AI Flow: Reverse Engineering

#!/usr/bin/env python3
"""Build the run video: about two minutes, 1920x1080 at 30 fps, H.264 with AAC, from a finished run folder alone.

Called by `adom-aiflow-reverse video build`, which writes facts.json (what only the program can compute: both scores as
they stand, the step states, the parts' outlines) and passes it here. Usage:

    build.py --run <project folder> --facts <facts.json> --out <reverse-flow/video> --name <file stem>
             [--no-voice] [--workers 2] [--only scene,scene] [--stills-only]

Writes <name>.mp4, <name>.vtt, <name>.srt, <name>-poster.jpg, stills/<nn>-<scene>.jpg (the middle frame of each scene)
and video.json: every scene with the files it showed, each narration line with its times, and every figure with the
file it was read from. The last line printed is `VIDEO OK: ...` or `VIDEO FAILED: ...`.
"""
import argparse
import json
import multiprocessing as mp
import os
import shutil
import subprocess
import sys
import time
import traceback

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import numpy as np            # noqa: E402
import cv2                    # noqa: E402
import style as st            # noqa: E402
import engine                 # noqa: E402
import voice                  # noqa: E402
import scenes as sc           # noqa: E402
from engine import Layer, Timeline, XF, SHOT_XF   # noqa: E402
from rundata import Run       # noqa: E402

TARGET = 120.0               # seconds the video aims for
FILL_TO = 125.0              # optional lines are added while the total stays under this
STRETCH_TO = 118.0           # a short video is held longer on its pictures up to this
LIMITS = (110.0, 130.0)      # the range the brief asks for
LEAD, GAP, TAIL, SHOT_LEAD = 0.55, 0.24, 0.3, 0.3


def build_scenes(run, only=None):
    made, skipped = [], []
    for key, fn in sc.SEQUENCE:
        if only and key not in only:
            continue
        run.scene = key
        try:
            made.append(fn(run))
        except sc.Skip as e:
            skipped.append({"scene": key, "why": str(e)})
        except Exception as e:   # one scene failing must not cost the whole video
            skipped.append({"scene": key, "why": f"the scene could not be built: {type(e).__name__}: {e}", "trace": traceback.format_exc(limit=4)})
    run.scene = "run"
    k = 0
    for s in made:
        if s.key in ("title", "end"):
            continue
        k += 1
        rgba, x, y = st.chapter_chip(k, s.title)
        s.layers.append(Layer(rgba, x, y, fade=0))
        s.index = k
    return made, skipped


def lay_out(scene, stretch=0.0):
    """Give every shot and spoken line of a scene its time. Returns the scene's length in seconds."""
    live = [b for b in scene.beats if b["on"]]
    shots = [s for s in scene.shots if s.min_len > 0 or any(b["shot"] == s.key for b in live)]
    if not shots:
        shots = scene.shots[:1]
    speech = sum(b["dur"] for b in live) or 1.0
    t = 0.0
    for k, s in enumerate(shots):
        s.start = t
        t += LEAD if k == 0 else SHOT_LEAD + SHOT_XF * 0.5
        mine = [b for b in live if b["shot"] == s.key]
        for b in mine:
            b["start"], b["end"] = t, t + b["dur"]
            t += b["dur"] + GAP
        t += stretch * (sum(b["dur"] for b in mine) / speech if live else 1.0 / len(shots))
        t = max(t, s.start + max(s.min_len, 2.0))
        s.end = t
    t += TAIL + XF - GAP
    shots[-1].end = t
    scene.live_shots = shots
    scene.dur = t
    return t


def total_of(scenes):
    return sum(s.dur for s in scenes) - XF * (len(scenes) - 1)


def fit(scenes):
    """Choose the optional lines and the holds so the whole video lands near two minutes."""
    for s in scenes:
        for b in s.beats:
            b["on"] = b["opt"] == 0
        lay_out(s)
    # optional lines, most important first across the whole video; one that does not fit is left out and the next is tried
    for prio in sorted({b["opt"] for s in scenes for b in s.beats if b["opt"] > 0}):
        for s in scenes:
            for b in s.beats:
                if b["opt"] != prio:
                    continue
                b["on"] = True
                lay_out(s)
                if total_of(scenes) > FILL_TO:
                    b["on"] = False
                    lay_out(s)
    total = total_of(scenes)
    if total < STRETCH_TO:
        # hold longer on the pictures, most where there is most to look at; no scene more than doubles
        room = {s.key: s.dur * 1.0 for s in scenes}
        need = STRETCH_TO - total
        w = sum(s.weight * s.dur for s in scenes)
        for s in scenes:
            lay_out(s, min(room[s.key], need * s.weight * s.dur / w))
    return total_of(scenes)


def scene_frame_shots(scene):
    """The scene's shots as laid out, for the renderer."""
    scene.shots = scene.live_shots


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run", required=True)
    ap.add_argument("--facts")
    ap.add_argument("--out")
    ap.add_argument("--name")
    ap.add_argument("--no-voice", action="store_true")
    ap.add_argument("--workers", type=int, default=2)
    ap.add_argument("--only")
    ap.add_argument("--stills-only", action="store_true")
    a = ap.parse_args()
    t0 = time.time()
    facts = {}
    if a.facts and os.path.isfile(a.facts):
        with open(a.facts, encoding="utf-8") as f:
            facts = json.load(f)
    run = Run(a.run, facts)
    out = os.path.abspath(a.out or os.path.join(run.flow, "video"))
    name = a.name or (facts.get("project") or os.path.basename(run.dir)) + "-run"
    work = os.path.join(out, "work")
    os.makedirs(work, exist_ok=True)
    os.makedirs(os.path.join(out, "stills"), exist_ok=True)
    run.out = out
    doc = {"what": "The run video, built by `video build` from the run's files. Every figure below names the file it was read from.",
           "project": facts.get("project") or os.path.basename(run.dir), "tool": {"version": facts.get("version"), "renderer": "tools/video (Python: PIL, OpenCV, numpy; ffmpeg)"},
           "builtAt": facts.get("builtAt"), "ok": False}

    def fail(msg):
        doc["error"] = msg
        with open(os.path.join(out, "video.json"), "w", encoding="utf-8") as f:
            json.dump(doc, f, indent=1)
        print(f"VIDEO FAILED: {msg}")
        sys.exit(1)

    if not shutil.which(engine.ffmpeg_exe()):
        fail("ffmpeg is not installed on this machine")
    only = set(a.only.split(",")) if a.only else None
    scenes, skipped = build_scenes(run, only)
    evidence = [s for s in scenes if s.key not in ("title", "end")]
    if not evidence:
        fail("no scene could be built: " + "; ".join(f"{x['scene']}: {x['why']}" for x in skipped))
    print(f"scenes: {', '.join(s.key for s in scenes)}" + (f"; skipped: {', '.join(x['scene'] for x in skipped)}" if skipped else ""), flush=True)

    # the voice: one take per line; with no voice the lines are timed by their length and shown as captions only
    lines = [b for s in scenes for b in s.beats]
    takes, why = voice.speak([b["text"] for b in lines], os.path.join(work, "vo"), enabled=not a.no_voice)
    for k, b in enumerate(lines):
        b["dur"] = (takes[k].size / voice.RATE) if takes else voice.estimate(b["text"])
        b["take"] = takes[k] if takes else None
    total = fit(scenes)
    for s in scenes:
        scene_frame_shots(s)
    tl = Timeline(scenes)
    print(f"timeline: {tl.total:.1f} s, {tl.frames} frames, voice {'on' if takes else 'off: ' + str(why)}", flush=True)

    # stills: the middle frame of every scene, straight from the renderer
    stills = {}
    for k, s in enumerate(scenes):
        f = tl.frame(int((s.start + s.dur / 2) * st.FPS))
        p = os.path.join(out, "stills", f"{k:02d}-{s.key}.jpg")
        cv2.imwrite(p, cv2.cvtColor(f, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, 90])
        stills[s.key] = os.path.relpath(p, run.dir)
    if a.stills_only:
        print(f"VIDEO OK: stills only, {len(scenes)} scenes, {tl.total:.1f} s planned")
        return

    # frames to H.264, in parts made side by side, then joined without encoding again
    # each part is one Python process making frames and one ffmpeg with two threads: two parts keep the build near four cores
    ranges = engine.split_ranges(tl.frames, max(1, min(a.workers, 2)))
    segs = [os.path.join(work, f"seg{k}.mp4") for k in range(len(ranges))]
    try:
        if len(ranges) == 1:
            engine.encode_range(tl, ranges[0][0], ranges[0][1], segs[0], threads=2)
        else:
            ctx = mp.get_context("fork")
            procs = [ctx.Process(target=_encode, args=(tl, r[0], r[1], p)) for r, p in zip(ranges, segs)]
            for p in procs:
                p.start()
            for p in procs:
                p.join()
            if any(p.exitcode != 0 for p in procs):
                raise RuntimeError("a render worker failed (see the lines above)")
    except Exception as e:
        fail(f"rendering failed: {e}")
    lst = os.path.join(work, "segs.txt")
    with open(lst, "w") as f:
        f.writelines(f"file '{p}'\n" for p in segs)
    mp4 = os.path.join(out, name + ".mp4")
    cues = []
    clips = []
    for s in scenes:
        for b in s.beats:
            if b["on"]:
                cues.append((s.start + b["start"], s.start + b["end"], b["text"]))
                if b["take"] is not None:
                    clips.append((s.start + b["start"], b["take"]))
    cues.sort()
    ff = [engine.ffmpeg_exe(), "-y", "-v", "error", "-f", "concat", "-safe", "0", "-i", lst]
    if clips:
        wav = os.path.join(work, "voice.wav")
        voice.write_wav(wav, tl.total, clips)
        ff += ["-i", wav]
    else:
        ff += ["-f", "lavfi", "-i", "anullsrc=r=48000:cl=mono"]
    ff += ["-map", "0:v", "-map", "1:a", "-c:v", "copy", "-c:a", "aac", "-b:a", "128k", "-t", f"{tl.frames / st.FPS:.3f}", "-movflags", "+faststart", mp4]
    r = subprocess.run(ff, capture_output=True)
    if r.returncode != 0 or not os.path.isfile(mp4):
        fail("ffmpeg could not join the video and the sound: " + r.stderr.decode(errors="replace").strip()[-300:])
    voice.write_captions(os.path.join(out, name), voice.caption_cues(cues))
    # the poster: the title card once its text is up
    poster = os.path.join(out, name + "-poster.jpg")
    pf = tl.frame(int(min(scenes[0].start + scenes[0].dur * 0.6, tl.total - 1) * st.FPS))
    cv2.imwrite(poster, cv2.cvtColor(pf, cv2.COLOR_RGB2BGR), [cv2.IMWRITE_JPEG_QUALITY, 92])
    probe = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration:stream=codec_name,width,height,r_frame_rate", "-of", "json", mp4], capture_output=True)
    try:
        pj = json.loads(probe.stdout.decode() or "{}")
        dur = float(pj["format"]["duration"])
        streams = [{k: v for k, v in s.items()} for s in pj.get("streams", [])]
    except Exception:
        dur, streams = tl.total, []

    doc.update({
        "ok": True, "file": os.path.relpath(mp4, run.dir), "poster": os.path.relpath(poster, run.dir),
        "captions": [os.path.relpath(os.path.join(out, name + e), run.dir) for e in (".vtt", ".srt")],
        "captionsBurnedIn": False, "seconds": round(dur, 2), "frames": tl.frames, "size": [st.W, st.H], "fps": st.FPS, "streams": streams,
        "targetSeconds": TARGET, "withinTarget": LIMITS[0] <= dur <= LIMITS[1],
        "voice": {"used": bool(takes), "service": "adom-tts say", "voice": voice.VOICE, "rate": "default", "why": why},
        "fonts": st.FONT_USED,
        "runComplete": facts.get("complete"), "steps": facts.get("steps"),
        "scenes": [{
            "n": getattr(s, "index", None), "id": s.key, "title": s.title, "start": round(s.start, 2), "end": round(s.start + s.dur, 2),
            "shots": [{"id": x.key, "start": round(s.start + x.start, 2), "end": round(s.start + x.end, 2)} for x in s.shots],
            "narration": [{"start": round(s.start + b["start"], 2), "end": round(s.start + b["end"], 2), "text": b["text"]} for b in s.beats if b["on"]],
            "linesLeftOutForTime": [b["text"] for b in s.beats if not b["on"]],
            "files": run.files.get(s.key, []), "still": stills.get(s.key),
            "numbers": [{k: v for k, v in n.items() if k != "scene"} for n in run.numbers if n["scene"] == s.key],
        } for s in scenes],
        "skipped": skipped, "notes": run.notes, "buildSeconds": round(time.time() - t0, 1),
    })
    if not doc["withinTarget"]:
        doc["notes"].append({"scene": "run", "note": f"the video is {dur:.0f} s, outside the 110 to 130 s it aims for: " + ("too few scenes had their inputs to fill two minutes" if dur < LIMITS[0] else "the lines that must be spoken run longer than two minutes")})
    if not takes:
        doc["notes"].append({"scene": "run", "note": f"built with captions only, no voice: {why}"})
    with open(os.path.join(out, "video.json"), "w", encoding="utf-8") as f:
        json.dump(doc, f, indent=1)
    for p in segs + [lst]:
        try:
            os.remove(p)
        except OSError:
            pass
    print(f"VIDEO OK: {os.path.relpath(mp4, run.dir)}, {dur:.1f} s, {len(scenes)} scenes ({', '.join(s.key for s in scenes)})"
          + (f", {len(skipped)} skipped ({', '.join(x['scene'] for x in skipped)})" if skipped else "")
          + (", voice " + voice.VOICE if takes else f", captions only ({why})") + f", built in {time.time() - t0:.0f} s")


def _encode(tl, a, b, path):
    try:
        engine.encode_range(tl, a, b, path, threads=2)
    except Exception as e:
        print(f"render worker: {e}", flush=True)
        sys.exit(1)


if __name__ == "__main__":
    main()