#!/usr/bin/env python3
"""Pick the vertical crop offset that keeps the speaker's head in frame.

CapCut calls this auto-reframe. A fixed centre crop (1920 -> 1350 at y=285)
decapitates anyone framed high in the shot, which is most talking-head takes —
people leave headroom, so the face sits well above centre.

Samples faces with Apple Vision (via reel-builder's detect_subjects.py) and
places the median face at FACE_AT of the crop height, so the eyeline lands
near the upper third rather than dead centre.

Usage:
  reframe_offset.py work/master.mp4 --height 1350
  reframe_offset.py work/master.mp4 --height 1350 --frames 12
  reframe_offset.py --selftest
"""
import argparse, statistics, subprocess, sys, tempfile
from pathlib import Path

FACE_AT = 0.38          # where the face centre sits within the crop, top-down
MIN_CONF = 0.60         # Vision's own floor is 0.30, which lets in passers-by
QUANTILE = 0.25         # NOT the median: see pick_y()

# Resolved from THIS file, never the CWD: every command in the skill runs after
# a cd into the dated project folder, where a repo-relative path resolves to
# nothing and the import below fails with a misleading ModuleNotFoundError.
DEFAULT_SCRIPTS = (Path(__file__).resolve().parents[2]
                   / "tessl__reel-builder" / "scripts")


def pick_y(face_ys, q=QUANTILE):
    """One representative face height from many samples, top-down normalized.

    NOT the median. The error here is asymmetric: guessing too low crops the
    speaker's head off, guessing too high only wastes headroom. With two people
    at different heights the median lands between them and matches neither —
    measured, a speaker at 0.20 and a second face at 0.80 yields an offset that
    puts the speaker 63px ABOVE the crop. A low quantile biases toward the
    higher face, which is the safe direction.
    """
    ys = sorted(face_ys)
    return ys[min(len(ys) - 1, int(len(ys) * q))]


def offset(face_ys, src_h, crop_h, face_at=FACE_AT):
    """Top-down crop origin in pixels, clamped inside the frame.

    face_ys are TOP-DOWN normalized (0 = top of frame). Vision reports
    bottom-up, so callers must flip before getting here.
    """
    if crop_h >= src_h:
        return 0
    if not face_ys:
        return (src_h - crop_h) // 2                  # no face: centre crop
    face_px = pick_y(face_ys) * src_h
    y = int(max(0, min(src_h - crop_h, face_px - crop_h * face_at)))
    return y & ~1                       # even offset for the yuv420p chroma grid


def probe(video, entries):
    """One ffprobe field, or exit with the real error. Never guess."""
    r = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0",
                        "-show_entries", entries, "-of", "csv=p=0", video],
                       capture_output=True, text=True)
    out = r.stdout.strip()
    if r.returncode != 0 or not out or out == "N/A":
        sys.exit(f"ffprobe could not read {entries} from {video}: "
                 f"{r.stderr.strip() or 'no value'}")
    return out


def top_down(dets):
    """Confident face centres, flipped from Vision's bottom-up origin."""
    return [1.0 - cy for kind, conf, _, cy in dets
            if kind == "face" and conf >= MIN_CONF]


def face_ys(video, n, scripts):
    """Sample n frames, return top-down normalized face centres."""
    sys.path.insert(0, str(scripts))
    try:
        from detect_subjects import detections, _vision   # noqa: E402
    except ImportError:
        sys.exit(f"detect_subjects.py not found under {scripts} — "
                 f"pass --scripts with the reel-builder scripts dir.")
    if not _vision():
        sys.exit("Apple Vision not available. Install it with:\n"
                 "  pip install pyobjc-framework-Vision pyobjc-framework-Quartz\n"
                 "Without it, pick the crop offset by eye from a frame grab — "
                 "a silent centre crop would behead the speaker.")

    dur = float(probe(video, "format=duration"))

    ys = []
    with tempfile.TemporaryDirectory() as tmp:
        for i in range(n):
            t = dur * (i + 0.5) / n
            png = Path(tmp) / f"{i:03d}.png"
            subprocess.run(["ffmpeg", "-v", "error", "-ss", str(t), "-i", video,
                            "-frames:v", "1", "-y", str(png)], check=True)
            # One y per FRAME, not per detection: a frame with two faces
            # must not outvote a frame with one. Topmost = smallest y.
            frame = top_down(detections(png))
            if frame:
                ys.append(min(frame))
    return ys


