"""Apply Roberto Campus's fine film grain to a finished SDR video.

Requires Python 3 and FFmpeg with ffprobe and libx264.

    python film-grain.py upscaled.mp4 finished.mp4
    python film-grain.py upscaled.mp4 grain-test.mp4 --sample 28 3 --compare

The grain is generated in FFmpeg. The output uses H.264 video and copies
the source audio. An SDR H.264/AAC MP4 is a good input for this workflow.
"""

import argparse
import json
import math
import shutil
import subprocess
from fractions import Fraction
from pathlib import Path


def even(value):
    return max(2, int(round(value / 2)) * 2)


def grain_filter(width, height, rate, seconds, opacity, compare):
    # Make fine grains, small clusters and faint larger variations.
    fine = f"color=c=0x808080:s={width}x{height}:r={rate}:d={seconds},format=gray,noise=alls=60:allf=t,gblur=sigma=0.45[gf]"
    half = f"color=c=0x808080:s={even(width / 2)}x{even(height / 2)}:r={rate}:d={seconds},format=gray,noise=alls=60:allf=t,scale={width}:{height}:flags=bicubic,gblur=sigma=0.5[gc]"
    quarter = f"color=c=0x808080:s={even(width / 4)}x{even(height / 4)}:r={rate}:d={seconds},format=gray,noise=alls=60:allf=t,scale={width}:{height}:flags=bicubic,gblur=sigma=0.8[gk]"
    graph = [
        f"[0:v:0]scale={width}:{height}:flags=lanczos,setsar=1,format=yuv420p,split=2[v][orig]",
        "[v]extractplanes=y+u+v[y][u][c]",
        fine,
        half,
        quarter,
        "[gf][gc]lut2=c0='128+(x-128)*0.55/0.9+(y-128)*0.35/0.9'[gfc]",
        "[gfc][gk]lut2=c0='128+(x-128)*0.9+(y-128)*0.1',lut=c0='clip(128+(val-128)*3,0,255)',"
        f"gblur=sigma=0.95:steps=2,tmix=frames=2,fps={rate}[g]",
        f"[y][g]blend=all_mode=overlay:all_opacity={opacity}:shortest=1[yg]",
        "[yg][u][c]mergeplanes=0x001020:yuv420p[gr]",
    ]
    if compare:
        # Align the split to the chroma samples in the encoded picture.
        split = width // 4 * 2
        graph += [
            f"[orig]crop={split}:{height}:0:0[left]",
            f"[gr]crop={width - split}:{height}:{split}:0[right]",
            "[left][right]hstack[vout]",
        ]
    else:
        graph += ["[orig]nullsink", "[gr]null[vout]"]
    return ";".join(graph)


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("source", type=Path)
    parser.add_argument("output", type=Path)
    parser.add_argument("--opacity", type=float, default=0.0914, help="grain opacity from 0 to 1 (default: 0.0914)")
    parser.add_argument("--sample", type=float, nargs=2, metavar=("START", "LENGTH"), help="test an excerpt in seconds")
    parser.add_argument("--compare", action="store_true", help="split the sample vertically, original left and grain right")
    args = parser.parse_args()

    ffmpeg, ffprobe = shutil.which("ffmpeg"), shutil.which("ffprobe")
    if not ffmpeg or not ffprobe:
        parser.error("Install FFmpeg and put both ffmpeg and ffprobe on your PATH.")
    if not args.source.is_file():
        parser.error(f"Input file does not exist: {args.source}")
    if args.output.exists():
        parser.error("Choose a new output path. Existing files are kept.")
    if not 0 <= args.opacity <= 1:
        parser.error("--opacity must be between 0 and 1.")
    if args.compare and args.sample is None:
        parser.error("--compare requires --sample START LENGTH.")

    try:
        info = json.loads(subprocess.check_output([
            ffprobe, "-v", "error", "-select_streams", "v:0", "-show_entries",
            "stream=width,height,r_frame_rate,color_transfer:format=duration", "-of", "json",
            str(args.source),
        ]))
        video = info["streams"][0]
        width, height = video["width"], video["height"]
        rate = Fraction(video["r_frame_rate"])
        duration = float(info["format"]["duration"])
    except (subprocess.CalledProcessError, KeyError, IndexError, ValueError, ZeroDivisionError):
        parser.error("Could not read the input video's size, frame rate and duration.")
    if video.get("color_transfer") in {"smpte2084", "arib-std-b67"}:
        parser.error("Use an SDR export for this Rec. 709 finishing pass.")
    if rate <= 0 or not math.isfinite(duration) or duration <= 0:
        parser.error("The input needs a valid frame rate and duration.")

    seek = []
    if args.sample:
        start, length = args.sample
        if not (math.isfinite(start) and math.isfinite(length) and 0 <= start < duration and length > 0):
            parser.error("Start must be inside the video and length must be positive.")
        duration = min(length, duration - start)
        seek = ["-ss", str(start), "-t", str(duration)]

    # Match the finishing dimensions when the upscale is a few pixels short.
    if abs(width - 1920) <= 4 and abs(height - 1080) <= 4:
        width, height = 1920, 1080
    else:
        width, height = even(width), even(height)
    if args.compare and width < 4:
        parser.error("Comparison mode needs a picture at least four pixels wide.")

    graph = grain_filter(width, height, rate, round(duration + 0.5, 3), args.opacity, args.compare)
    try:
        subprocess.run([
            ffmpeg, "-hide_banner", "-loglevel", "error", "-stats", "-nostdin", "-n",
            *seek, "-i", str(args.source), "-filter_complex", graph,
            "-map", "[vout]", "-map", "0:a?", "-c:v", "libx264", "-preset", "slow", "-crf", "14",
            "-profile:v", "high", "-tune", "grain", "-color_primaries", "bt709",
            "-color_trc", "bt709", "-colorspace", "bt709", "-c:a", "copy", "-movflags", "+faststart",
            str(args.output),
        ], check=True)
    except subprocess.CalledProcessError as error:
        parser.exit(error.returncode, "FFmpeg did not finish writing the video. See its error above.\n")
    print(f"Wrote {args.output} ({width}x{height}, {rate} fps, grain opacity {args.opacity})")


if __name__ == "__main__":
    main()