#!/bin/sh
"exec" "`dirname $0`/../venvs/roto/bin/python" "$0" "$@"
"""vg-roto — one-click actor cutout: SAM 2.1 + MatAnyone on Apple Silicon.

Usage:
  vg-roto CLIP.mp4 --point 640,360                     click on the actor (frame 0 coords)
  vg-roto CLIP.mp4 --point 640,360 --point 700,200     multiple clicks refine the pick
  vg-roto CLIP.mp4 --box 400,100,900,700               or a rough box around them
  vg-roto CLIP.mp4 --point ... --neg 100,100           negative click excludes a region
  vg-roto CLIP.mp4 --grab-frame                        just save frame 0 as a PNG so you
                                                       can find click coords, then exit

Options: --out DIR (default CLIP_roto/), --frame N (prompt+start on frame N),
         --mode matte|mask (matte = soft alpha via MatAnyone, default;
                            mask  = hard binary via SAM2 video propagation),
         --max-size N (downscale min side for inference, e.g. 720 for speed)

Output in DIR:
  CLIP_alpha.mov    ProRes 4444 with real alpha  -> drop straight into Resolve
  CLIP_matte.mov    the matte alone (grayscale)
  CLIP_green.mp4    preview comp over green
"""
import argparse, os, shutil, subprocess, sys
from pathlib import Path

os.environ["PATH"] = "/opt/homebrew/bin:" + os.environ.get("PATH", "")
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")

VG = Path(__file__).resolve().parent.parent
SAM2_CKPT = VG / "models/sam2.1_hiera_large.pt"
SAM2_CFG = "configs/sam2.1/sam2.1_hiera_l.yaml"
MATANYONE = VG / "tools/MatAnyone"
RPY = VG / "venvs/roto/bin/python"


def run(cmd, **kw):
    r = subprocess.run([str(c) for c in cmd], **kw)
    if r.returncode != 0:
        sys.exit(f"command failed: {' '.join(str(c) for c in cmd)}")
    return r


def probe_video(path):
    out = subprocess.run(
        ["ffprobe", "-v", "error", "-select_streams", "v:0",
         "-show_entries", "stream=width,height,avg_frame_rate,nb_frames",
         "-of", "csv=p=0", str(path)], capture_output=True, text=True).stdout.strip()
    w, h, fr = out.split(",")[:3]
    num, den = fr.split("/")
    return int(w), int(h), (float(num) / float(den) if float(den) else 25.0)


def grab_frame(video, n, fps, dest):
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
         "-ss", f"{n / fps:.6f}", "-i", video, "-frames:v", "1", dest])


def sam2_first_frame_mask(frame_png, points, neg_points, box, mask_png):
    import numpy as np, torch
    from PIL import Image
    from sam2.build_sam import build_sam2
    from sam2.sam2_image_predictor import SAM2ImagePredictor

    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"SAM2: loading {SAM2_CKPT.name} on {device} ...")
    model = build_sam2(SAM2_CFG, str(SAM2_CKPT), device=device)
    pred = SAM2ImagePredictor(model)
    img = np.array(Image.open(frame_png).convert("RGB"))
    pred.set_image(img)

    pc = pl = bx = None
    if points or neg_points:
        pc = np.array(points + neg_points, dtype=np.float32)
        pl = np.array([1] * len(points) + [0] * len(neg_points), dtype=np.int32)
    if box:
        bx = np.array(box, dtype=np.float32)
    masks, scores, _ = pred.predict(point_coords=pc, point_labels=pl, box=bx, multimask_output=True)
    best = int(np.argmax(scores))
    m = (masks[best] > 0).astype(np.uint8) * 255
    Image.fromarray(m).save(mask_png)
    cov = 100.0 * (m > 0).mean()
    print(f"SAM2: mask score {scores[best]:.3f}, covers {cov:.1f}% of frame")
    if cov < 0.05 or cov > 95:
        print("WARNING: mask looks degenerate — check your click point (--grab-frame to inspect)")
    return mask_png


def matanyone_matte(video, mask_png, outdir, max_size):
    env = dict(os.environ)
    cmd = [RPY, "inference_matanyone.py", "-i", str(video), "-m", str(mask_png),
           "-o", str(outdir), "--max_size", str(max_size)]
    print("MatAnyone: propagating matte (this is the slow part) ...")
    r = subprocess.run([str(c) for c in cmd], cwd=MATANYONE, env=env)
    if r.returncode != 0:
        sys.exit("MatAnyone failed")
    stem = Path(video).stem
    pha = outdir / f"{stem}_pha.mp4"
    if not pha.exists():
        cands = sorted(outdir.glob("*_pha.mp4"))
        if not cands:
            sys.exit(f"MatAnyone produced no *_pha.mp4 in {outdir}")
        pha = cands[-1]
    return pha