def pytest_approx(x, tol=1e-9):
    class _A(float):
        def __eq__(self, o): return abs(float(self) - o) < tol
    return _A(x)


def selftest():
    # Face high in a 1920 frame (y=0.25 -> 480px) cropped to 1350:
    # 480 - 1350*0.38 = -33 -> clamps to 0, head stays in frame.
    assert offset([0.25], 1920, 1350) == 0
    # Face at 0.5 (960px): 960 - 513 = 447 -> 446 on the chroma grid,
    # well below the naive 285 centre crop.
    assert offset([0.5], 1920, 1350) == 446
    # Face low (0.9 -> 1728px) would want 1215, clamps to the 570 maximum.
    assert offset([0.9], 1920, 1350) == 570
    # No detections falls back to the centre crop.
    assert offset([], 1920, 1350) == 285
    # Crop taller than source is a no-op.
    assert offset([0.3], 1080, 1350) == 0
    # Two faces at different heights: the speaker must stay inside the crop.
    # Median would return 447 and put a face at 384px 63px ABOVE the frame.
    y = offset([0.20] * 4 + [0.80] * 4, 1920, 1350)
    assert y <= 0.20 * 1920, f"decapitates the higher face: {y}"
    # A single low outlier must not drag the crop to the bottom of the frame.
    assert offset([0.20, 0.20, 0.85], 1920, 1350) <= 0.20 * 1920
    # Vision's bottom-up origin flips to top-down, and low confidence is dropped.
    assert top_down([("face", 0.9, 0.5, 0.8)]) == [pytest_approx(0.2)]
    assert top_down([("face", 0.4, 0.5, 0.8)]) == []
    assert top_down([("human", 0.9, 0.5, 0.8)]) == []
    # Offsets stay on the chroma grid.
    assert offset([0.31], 1920, 1350) % 2 == 0
    print("ok")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("video", nargs="?")
    ap.add_argument("--height", type=int, default=1350, help="Crop height")
    ap.add_argument("--width", type=int, default=1080)
    ap.add_argument("--frames", type=int, default=8)
    ap.add_argument("--allow-center", action="store_true",
                    help="Accept a centre crop when no face is detected")
    ap.add_argument("--scripts", default=str(DEFAULT_SCRIPTS),
                    help="reel-builder scripts dir (holds detect_subjects.py)")
    ap.add_argument("--selftest", action="store_true")
    a = ap.parse_args()

    if a.selftest:
        return selftest()
    if not a.video:
        ap.error("pass a video, or --selftest")

    src_w = int(probe(a.video, "stream=width"))
    src_h = int(probe(a.video, "stream=height"))
    if a.height > src_h:
        sys.exit(f"--height {a.height} exceeds the source height {src_h}")
    if a.width > src_w:
        sys.exit(f"--width {a.width} exceeds the source width {src_w}")

    ys = face_ys(a.video, a.frames, Path(a.scripts).resolve())
    if not ys and not a.allow_center:
        sys.exit(f"No face found in {a.frames} frames of {a.video}. A centre "
                 f"crop would likely behead the speaker — check the file, "
                 f"raise --frames, or pass --allow-center to accept it.")

    x = ((src_w - a.width) // 2) & ~1
    y = offset(ys, src_h, a.height)
    print(f"crop={a.width}:{a.height}:{x}:{y}")
    print(f"  {len(ys)}/{a.frames} frames had a confident face"
          f"{' — none, centre crop' if not ys else ''}", file=sys.stderr)


if __name__ == "__main__":
    main()