"""Original deterministic luma sequence and native FFmpeg encode study."""

import hashlib
import json
import subprocess
import sys
from pathlib import Path

import numpy as np
from PIL import Image

out = (Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).parent).resolve()
out.mkdir(parents=True, exist_ok=True)
w, h, count, fps = 640, 360, 72, 24
y, x = np.mgrid[:h, :w]
frames = []
for f in range(count):
    ramp = 38 + 105 * x / (w - 1) + 15 * np.sin((y + f * 1.5) / 31)
    texture = ((x * 73 + y * 131 + (x * y) % 67) % 81) - 40
    ramp = np.where((x >= 352) & (x < 544) & (y >= 120) & (y < 264), 130 + texture, ramp)
    bars = (x > 60) & (x < 260) & (y > 90) & (y < 270) & (((x + 2 * y + f * 2) % 17) < 3)
    ramp = np.where(bars, 219, ramp)
    disc = (x - (70 + f * 6.3)) ** 2 + (y - 295 + 22 * np.sin(f / 9)) ** 2 < 24**2
    ramp = np.where(disc, 228, ramp)
    frames.append(np.clip(np.floor(ramp + 0.5), 16, 235).astype(np.uint8))
raw = b"".join(a.tobytes() + bytes([128]) * (w * h // 2) for a in frames)
raw_path = out / "compression-source.yuv"
raw_path.write_bytes(raw)
version = subprocess.check_output(["ffmpeg", "-version"], text=True).splitlines()[0]
commands = []
for name, codec in [
    ("source", ["-c:v", "ffv1", "-level", "3"]),
    (
        "crf18",
        [
            "-c:v",
            "libx264",
            "-preset",
            "medium",
            "-crf",
            "18",
            "-g",
            "24",
            "-keyint_min",
            "24",
            "-sc_threshold",
            "0",
        ],
    ),
    (
        "crf44",
        [
            "-c:v",
            "libx264",
            "-preset",
            "medium",
            "-crf",
            "44",
            "-g",
            "24",
            "-keyint_min",
            "24",
            "-sc_threshold",
            "0",
        ],
    ),
]:
    ext = "mkv" if name == "source" else "mp4"
    target = out / f"compression-{name}.{ext}"
    cmd = [
        "ffmpeg",
        "-y",
        "-v",
        "error",
        "-f",
        "rawvideo",
        "-pix_fmt",
        "yuv420p",
        "-s",
        f"{w}x{h}",
        "-framerate",
        str(fps),
        "-i",
        raw_path.name,
        "-an",
        *codec,
        "-threads",
        "1",
        "-pix_fmt",
        "yuv420p",
        target.name,
    ]
    subprocess.run(cmd, check=True, cwd=out)
    commands.append(cmd)
records = []
decoded = {}
for name, ext in [("source", "mkv"), ("crf18", "mp4"), ("crf44", "mp4")]:
    p = out / f"compression-{name}.{ext}"
    b = subprocess.check_output(
        ["ffmpeg", "-v", "error", "-i", str(p), "-f", "rawvideo", "-pix_fmt", "yuv420p", "-"]
    )
    assert len(b) == len(raw)
    a = (
        np.frombuffer(b, dtype=np.uint8)
        .reshape(count, w * h * 3 // 2)[:, : w * h]
        .reshape(count, h, w)
    )
    decoded[name] = a
    delta = np.abs(a.astype(np.int16) - np.stack(frames).astype(np.int16))
    probe = json.loads(
        subprocess.check_output(
            [
                "ffprobe",
                "-v",
                "error",
                "-count_frames",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=width,height,pix_fmt,codec_name,nb_read_frames,r_frame_rate,duration",
                "-of",
                "json",
                str(p),
            ],
            text=True,
        )
    )["streams"][0]
    assert (
        probe["width"] == w
        and probe["height"] == h
        and probe["nb_read_frames"] == str(count)
        and probe["pix_fmt"] == "yuv420p"
    )
    if name == "source":
        assert b == raw
    per_frame = delta.mean(axis=(1, 2)).tolist()
    records.append(
        {
            "file": p.name,
            "bytes": p.stat().st_size,
            "sha256": hashlib.sha256(p.read_bytes()).hexdigest(),
            "probe": probe,
            "decodedFrames": count,
            "durationSeconds": 3,
            "meanAbsoluteStoredLumaError": float(delta.mean()),
            "perFrameLumaError": per_frame,
        }
    )
    Image.fromarray(a[36]).save(out / f"compression-{name}-luma-frame.png")
    Image.fromarray(a[36, 120:264, 352:544]).resize((576, 432), Image.Resampling.NEAREST).save(
        out / f"compression-{name}-luma-crop.png"
    )
assert records[2]["meanAbsoluteStoredLumaError"] > records[1]["meanAbsoluteStoredLumaError"] > 0
report = {
    "version": 1,
    "scope": "Original authored YUV420 luma and actual native codec outputs no AI model inference or perceptual quality ranking",
    "runtime": version,
    "width": w,
    "height": h,
    "frames": count,
    "fps": fps,
    "rawSHA256": hashlib.sha256(raw).hexdigest(),
    "frameIndex": 36,
    "frameTimeSeconds": 1.5,
    "crop": [352, 120, 192, 144],
    "cropScale": 3,
    "displayConvention": "Stored limited range luma bytes mapped directly to grayscale PNG bytes not a full RGB playback simulation",
    "records": records,
    "commands": commands,
}
(out / "compression-measurements.json").write_text(json.dumps(report, indent=2) + "\n")
(out / "compression-measurements.csv").write_text(
    "file,bytes,decoded_frames,duration_seconds,mean_absolute_stored_luma_error\n"
    + "".join(
        f"{r['file']},{r['bytes']},72,3,{r['meanAbsoluteStoredLumaError']}\n" for r in records
    )
)
print(
    json.dumps(
        [{k: r[k] for k in ["file", "bytes", "meanAbsoluteStoredLumaError"]} for r in records]
    )
)