"""stereo_depth — SBS stereoscopic video/frame → SGBM disparity maps + mid-frame point cloud.

True geometric depth from the built-in L/R parallax (cv2.StereoSGBM). ffmpeg (system)
decodes/splits; depth PNGs + preview mp4 + a colored .ply land as the outputs.
ponytail: uncalibrated disparity (SBS rips have no camera intrinsics) — z is relative,
perfect for Blender displace planes; not metric reconstruction.
"""
import argparse
import glob
import json
import os
import shutil
import subprocess
import sys
import tempfile
import zipfile

import cv2
import numpy as np

FFMPEG = shutil.which("ffmpeg") or "/opt/homebrew/bin/ffmpeg"   # headless ssh has no brew PATH

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

if not a.input:
    print("ERROR: no input")
    sys.exit(1)

src = a.input[0]
outdir = a.outdir
max_frames = int(p.get("max_frames", 120))
num_disp = max(16, (int(p.get("num_disp", 128)) // 16) * 16)
block = int(p.get("block", 5)) | 1
tab = p.get("layout", "sbs") == "tab"

work = tempfile.mkdtemp(prefix="stereo_")
frames_dir = os.path.join(work, "frames")
depth_dir = os.path.join(work, "depth")
os.makedirs(frames_dir)
os.makedirs(depth_dir)

# decode → frames (a single image just becomes frame 1)
subprocess.run([FFMPEG, "-y", "-loglevel", "error", "-i", src,
                "-frames:v", str(max_frames), os.path.join(frames_dir, "f%05d.png")], check=True)
frames = sorted(glob.glob(os.path.join(frames_dir, "*.png")))
print(f"[stereo] {len(frames)} frames, num_disp={num_disp} block={block}", flush=True)

min_disp = -(num_disp // 2)          # movie stereo converges on the screen plane — content
sgbm = cv2.StereoSGBM_create(         # lives BOTH sides of it, so the window must go negative
    minDisparity=min_disp, numDisparities=num_disp, blockSize=block,
    P1=8 * 3 * block * block, P2=32 * 3 * block * block,
    disp12MaxDiff=1, uniquenessRatio=10, speckleWindowSize=100, speckleRange=2,
    mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY)

mid = frames[len(frames) // 2]
for i, f in enumerate(frames):
    img = cv2.imread(f)
    h, w = img.shape[:2]
    if tab:
        L, R = img[: h // 2], img[h // 2:]
    else:
        L, R = img[:, : w // 2], img[:, w // 2:]
    gl, gr = cv2.cvtColor(L, cv2.COLOR_BGR2GRAY), cv2.cvtColor(R, cv2.COLOR_BGR2GRAY)
    disp = sgbm.compute(gl, gr).astype(np.float32) / 16.0
    valid = disp > (min_disp - 0.5)               # cv2 marks unmatched pixels min_disp-1
    if valid.sum() > 100:
        lo, hi = np.percentile(disp[valid], [2, 98])
    else:
        lo, hi = 0.0, 1.0
    norm = np.zeros(disp.shape, np.uint8)
    norm[valid] = (np.clip((disp[valid] - lo) / max(1e-6, hi - lo), 0, 1) * 255).astype(np.uint8)
    cv2.imwrite(os.path.join(depth_dir, os.path.basename(f)), norm)
    if f == mid:                                  # colored point cloud of the middle frame
        ys, xs = np.where(valid & (norm > 2))
        step = max(1, len(ys) // 400_000)         # cap cloud size
        ys, xs = ys[::step], xs[::step]
        z = disp[ys, xs]
        rgb = L[ys, xs][:, ::-1]                  # BGR→RGB
        ply = os.path.join(outdir, "midframe_cloud.ply")
        with open(ply, "w") as fh:
            fh.write("ply\nformat ascii 1.0\n"
                     f"element vertex {len(ys)}\n"
                     "property float x\nproperty float y\nproperty float z\n"
                     "property uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n")
            for (yy, xx, zz, (r, g, b)) in zip(ys, xs, z, rgb):
                fh.write(f"{xx} {-yy} {zz:.1f} {r} {g} {b}\n")
        print(f"[stereo] wrote midframe_cloud.ply ({len(ys)} pts)", flush=True)
    if i % 24 == 0:
        print(f"[stereo] frame {i + 1}/{len(frames)}", flush=True)

outputs = [{"path": "midframe_cloud.ply", "meta": {"tool": "stereo_sgbm"}}]
if len(frames) > 1:                               # depth preview video
    prev = os.path.join(outdir, "depth_preview.mp4")
    subprocess.run([FFMPEG, "-y", "-loglevel", "error", "-framerate", "24",
                    "-i", os.path.join(depth_dir, "f%05d.png"),
                    "-pix_fmt", "yuv420p", prev], check=True)
    outputs.append({"path": "depth_preview.mp4", "meta": {"tool": "stereo_sgbm"}})

zpath = os.path.join(outdir, "depth_frames.zip")
with zipfile.ZipFile(zpath, "w", zipfile.ZIP_DEFLATED) as zf:
    for f in sorted(glob.glob(os.path.join(depth_dir, "*.png"))):
        zf.write(f, os.path.basename(f))
outputs.append({"path": "depth_frames.zip", "meta": {"tool": "stereo_sgbm"}})

json.dump({"outputs": outputs}, open(os.path.join(outdir, "result.json"), "w"))
shutil.rmtree(work, ignore_errors=True)
print("done:", ", ".join(o["path"] for o in outputs), flush=True)