def sam2_video_masks(video, points, neg_points, box, frame_n, workdir, fps):
    import numpy as np, torch
    from PIL import Image
    from sam2.build_sam import build_sam2_video_predictor

    frames = workdir / "frames"
    frames.mkdir(parents=True, exist_ok=True)
    if not list(frames.glob("*.jpg")):
        run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", video,
             "-q:v", "2", "-start_number", "0", frames / "%05d.jpg"])
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    print(f"SAM2 video: loading on {device} ...")
    pred = build_sam2_video_predictor(SAM2_CFG, str(SAM2_CKPT), device=device)
    state = pred.init_state(video_path=str(frames), offload_video_to_cpu=True,
                            offload_state_to_cpu=True)
    pc = pl = bx = None
    if points or neg_points:
        pc = np.array(points + neg_points, dtype=np.float32)
        pl = np.array([1] * len(points) + [0] * len(neg_points), dtype=np.int32)
    if box:
        bx = np.array(box, dtype=np.float32)
    pred.add_new_points_or_box(state, frame_idx=frame_n, obj_id=1,
                               points=pc, labels=pl, box=bx)
    masks_dir = workdir / "masks"
    masks_dir.mkdir(exist_ok=True)
    n = 0
    for fidx, _, logits in pred.propagate_in_video(state):
        m = (logits[0] > 0).cpu().numpy().squeeze().astype(np.uint8) * 255
        Image.fromarray(m).save(masks_dir / f"{fidx:05d}.png")
        n += 1
        if n % 50 == 0:
            print(f"  propagated {n} frames")
    print(f"SAM2 video: {n} frames masked")
    matte = workdir / "matte_raw.mp4"
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-framerate", f"{fps}",
         "-start_number", str(frame_n), "-i", masks_dir / "%05d.png",
         "-c:v", "libx264", "-crf", "12", "-pix_fmt", "yuv420p", matte])
    return matte


def compose_outputs(video, pha, outdir, stem, fps, w, h):
    alpha_mov = outdir / f"{stem}_alpha.mov"
    matte_mov = outdir / f"{stem}_matte.mov"
    green_mp4 = outdir / f"{stem}_green.mp4"
    scale = f"scale={w}:{h}:flags=bicubic,format=gray"
    print("composing ProRes 4444 + previews ...")
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", video, "-i", pha,
         "-filter_complex", f"[1:v]{scale}[a];[0:v][a]alphamerge,format=yuva444p10le[out]",
         "-map", "[out]", "-map", "0:a:0?", "-c:v", "prores_ks", "-profile:v", "4444",
         "-c:a", "pcm_s16le", "-shortest", alpha_mov])
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", pha,
         "-vf", scale, "-c:v", "prores_ks", "-profile:v", "1", matte_mov])
    run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y", "-i", video, "-i", pha,
         "-filter_complex",
         f"color=0x00b140:size={w}x{h}:rate={fps}[bg];"
         f"[1:v]{scale}[a];[0:v][a]alphamerge[fg];[bg][fg]overlay=shortest=1,format=yuv420p[out]",
         "-map", "[out]", "-c:v", "libx264", "-crf", "18", green_mp4])
    return alpha_mov, matte_mov, green_mp4


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("video")
    ap.add_argument("--point", action="append", default=[], help="x,y positive click (repeatable)")
    ap.add_argument("--neg", action="append", default=[], help="x,y negative click (repeatable)")
    ap.add_argument("--box", help="x1,y1,x2,y2")
    ap.add_argument("--frame", type=int, default=0, help="frame to prompt on / start from")
    ap.add_argument("--mode", choices=["matte", "mask"], default="matte")
    ap.add_argument("--max-size", type=int, default=-1, help="downscale min side for inference")
    ap.add_argument("--out", default=None)
    ap.add_argument("--grab-frame", action="store_true", help="save the prompt frame as PNG and exit")
    args = ap.parse_args()

    video = Path(args.video).expanduser().resolve()
    if not video.exists():
        sys.exit(f"no such file: {video}")
    stem = video.stem
    outdir = Path(args.out).expanduser() if args.out else video.parent / f"{stem}_roto"
    outdir.mkdir(parents=True, exist_ok=True)
    work = outdir / "work"
    work.mkdir(exist_ok=True)

    w, h, fps = probe_video(video)
    frame_png = work / f"frame{args.frame:05d}.png"
    grab_frame(video, args.frame, fps, frame_png)
    if args.grab_frame:
        print(f"prompt frame saved: {frame_png}  ({w}x{h}) — open it, note x,y of your click")
        return

    points = [tuple(float(v) for v in p.split(",")) for p in args.point]
    negs = [tuple(float(v) for v in p.split(",")) for p in args.neg]
    box = [float(v) for v in args.box.split(",")] if args.box else None
    if not points and not box:
        sys.exit("give me --point x,y or --box x1,y1,x2,y2 (use --grab-frame to find coords)")

    src = video
    if args.frame > 0 and args.mode == "matte":
        src = work / f"{stem}_from{args.frame}.mp4"
        print(f"trimming from frame {args.frame} (MatAnyone propagates forward from its mask)")
        run(["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
             "-ss", f"{args.frame / fps:.6f}", "-i", video,
             "-c:v", "libx264", "-crf", "12", "-c:a", "aac", src])

    mask_png = work / "first_mask.png"
    sam2_first_frame_mask(frame_png, points, negs, box, mask_png)

    if args.mode == "matte":
        pha = matanyone_matte(src, mask_png, work, args.max_size)
    else:
        pha = sam2_video_masks(src, points, negs, box, args.frame, work, fps)

    alpha, matte, green = compose_outputs(src, pha, outdir, stem, fps, w, h)
    print(f"\ndone:\n  {alpha}   <- ProRes 4444 with alpha, drop into Resolve\n  {matte}\n  {green}   <- preview")


if __name__ == "__main__":
    main()