"""Export the run 2 / environment 2 grasp comparison from original policy cameras.

Usage: python3 grasp_comparison.py /path/to/ft-spatula-01/eval
Requires ffmpeg and ffprobe. Both full camera views are retained; the source
camera mosaic is rearranged into engine columns, scene above wrist camera.
"""
import argparse
import hashlib
import json
from pathlib import Path
import subprocess


def run(*args):
    subprocess.run(["ffmpeg", "-v", "error", "-y", *map(str, args)], check=True)


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("evaluation_root", type=Path)
    args = parser.parse_args()
    output = Path(__file__).resolve().parent
    media = output / "media"
    inputs, durations, sources = [], [], []
    for engine in ["physx", "newton_mjwarp"]:
        root = args.evaluation_root / f"ft-spatula-01-eval-base-{engine}"
        video = root / "WoodSpatulaToBowlTask/Put_the_wooden_spatula_in_the_bowl_2_env2.mp4"
        metadata = json.loads(subprocess.check_output([
            "ffprobe", "-v", "error", "-show_entries",
            "stream=width,height,r_frame_rate:format=duration", "-of", "json", str(video)
        ], text=True))
        assert metadata["streams"][0]["width"] == 960
        assert metadata["streams"][0]["height"] == 270
        rows = [json.loads(line) for line in (root / "episode_results.jsonl").read_text().splitlines()]
        selected = next(row for row in rows if (row["run"], row["env_id"]) == (2, 2))
        log = root / "WoodSpatulaToBowlTask/log_2_env2.json"
        fields = ["run", "env_id", "success", "duration", "events", "physics_artifact", "reason"]
        sources.append({"engine": engine, "path": str(video.relative_to(args.evaluation_root)),
                        "sha256": digest(video), "video": metadata,
                        "episode": {key: selected[key] for key in fields},
                        "event_log": json.loads(log.read_text())})
        inputs.append(video)
        durations.append(float(metadata["format"]["duration"]))
    filters = []
    for i, length in enumerate(durations):
        # Extract each complete camera image, preserving its entire field of view.
        filters.append(f"[{i}:v]setpts=PTS-STARTPTS,tpad=stop_mode=clone:stop_duration={max(durations)-length:.6f},split=2[s{i}][w{i}]")
        filters.extend([f"[s{i}]crop=480:270:0:0[scene{i}]", f"[w{i}]crop=480:270:480:0[wrist{i}]"])
    filters.extend(["[scene0][scene1]hstack[top]", "[wrist0][wrist1]hstack[bottom]", "[top][bottom]vstack[v]"])
    video = media / "spatula-grasp-comparison.mp4"
    run("-i", inputs[0], "-i", inputs[1], "-filter_complex", ";".join(filters), "-map", "[v]",
        "-an", "-c:v", "libx264", "-crf", 20, "-pix_fmt", "yuv420p", "-movflags", "+faststart", video)
    poster = video.with_suffix(".jpg")
    run("-ss", 6, "-i", video, "-frames:v", 1, "-q:v", 2, poster)
    (output / "grasp-provenance.json").write_text(json.dumps({
        "study": "ft-spatula-01", "policy": "base", "selection": {"run": 2, "env_id": 2},
        "selection_note": "Selected to show attempts on the wooden target in both engines, successful placement in PhysX and repeated grasp failures followed by a drop in Newton.",
        "pairing": "Same policy sampling-noise index.",
        "rendering": "Original RTX policy cameras at 15 fps, original speed. PhysX left, Newton right. Complete scene cameras above complete wrist cameras; no camera field of view is cropped. PhysX holds its final frame after 10.2 s of recording. Newton recording ends at 40.733 s.",
        "poster_time_s": 6,
        "sources": sources,
        "outputs": [{"path": str(p.relative_to(output)), "sha256": digest(p), "bytes": p.stat().st_size} for p in [video, poster]],
    }, indent=2) + "\n")


if __name__ == "__main__":
    main()