#!/usr/bin/env python3
"""
pipeline_360.py -- End-to-end 360 photogrammetry prep pipeline for COLMAP.

Merges what used to be four separate scripts (extract_sharpest_frame.py,
generate_sfm_masks.py, project_to_cubemap.py, generate_rig_config.py) into
one tool with one project directory layout. The old multi-session merge
step (merge_frames_masks.py) is gone -- this pipeline is single-session.

Also gone: the Metashape/RealityCapture mask "convention" options that used
to be threaded through every stage. Everything downstream commits to
COLMAP's rig-based reconstruction now, so there was never a real reason for
a mask written for stage 2 to end up in a different format than stage 3
expects. One convention (COLMAP's: black=excluded, filename=<image>.png)
everywhere.

PIPELINE
--------
    video.mp4
      |  [1: extract]  sharpest frame per --extract-chunk-size window
      v
    01_frames/            equirect frames (output_frame_00001.png, ...)
      |  [2: masks]       YOLO-seg, cubemap-projected detection to fix
      |                   pole distortion, reprojected back to equirect
      v
    02_masks/              equirect masks (output_frame_00001.png.png, ...)
      |  [3: cubemap]     each equirect frame+mask -> 6 perspective faces
      v
    03_cubemap/
      images/               output_frame_00001_F.png, ..._R.png, ...
      masks/                output_frame_00001_F.png.png, ...
      |  [4: rig]         COLMAP rig config + nested {FACE}/{stem} tree
      v
    04_colmap/
      images/F/output_frame_00001.png, images/R/..., ...
      masks/F/output_frame_00001.png.png, ...
      rig_config.json

Run the whole thing:
    python pipeline_360.py all --project-dir ./project --video ./raw.mp4

Resume from a later stage (e.g. you already have frames+masks and just
want to redo the cubemap/rig stages with a different --cubemap-face-size):
    python pipeline_360.py all --project-dir ./project --start-stage cubemap

Run stages individually (mirrors the original standalone scripts, useful
for debugging a single stage in isolation):
    python pipeline_360.py extract --video raw.mp4 --output-dir 01_frames
    python pipeline_360.py masks --input-dir 01_frames --output-dir 02_masks
    python pipeline_360.py cubemap --frames-dir 01_frames --masks-dir 02_masks \\
        --output-frames-dir 03_cubemap/images --output-masks-dir 03_cubemap/masks
    python pipeline_360.py rig --images-dir 03_cubemap/images --masks-dir 03_cubemap/masks \\
        --output 04_colmap/rig_config.json --link-tree 04_colmap/images \\
        --masks-link-tree 04_colmap/masks

Self-test for the rig-config geometry/schema logic (no video/GPU needed):
    python pipeline_360.py rig --self-test

Requirements:
    pip install opencv-python numpy tqdm py360convert ultralytics --break-system-packages
    (ultralytics is only needed for the masks stage; the other three stages
    work without it.)
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import multiprocessing
import re
import struct
import sys
import zlib
from collections import defaultdict
from concurrent.futures import FIRST_COMPLETED, ProcessPoolExecutor, as_completed, wait
from math import ceil
from pathlib import Path
from time import monotonic
from typing import Callable, Dict, List, Optional, Tuple

try:
    import cv2
except ImportError:
    cv2 = None

import numpy as np

try:
    from tqdm import tqdm
except ImportError:
    tqdm = None


# ==========================================================================
# Shared: logging / progress / cancellation
#
# One shared plumbing layer used by every stage, so the GUI only has to
# understand one line protocol (plain lines = log, "[progress] "-prefixed
# lines = replace the last line) regardless of which stage is running.
# ==========================================================================

Logger = Optional[Callable[[str], None]]
CancelChecker = Optional[Callable[[], bool]]
PROGRESS_PREFIX = "[progress] "
SPINNER_FRAMES = "|/-\\"
MULTIPROCESS_PROGRESS_INTERVAL = 0.2

IMAGE_EXTENSIONS = frozenset({".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"})
FACES: Tuple[str, ...] = ("F", "R", "B", "L", "U", "D")
FACE_NAMES = {"F": "front", "R": "right", "B": "back", "L": "left", "U": "up", "D": "down"}


class PipelineError(Exception):
    pass


class PipelineCancelled(PipelineError):
    pass


class GuiTqdm(tqdm if tqdm is not None else object):
    def __init__(self, *args, logger: Logger = None, **kwargs) -> None:
        self.logger = logger
        self._last_displayed_message = ""
        super().__init__(*args, **kwargs)

    def update(self, n=1):
        displayed = super().update(n)
        if self.logger is not None:
            self.display()
        return displayed

    def display(self, msg=None, pos=None) -> None:
        if self.logger is None:
            super().display(msg=msg, pos=pos)
            return
        progress_message = msg if msg is not None else str(self)
        cleaned = progress_message.strip()
        if cleaned and cleaned != self._last_displayed_message:
            self._last_displayed_message = cleaned
            emit_progress(cleaned, self.logger)


def emit(message: str, logger: Logger = None) -> None:
    if logger is None:
        print(message)
    else:
        logger(message)


def emit_progress(message: str, logger: Logger = None) -> None:
    if logger is not None:
        logger(f"{PROGRESS_PREFIX}{message}")


def print_stage_header(index: int, total: int, title: str, logger: Logger = None) -> None:
    emit("", logger)
    emit(f"==== STAGE {index}/{total}: {title} ====", logger)


def create_progress(total: Optional[int], desc: str, logger: Logger = None):
    ensure_tqdm_available()
    kwargs = {
        "total": total, "desc": desc, "unit": "frame",
        "dynamic_ncols": True, "leave": False,
        "bar_format": "{l_bar}{bar:20}{r_bar}",
    }
    if logger is None:
        return tqdm(**kwargs)
    kwargs.update(mininterval=0, miniters=1, ascii=False)
    return GuiTqdm(logger=logger, **kwargs)


def ensure_opencv_available() -> None:
    if cv2 is None:
        raise PipelineError("opencv-python is required. Install with: pip install opencv-python --break-system-packages")


def ensure_tqdm_available() -> None:
    if tqdm is None:
        raise PipelineError("tqdm is required. Install with: pip install tqdm --break-system-packages")


def ensure_py360convert_available():
    try:
        import py360convert
    except ImportError:
        raise PipelineError("py360convert is required. Install with: pip install py360convert --break-system-packages")
    return py360convert


def raise_if_cancelled(should_cancel: CancelChecker = None) -> None:
    if should_cancel is not None and should_cancel():
        raise PipelineCancelled("Processing was cancelled.")


def gather_images(input_dir: Path) -> List[Path]:
    frames = sorted(f for f in input_dir.iterdir() if f.suffix.lower() in IMAGE_EXTENSIONS)
    if not frames:
        raise PipelineError(f"No image files found in {input_dir}")
    return frames


# ==========================================================================
# STAGE 1: extract sharpest frame per chunk from a video
# ==========================================================================
#
# Picks the sharpest (least motion-blurred) frame out of every
# --extract-chunk-size window, using variance-of-Laplacian as the sharpness
# metric. Two-pass: analyze every frame's sharpness to a CSV (parallelized
# across --extract-workers, since this is the slow part), then extract just
# the winners.

def get_frame_count(capture) -> Optional[int]:
    total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    return total if total > 0 else None


def compute_sharpness(frame, scale_width: int) -> float:
    height, width = frame.shape[:2]
    if scale_width > 0 and width > scale_width:
        ratio = scale_width / float(width)
        frame = cv2.resize(frame, (scale_width, max(1, int(height * ratio))), interpolation=cv2.INTER_AREA)
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    return float(cv2.Laplacian(gray, cv2.CV_64F).var())


def build_analysis_ranges(total_frames: int, workers: int) -> List[Tuple[int, int]]:
    task_count = max(workers * 8, workers)
    range_size = max(1, ceil(total_frames / task_count))
    ranges, start = [], 0
    while start < total_frames:
        end = min(total_frames, start + range_size)
        ranges.append((start, end))
        start = end
    return ranges


def analyze_video_range(video_path: str, start_frame: int, end_frame: int, scale_width: int) -> List[Tuple[int, float]]:
    ensure_opencv_available()
    capture = cv2.VideoCapture(video_path)
    if not capture.isOpened():
        raise PipelineError(f"Failed to open video: {video_path}")
    capture.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
    results, frame_number = [], start_frame
    while frame_number < end_frame:
        ok, frame = capture.read()
        if not ok:
            break
        results.append((frame_number, compute_sharpness(frame, scale_width)))
        frame_number += 1
    capture.release()
    return results


def analyze_video_single_process(video_file, metadata_path, scale_width, total_frames, logger=None, should_cancel=None) -> int:
    capture = cv2.VideoCapture(str(video_file))
    if not capture.isOpened():
        raise PipelineError(f"Failed to open video: {video_file}")
    frame_number = 0
    try:
        with metadata_path.open("w", newline="", encoding="utf-8") as fh:
            writer = csv.writer(fh)
            writer.writerow(["frame", "sharpness"])
            with create_progress(total=total_frames, desc="Analyze", logger=logger) as bar:
                while True:
                    raise_if_cancelled(should_cancel)
                    ok, frame = capture.read()
                    if not ok:
                        break
                    sharpness = compute_sharpness(frame, scale_width)
                    writer.writerow([frame_number, f"{sharpness:.10f}"])
                    frame_number += 1
                    bar.update(1)
    except PipelineCancelled:
        capture.release()
        if metadata_path.exists():
            metadata_path.unlink()
        raise
    capture.release()
    return frame_number


def analyze_video_multi_process(video_file, metadata_path, scale_width, total_frames, workers, logger=None, should_cancel=None) -> int:
    ranges = build_analysis_ranges(total_frames, workers)
    completed: Dict[int, List[Tuple[int, float]]] = {}
    next_index, written, spinner, last_refresh = 0, 0, 0, 0.0

    try:
        with metadata_path.open("w", newline="", encoding="utf-8") as fh:
            writer = csv.writer(fh)
            writer.writerow(["frame", "sharpness"])
            with create_progress(total=total_frames, desc="Analyze", logger=logger) as bar:
                with ProcessPoolExecutor(max_workers=workers) as executor:
                    future_map = {
                        executor.submit(analyze_video_range, str(video_file), s, e, scale_width): i
                        for i, (s, e) in enumerate(ranges)
                    }
                    bar.set_postfix_str(f"{SPINNER_FRAMES[spinner]} {len(future_map)} ranges")
                    bar.display()
                    last_refresh = monotonic()

                    while future_map:
                        raise_if_cancelled(should_cancel)
                        done, _ = wait(future_map, timeout=0.1, return_when=FIRST_COMPLETED)
                        if not done:
                            now = monotonic()
                            if now - last_refresh >= MULTIPROCESS_PROGRESS_INTERVAL:
                                spinner = (spinner + 1) % len(SPINNER_FRAMES)
                                bar.set_postfix_str(f"{SPINNER_FRAMES[spinner]} {len(future_map)} ranges")
                                bar.display()
                                last_refresh = now
                            continue
                        for fut in done:
                            idx = future_map.pop(fut)
                            completed[idx] = fut.result()
                            bar.update(len(completed[idx]))
                        spinner = (spinner + 1) % len(SPINNER_FRAMES)
                        bar.set_postfix_str(f"{SPINNER_FRAMES[spinner]} {len(future_map)} ranges")
                        bar.display()
                        last_refresh = monotonic()
                        while next_index in completed:
                            for frame_number, sharpness in completed.pop(next_index):
                                writer.writerow([frame_number, f"{sharpness:.10f}"])
                                written += 1
                            next_index += 1
                    bar.set_postfix_str("done")
                    bar.display()
    except PipelineCancelled:
        if metadata_path.exists():
            metadata_path.unlink()
        raise
    return written


def analyze_video(video_file, metadata_path, scale_width, workers=4, logger=None, should_cancel=None) -> None:
    ensure_opencv_available()
    ensure_tqdm_available()
    metadata_path.parent.mkdir(parents=True, exist_ok=True)
    emit("[1/3] Analyzing sharpness...", logger)

    probe = cv2.VideoCapture(str(video_file))
    if not probe.isOpened():
        raise PipelineError(f"Failed to open video: {video_file}")
    total_frames = get_frame_count(probe)
    probe.release()

    if workers <= 1 or total_frames is None or total_frames <= 1:
        n = analyze_video_single_process(video_file, metadata_path, scale_width, total_frames, logger, should_cancel)
    else:
        n = analyze_video_multi_process(video_file, metadata_path, scale_width, total_frames, workers, logger, should_cancel)

    if n == 0:
        raise PipelineError("No frames could be read from the video.")


def parse_best_frames(metadata_path: Path, chunk_size: int, logger: Logger = None, should_cancel: CancelChecker = None) -> List[int]:
    ensure_tqdm_available()
    frame_data: List[Tuple[int, float]] = []
    try:
        with metadata_path.open("r", encoding="utf-8") as fh:
            total_lines = max(0, sum(1 for _ in fh) - 1)
        with metadata_path.open("r", newline="", encoding="utf-8") as fh:
            reader = csv.DictReader(fh)
            with create_progress(total=total_lines or None, desc="Parse", logger=logger) as bar:
                for row in reader:
                    raise_if_cancelled(should_cancel)
                    try:
                        frame_data.append((int(row["frame"]), float(row["sharpness"])))
                    except (KeyError, TypeError, ValueError):
                        pass
                    finally:
                        bar.update(1)
    except FileNotFoundError as exc:
        raise PipelineError(f"Metadata file not found: {metadata_path}") from exc

    if not frame_data:
        raise PipelineError("No valid frame data was found in metadata.")

    best, chunk = [], []
    for item in frame_data:
        chunk.append(item)
        if len(chunk) == chunk_size:
            best.append(max(chunk, key=lambda x: x[1])[0])
            chunk = []
    if chunk:
        best.append(max(chunk, key=lambda x: x[1])[0])
    return best


def format_output_filename(pattern: str, index: int, output_format: str) -> str:
    try:
        name = pattern % index
    except (TypeError, ValueError):
        if "%d" not in pattern:
            name = pattern
        else:
            raise PipelineError("--extract-output-pattern must be a valid printf-style pattern such as output_frame_%05d.png")

    path = Path(name)
    target_suffix = f".{output_format}"
    if path.suffix.lower() in {".jpg", ".jpeg", ".png"}:
        return str(path.with_suffix(target_suffix))
    if path.suffix:
        raise PipelineError("--extract-output-pattern must end with .png or .jpg, or omit the extension entirely.")
    return f"{name}{target_suffix}"


def save_frame(frame, output_file: Path, jpeg_quality: int, output_format: str) -> None:
    output_file.parent.mkdir(parents=True, exist_ok=True)
    if output_format == "jpg":
        ok = cv2.imwrite(str(output_file), frame, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
    else:
        ok = cv2.imwrite(str(output_file), frame)
    if not ok:
        raise PipelineError(f"Failed to write image: {output_file}")


def extract_frames(video_file, frame_numbers, output_dir, output_pattern, jpeg_quality, output_format, logger=None, should_cancel=None) -> None:
    ensure_opencv_available()
    ensure_tqdm_available()
    if not frame_numbers:
        emit("No frames selected for extraction.", logger)
        return
    output_dir.mkdir(parents=True, exist_ok=True)
    targets = set(frame_numbers)

    capture = cv2.VideoCapture(str(video_file))
    if not capture.isOpened():
        raise PipelineError(f"Failed to open video: {video_file}")
    emit(f"[3/3] Extracting frames... ({len(frame_numbers)} frames)", logger)
    total = get_frame_count(capture)

    frame_number, saved, output_index = 0, 0, 1
    with create_progress(total=total, desc="Extract", logger=logger) as bar:
        while True:
            raise_if_cancelled(should_cancel)
            ok, frame = capture.read()
            if not ok:
                break
            if frame_number in targets:
                name = format_output_filename(output_pattern, output_index, output_format)
                save_frame(frame, output_dir / name, jpeg_quality, output_format)
                output_index += 1
                saved += 1
                if saved == len(targets):
                    bar.update(1)
                    break
            frame_number += 1
            bar.update(1)
    capture.release()

    if saved != len(targets):
        raise PipelineError(f"Some selected frames could not be extracted. Expected {len(targets)}, saved {saved}.")


def run_extract_stage(
    video: Path, output_dir: Path, chunk_size: int = 30, scale_width: int = 1920, workers: int = 4,
    output_pattern: str = "output_frame_%05d.png", output_format: str = "png", jpeg_quality: int = 95,
    reuse_metadata: bool = True, logger: Logger = None, should_cancel: CancelChecker = None,
) -> Path:
    """Returns output_dir (where the extracted equirect frames now live)."""
    ensure_opencv_available()
    ensure_tqdm_available()
    if not video.exists():
        raise PipelineError(f"Input video was not found: {video}")
    if chunk_size <= 0:
        raise PipelineError("--extract-chunk-size must be 1 or greater.")
    if workers <= 0:
        raise PipelineError("--extract-workers must be 1 or greater.")

    output_dir.mkdir(parents=True, exist_ok=True)
    metadata_path = output_dir / "_sharpness_metadata.csv"

    raise_if_cancelled(should_cancel)
    if metadata_path.exists() and reuse_metadata:
        emit(f"Using existing sharpness metadata: {metadata_path}", logger)
    else:
        emit(f"Generating sharpness metadata: {metadata_path}", logger)
        analyze_video(video, metadata_path, scale_width, workers, logger, should_cancel)

    raise_if_cancelled(should_cancel)
    emit("[2/3] Parsing metadata...", logger)
    best_frames = parse_best_frames(metadata_path, chunk_size, logger, should_cancel)

    extract_frames(video, best_frames, output_dir, output_pattern, jpeg_quality, output_format, logger, should_cancel)

    emit("Done: sharp frames extracted.", logger)
    emit(f"Output directory: {output_dir.resolve()}", logger)
    return output_dir


# ==========================================================================
# STAGE 2: YOLO segmentation masks (cubemap-projected detection)
# ==========================================================================
#
# Masks excluded/movable objects (people, vehicles, ...) so they don't
# pollute the SfM point cloud. Detection is run on cube-face reprojections
# rather than the raw equirect frame by default, since equirect projection
# stretches anything near the poles enough that YOLO misses it -- then
# masks are reprojected back to equirect space and merged. Output is always
# COLMAP convention: black(0)=excluded, filename is "<image filename>.png".

COCO_MOVABLE_DEFAULTS = ["person", "car", "bus", "truck", "bicycle", "motorcycle"]


def load_yolo_model(model_name: str, device: str):
    try:
        from ultralytics import YOLO
    except ImportError:
        raise PipelineError(
            "ultralytics is not installed. Run:\n    pip install ultralytics --break-system-packages\n"
            "(ROCm users: install the ROCm torch build first so this doesn't pull in a CUDA-only wheel.)"
        )
    model = YOLO(model_name)
    try:
        model.to(device)
    except Exception as exc:
        print(f"[warn] could not move model to device={device} ({exc}); falling back to cpu", file=sys.stderr)
        device = "cpu"
    return model, device


def build_mask(result, shape_hw, target_ids, dilate_px) -> Tuple[np.ndarray, int]:
    h, w = shape_hw
    mask = np.zeros((h, w), dtype=np.uint8)
    if result.masks is None or result.boxes is None or len(result.boxes) == 0:
        return mask, 0

    cls_ids = result.boxes.cls.cpu().numpy().astype(int)
    seg_masks = result.masks.data.cpu().numpy()
    hits = 0
    for i, cid in enumerate(cls_ids):
        if cid not in target_ids:
            continue
        m = seg_masks[i]
        if m.shape != (h, w):
            m = cv2.resize(m, (w, h), interpolation=cv2.INTER_NEAREST)
        mask[m > 0.5] = 255
        hits += 1

    if hits and dilate_px > 0:
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilate_px * 2 + 1, dilate_px * 2 + 1))
        mask = cv2.dilate(mask, kernel)
    return mask, hits


def _chunked(seq, size):
    for i in range(0, len(seq), size):
        yield seq[i:i + size]


def run_predict_chunked(model, images, device, conf, imgsz, tile_batch):
    all_results = []
    for chunk in _chunked(images, max(1, tile_batch)):
        results = model.predict(source=chunk, device=device, conf=conf, imgsz=imgsz, verbose=False, retina_masks=True)
        all_results.extend(results)
        if device != "cpu":
            try:
                import torch
                torch.cuda.empty_cache()
            except Exception:
                pass
    return all_results


def _process_mask_batch_equirect(batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate_px):
    results = run_predict_chunked(model, [str(p) for p in batch_paths], device, conf, imgsz, tile_batch)
    out = []
    for result in results:
        h, w = result.orig_shape
        raw_mask, hits = build_mask(result, (h, w), target_ids, dilate_px)
        out.append((raw_mask, hits, (h, w)))
    return out


def _process_mask_batch_cubemap(batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate_px, face_size, skip_faces):
    py360convert = ensure_py360convert_available()

    frame_records, flat_images, flat_index = [], [], []
    for local_idx, frame_path in enumerate(batch_paths):
        img = cv2.imread(str(frame_path))
        if img is None:
            print(f"[warn] could not read {frame_path}, skipping", file=sys.stderr)
            frame_records.append(None)
            continue
        h, w = img.shape[:2]
        faces = py360convert.e2c(img, face_w=face_size, mode="bilinear", cube_format="dict")
        frame_records.append({"h": h, "w": w, "faces": faces})
        for face_name, face_img in faces.items():
            if face_name in skip_faces:
                continue
            flat_images.append(face_img)
            flat_index.append((local_idx, face_name))

    results = run_predict_chunked(model, flat_images, device, conf, imgsz, tile_batch) if flat_images else []

    per_frame_face_masks = [dict() if rec is not None else None for rec in frame_records]
    per_frame_hits = [0] * len(frame_records)
    for (local_idx, face_name), result in zip(flat_index, results):
        face_img = frame_records[local_idx]["faces"][face_name]
        fh, fw = face_img.shape[:2]
        raw_mask, hits = build_mask(result, (fh, fw), target_ids, dilate_px)
        per_frame_face_masks[local_idx][face_name] = raw_mask
        per_frame_hits[local_idx] += hits

    out = []
    for local_idx, rec in enumerate(frame_records):
        if rec is None:
            out.append((None, 0, None))
            continue
        h, w = rec["h"], rec["w"]
        mask_dict_3ch = {}
        for face_name in FACES:
            face_shape = rec["faces"][face_name].shape[:2]
            m = per_frame_face_masks[local_idx].get(face_name)
            if m is None:
                m = np.zeros(face_shape, dtype=np.uint8)
            mask_dict_3ch[face_name] = cv2.merge([m, m, m])
        equirect_mask_3ch = py360convert.c2e(mask_dict_3ch, h=h, w=w, mode="nearest", cube_format="dict")
        raw_equirect_mask = equirect_mask_3ch[:, :, 0]
        _, raw_equirect_mask = cv2.threshold(raw_equirect_mask, 127, 255, cv2.THRESH_BINARY)
        out.append((raw_equirect_mask, per_frame_hits[local_idx], (h, w)))
    return out


def mask_path_for(frame_path: Path, output_dir: Path) -> Path:
    """COLMAP convention only: mask_path/<image filename>.png"""
    return output_dir / f"{frame_path.name}.png"


def run_masks_stage(
    input_dir: Path, output_dir: Path,
    classes: str = ",".join(COCO_MOVABLE_DEFAULTS), model_name: str = "yolo11x-seg.pt", device: str = "cuda",
    conf: float = 0.35, imgsz: int = 1280, dilate: int = 6,
    projection: str = "cubemap", face_size: int = 1024, skip_faces_raw: str = "",
    batch: int = 8, tile_batch: int = 8, empty_if_none: bool = False,
    logger: Logger = None, should_cancel: CancelChecker = None,
) -> Path:
    """Returns output_dir (where equirect-space COLMAP-convention masks now live)."""
    ensure_opencv_available()
    output_dir.mkdir(parents=True, exist_ok=True)
    frames = gather_images(input_dir)
    emit(f"Found {len(frames)} frames in {input_dir}", logger)

    if projection == "cubemap":
        ensure_py360convert_available()

    skip_faces = {f.strip().upper() for f in skip_faces_raw.split(",") if f.strip()}
    unknown = skip_faces - set(FACES)
    if unknown:
        raise PipelineError(f"--mask-skip-faces has unrecognised face(s) {sorted(unknown)}; valid: {list(FACES)}")

    model, device = load_yolo_model(model_name, device)

    wanted = {c.strip().lower() for c in classes.split(",") if c.strip()}
    name_to_id = {v.lower(): k for k, v in model.names.items()}
    unknown_classes = wanted - set(name_to_id.keys())
    if unknown_classes:
        print(f"[warn] class names not in model label set, ignored: {sorted(unknown_classes)}", file=sys.stderr)
    target_ids = {name_to_id[n] for n in wanted if n in name_to_id}
    if not target_ids:
        raise PipelineError("No valid target classes resolved. Check --mask-classes against the model's class list.")
    emit(f"Masking classes: {sorted(wanted & set(name_to_id.keys()))} -> ids {sorted(target_ids)}", logger)
    emit(f"Projection: {projection}" + (f" (skipping faces: {sorted(skip_faces)})" if skip_faces else ""), logger)

    written = skipped = total_detections = 0
    with create_progress(total=len(frames), desc="Masks", logger=logger) as bar:
        for batch_start in range(0, len(frames), batch):
            raise_if_cancelled(should_cancel)
            batch_paths = frames[batch_start:batch_start + batch]

            if projection == "cubemap":
                batch_results = _process_mask_batch_cubemap(
                    batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate, face_size, skip_faces
                )
            else:
                batch_results = _process_mask_batch_equirect(batch_paths, model, device, conf, imgsz, tile_batch, target_ids, dilate)

            for frame_path, (raw_mask, hit_count, _shape) in zip(batch_paths, batch_results):
                if raw_mask is None:
                    bar.update(1)
                    continue
                total_detections += hit_count
                out_path = mask_path_for(frame_path, output_dir)
                if hit_count == 0 and not empty_if_none:
                    skipped += 1
                    bar.update(1)
                    continue
                final_mask = 255 - raw_mask  # COLMAP: black(0) = excluded
                cv2.imwrite(str(out_path), final_mask)
                written += 1
                bar.update(1)

    emit(f"Done. {written} masks written, {skipped} frame(s) had no target objects (unmasked), "
         f"{total_detections} total target-object instances masked.", logger)
    emit(f"Masks in: {output_dir}", logger)
    return output_dir


# ==========================================================================
# STAGE 3: project equirect frames+masks to 6 cube faces
# ==========================================================================
#
# Neither COLMAP nor Brush has an equirectangular camera model, so this
# splits each equirect frame (and its matching mask, if present) into 6
# perspective cube-face crops -- the "virtual camera rig" images that
# stage 4's rig config describes.

def _cubemap_input_mask_path(frame_path: Path, masks_dir: Path) -> Path:
    return masks_dir / f"{frame_path.name}.png"


def cubemap_process_one(args_tuple):
    (frame_path_str, masks_dir_str, out_frames_dir_str, out_masks_dir_str,
     face_size, skip_faces, jpeg_quality, resume) = args_tuple

    import py360convert  # per-worker import: required for ProcessPoolExecutor pickling safety

    frame_path = Path(frame_path_str)
    out_frames_dir = Path(out_frames_dir_str)
    out_masks_dir = Path(out_masks_dir_str) if out_masks_dir_str else None
    masks_dir = Path(masks_dir_str) if masks_dir_str else None

    active_faces = [f for f in FACES if f not in skip_faces]

    if resume:
        expected = [out_frames_dir / f"{frame_path.stem}_{f}{frame_path.suffix}" for f in active_faces]
        if all(e.exists() for e in expected):
            return frame_path.name, len(active_faces), 0, "skipped (resume)"

    img = cv2.imread(str(frame_path))
    if img is None:
        return frame_path.name, 0, 0, "unreadable"
    h, w = img.shape[:2]

    faces = py360convert.e2c(img, face_w=face_size, mode="bilinear", cube_format="dict")

    mask_faces = None
    if masks_dir is not None and out_masks_dir is not None:
        mask_path = _cubemap_input_mask_path(frame_path, masks_dir)
        if mask_path.exists():
            mask_img = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
            if mask_img is not None and mask_img.shape[:2] == (h, w):
                mask_3ch = cv2.merge([mask_img, mask_img, mask_img])
                mask_faces_3ch = py360convert.e2c(mask_3ch, face_w=face_size, mode="nearest", cube_format="dict")
                mask_faces = {}
                for fname, m3 in mask_faces_3ch.items():
                    m = m3[:, :, 0]
                    _, m = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
                    mask_faces[fname] = m

    frames_written = masks_written = 0
    for face_name in active_faces:
        face_img = faces[face_name]
        out_frame_name = f"{frame_path.stem}_{face_name}{frame_path.suffix}"
        out_path = out_frames_dir / out_frame_name
        if frame_path.suffix.lower() in (".jpg", ".jpeg"):
            cv2.imwrite(str(out_path), face_img, [cv2.IMWRITE_JPEG_QUALITY, jpeg_quality])
        else:
            cv2.imwrite(str(out_path), face_img)
        frames_written += 1

        if mask_faces is not None:
            cv2.imwrite(str(out_masks_dir / f"{out_frame_name}.png"), mask_faces[face_name])
            masks_written += 1

    return frame_path.name, frames_written, masks_written, "ok"


def run_cubemap_stage(
    frames_dir: Path, output_frames_dir: Path, masks_dir: Optional[Path] = None, output_masks_dir: Optional[Path] = None,
    face_size: int = 1600, skip_faces_raw: str = "", jpeg_quality: int = 95, workers: int = 0, resume: bool = False,
    logger: Logger = None, should_cancel: CancelChecker = None,
) -> Tuple[Path, Optional[Path]]:
    """Returns (output_frames_dir, output_masks_dir_or_None)."""
    ensure_opencv_available()
    ensure_py360convert_available()

    if masks_dir is not None and output_masks_dir is None:
        raise PipelineError("masks_dir given but output_masks_dir missing")

    output_frames_dir.mkdir(parents=True, exist_ok=True)
    if output_masks_dir is not None:
        output_masks_dir.mkdir(parents=True, exist_ok=True)

    skip_faces = {f.strip().upper() for f in skip_faces_raw.split(",") if f.strip()}
    unknown = skip_faces - set(FACES)
    if unknown:
        raise PipelineError(f"--cubemap-skip-faces has unrecognised face(s) {sorted(unknown)}; valid: {list(FACES)}")

    frames = gather_images(frames_dir)
    emit(f"Found {len(frames)} equirect frames in {frames_dir}", logger)
    active_faces = [f for f in FACES if f not in skip_faces]
    emit(f"Active faces: {active_faces}", logger)
    emit(f"Expected output frames: {len(frames) * len(active_faces)}", logger)

    tasks = [
        (str(f), str(masks_dir) if masks_dir else None, str(output_frames_dir),
         str(output_masks_dir) if output_masks_dir else None, face_size, skip_faces, jpeg_quality, resume)
        for f in frames
    ]
    pool_workers = workers if workers > 0 else None

    total_frames_written = total_masks_written = unreadable = 0
    with create_progress(total=len(tasks), desc="Cubemap", logger=logger) as bar:
        with ProcessPoolExecutor(max_workers=pool_workers) as executor:
            futures = [executor.submit(cubemap_process_one, t) for t in tasks]
            for fut in as_completed(futures):
                raise_if_cancelled(should_cancel)
                name, fw, mw, status = fut.result()
                total_frames_written += fw
                total_masks_written += mw
                if status == "unreadable":
                    unreadable += 1
                    print(f"[warn] could not read {name}", file=sys.stderr)
                bar.update(1)

    emit(f"Done. {total_frames_written} face-crop images written to {output_frames_dir}", logger)
    if output_masks_dir:
        emit(f"{total_masks_written} face-crop masks written to {output_masks_dir}", logger)
    if unreadable:
        emit(f"{unreadable} frames were unreadable and skipped", logger)
    return output_frames_dir, (output_masks_dir if output_masks_dir else None)


# ==========================================================================
# STAGE 4: COLMAP rig config + nested {FACE}/{stem} tree
# ==========================================================================
#
# Declares that the 6 cube faces from one frame share a single optical
# center with fixed relative rotations, so bundle adjustment solves one
# 6-DoF rig pose per frame instead of 6 independent poses. Also relinks the
# flat "{stem}_{FACE}.ext" output from stage 3 into the nested
# "{FACE}/{stem}.ext" tree COLMAP's rig_configurator requires (it matches
# rig cameras by filename PREFIX only -- see build_colmap_rig_config()'s
# docstring for the schema this assumes and why).
#
# Conventions verified against py360convert source: xyzcube() puts
# +X=right, +Y=up, +Z=front in an endpoint-inclusive linspace grid, so
# fx=fy=(N-1)/2 (not the naive N/2) and F->R is a clean 90deg yaw. See
# FACE_AXES / compute_intrinsics below for the full derivation.

Vec3 = Tuple[float, float, float]
Mat3 = Tuple[Vec3, Vec3, Vec3]

FACE_AXES: Dict[str, Tuple[Vec3, Vec3, Vec3]] = {
    #      image right       image down       view direction
    "F": ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)),
    "R": ((0.0, 0.0, -1.0), (0.0, 1.0, 0.0), (1.0, 0.0, 0.0)),
    "B": ((-1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, -1.0)),
    "L": ((0.0, 0.0, 1.0), (0.0, 1.0, 0.0), (-1.0, 0.0, 0.0)),
    "U": ((1.0, 0.0, 0.0), (0.0, 0.0, 1.0), (0.0, -1.0, 0.0)),
    "D": ((1.0, 0.0, 0.0), (0.0, 0.0, -1.0), (0.0, 1.0, 0.0)),
}


def face_rotation_matrix(face: str) -> Mat3:
    right, down, forward = FACE_AXES[face]
    return (right, down, forward)


def mat_mul(a: Mat3, b: Mat3) -> Mat3:
    return tuple(tuple(sum(a[i][k] * b[k][j] for k in range(3)) for j in range(3)) for i in range(3))  # type: ignore


def mat_transpose(m: Mat3) -> Mat3:
    return tuple(tuple(m[j][i] for j in range(3)) for i in range(3))  # type: ignore


def mat_det(m: Mat3) -> float:
    return (m[0][0] * (m[1][1] * m[2][2] - m[1][2] * m[2][1])
            - m[0][1] * (m[1][0] * m[2][2] - m[1][2] * m[2][0])
            + m[0][2] * (m[1][0] * m[2][1] - m[1][1] * m[2][0]))


def assert_valid_rotation(m: Mat3, label: str, tol: float = 1e-9) -> None:
    det = mat_det(m)
    if abs(det - 1.0) > tol:
        raise ValueError(f"{label}: det={det!r}, expected +1 (not a proper rotation)")
    ident = mat_mul(m, mat_transpose(m))
    for i in range(3):
        for j in range(3):
            expected = 1.0 if i == j else 0.0
            if abs(ident[i][j] - expected) > tol:
                raise ValueError(f"{label}: R @ R.T is not identity ({ident!r})")


def quaternion_from_matrix(m: Mat3) -> Tuple[float, float, float, float]:
    """Shepperd's method -> COLMAP's (w, x, y, z) order, sign-canonicalised to w>=0."""
    m00, m01, m02 = m[0]
    m10, m11, m12 = m[1]
    m20, m21, m22 = m[2]
    trace = m00 + m11 + m22

    if trace > 0.0:
        s = math.sqrt(trace + 1.0) * 2.0
        w, x, y, z = 0.25 * s, (m21 - m12) / s, (m02 - m20) / s, (m10 - m01) / s
    elif m00 > m11 and m00 > m22:
        s = math.sqrt(1.0 + m00 - m11 - m22) * 2.0
        w, x, y, z = (m21 - m12) / s, 0.25 * s, (m01 + m10) / s, (m02 + m20) / s
    elif m11 > m22:
        s = math.sqrt(1.0 + m11 - m00 - m22) * 2.0
        w, x, y, z = (m02 - m20) / s, (m01 + m10) / s, 0.25 * s, (m12 + m21) / s
    else:
        s = math.sqrt(1.0 + m22 - m00 - m11) * 2.0
        w, x, y, z = (m10 - m01) / s, (m02 + m20) / s, (m12 + m21) / s, 0.25 * s

    norm = math.sqrt(w * w + x * x + y * y + z * z)
    w, x, y, z = w / norm, x / norm, y / norm, z / norm
    if w < 0.0:
        w, x, y, z = -w, -x, -y, -z
    return (w, x, y, z)


def cam_from_rig_rotation(face: str, ref_face: str) -> Tuple[float, float, float, float]:
    r_face = face_rotation_matrix(face)
    r_ref = face_rotation_matrix(ref_face)
    assert_valid_rotation(r_face, f"FACE_AXES[{face}]")
    assert_valid_rotation(r_ref, f"FACE_AXES[{ref_face}]")
    r = mat_mul(r_face, mat_transpose(r_ref))
    assert_valid_rotation(r, f"cam_from_rig[{face}<-{ref_face}]")
    return quaternion_from_matrix(r)


def rotation_angle_deg(q: Tuple[float, float, float, float]) -> float:
    w = min(1.0, max(-1.0, abs(q[0])))
    return math.degrees(2.0 * math.acos(w))


def compute_intrinsics(face_size: int, convention: str) -> Tuple[float, float, float, float]:
    """'linspace' (default, exact for py360convert): f=(N-1)/2. 'half': f=N/2."""
    if face_size < 2:
        raise ValueError(f"face_size must be >= 2, got {face_size}")
    if convention == "linspace":
        f = (face_size - 1) / 2.0
    elif convention == "half":
        f = face_size / 2.0
    else:
        raise ValueError(f"unknown intrinsics convention: {convention!r}")
    c = face_size / 2.0
    return (f, f, c, c)


FLAT_RE = re.compile(r"^(?P<stem>.+)_(?P<face>[FRBLUD])$")
MASK_FLAT_RE = re.compile(r"^(?P<stem>.+)_(?P<face>[FRBLUD])\.[A-Za-z0-9]+$")


def _iter_images(root: Path):
    for p in sorted(root.rglob("*")):
        if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS:
            yield p


def scan_cubemap_images(images_dir: Path, layout: str) -> Tuple[Dict[str, Dict[str, Path]], str]:
    """{stem: {FACE: path}}. flat='{stem}_{FACE}.ext', nested='{FACE}/{stem}.ext', auto detects."""
    if not images_dir.is_dir():
        raise PipelineError(f"--images-dir is not a directory: {images_dir}")
    if layout == "auto":
        subdirs = {p.name.upper() for p in images_dir.iterdir() if p.is_dir()}
        layout = "nested" if subdirs and subdirs <= set(FACES) else "flat"

    frames: Dict[str, Dict[str, Path]] = defaultdict(dict)
    if layout == "flat":
        for path in _iter_images(images_dir):
            m = FLAT_RE.match(path.stem)
            if m:
                frames[m.group("stem")][m.group("face")] = path
    else:
        for face in FACES:
            face_dir = images_dir / face
            if face_dir.is_dir():
                for path in _iter_images(face_dir):
                    frames[path.stem][face] = path

    if not frames:
        expected = "'{stem}_{FACE}.png'" if layout == "flat" else "'{FACE}/{stem}.png'"
        raise PipelineError(f"No cubemap face images found under {images_dir} (layout={layout}). Expected {expected}.")
    return dict(frames), layout


def scan_cubemap_masks(masks_dir: Path, layout: str) -> Tuple[Dict[str, Dict[str, Path]], str]:
    """Mirrors scan_cubemap_images(); a frame/face with no mask is expected, not an error."""
    if not masks_dir.is_dir():
        raise PipelineError(f"--masks-dir is not a directory: {masks_dir}")
    if layout == "auto":
        subdirs = {p.name.upper() for p in masks_dir.iterdir() if p.is_dir()}
        layout = "nested" if subdirs and subdirs <= set(FACES) else "flat"

    frames: Dict[str, Dict[str, Path]] = defaultdict(dict)
    if layout == "flat":
        for p in sorted(masks_dir.rglob("*")):
            if not p.is_file() or p.suffix.lower() != ".png":
                continue
            m = MASK_FLAT_RE.match(p.stem)
            if m:
                frames[m.group("stem")][m.group("face")] = p
    else:
        for face in FACES:
            face_dir = masks_dir / face
            if not face_dir.is_dir():
                continue
            for p in sorted(face_dir.rglob("*")):
                if not p.is_file() or p.suffix.lower() != ".png":
                    continue
                frames[Path(p.stem).stem][face] = p  # strip outer ".png" mask suffix only

    return dict(frames), layout


def _png_dimensions(path: Path) -> Optional[Tuple[int, int]]:
    try:
        with path.open("rb") as fh:
            head = fh.read(24)
    except OSError:
        return None
    if len(head) < 24 or head[:8] != b"\x89PNG\r\n\x1a\n" or head[12:16] != b"IHDR":
        return None
    return struct.unpack(">II", head[16:24])


def image_dimensions(path: Path) -> Optional[Tuple[int, int]]:
    dims = _png_dimensions(path)
    if dims is not None:
        return dims
    try:
        from PIL import Image
    except ImportError:
        return None
    try:
        with Image.open(path) as im:
            return im.size
    except OSError:
        return None


def detect_face_size(frames: Dict[str, Dict[str, Path]], sample_limit: int = 24) -> int:
    seen: Dict[Tuple[int, int], Path] = {}
    checked = 0
    for faces in frames.values():
        for path in faces.values():
            dims = image_dimensions(path)
            if dims is None:
                continue
            seen.setdefault(dims, path)
            checked += 1
            if checked >= sample_limit:
                break
        if checked >= sample_limit:
            break
    if not seen:
        raise PipelineError("could not read any image dimensions (non-PNG without Pillow?). Pass --rig-face-size explicitly.")
    if len(seen) > 1:
        detail = ", ".join(f"{w}x{h} (e.g. {p.name})" for (w, h), p in seen.items())
        raise PipelineError(f"inconsistent face dimensions: {detail}. Pass --rig-face-size explicitly.")
    (width, height), example = next(iter(seen.items()))
    if width != height:
        raise PipelineError(f"faces must be square, got {width}x{height} ({example.name}).")
    return width


def build_colmap_rig_config(present_faces: List[str], ref_face: str, intrinsics: Tuple[float, float, float, float],
                             camera_model: str = "PINHOLE") -> List[dict]:
    """
    SCHEMA (COLMAP >= 3.12, `colmap rig_configurator --rig_config_path`)
    Top level: JSON array holding ONE rig definition, reused for every frame
    (rig_configurator groups images by name-after-prefix into per-frame
    Frames -- you do not emit one rig per source panorama).

        [{"cameras": [
            {"image_prefix": "F/", "ref_sensor": true,
             "camera_model_name": "PINHOLE", "camera_params": [fx,fy,cx,cy]},
            {"image_prefix": "R/",
             "cam_from_rig_rotation": [w,x,y,z], "cam_from_rig_translation": [0,0,0],
             "camera_model_name": "PINHOLE", "camera_params": [fx,fy,cx,cy]}
        ]}]

    Read from colmap/doc/rigs.rst + src/colmap/scene/rig.cc:
      * image_prefix: matched via StringStartsWith(); no suffix field exists,
        which is why the nested {FACE}/{stem} layout is mandatory.
      * exactly one ref_sensor per rig (identity pose, no cam_from_rig_*).
      * cam_from_rig_rotation is W-FIRST (Eigen::Quaterniond(w,x,y,z) ctor).
      * cam_from_rig_translation all-zero: every face shares one optical center.
    Not run against a real COLMAP binary -- verify against your build if
    rig_configurator rejects this.
    """
    fx, fy, cx, cy = intrinsics
    if ref_face not in present_faces:
        raise ValueError(f"reference face {ref_face!r} not among present faces {present_faces}")

    cameras = []
    ordered = [ref_face] + [f for f in FACES if f in present_faces and f != ref_face]
    for face in ordered:
        entry: dict = {"image_prefix": f"{face}/"}
        if face == ref_face:
            entry["ref_sensor"] = True
        else:
            qw, qx, qy, qz = cam_from_rig_rotation(face, ref_face)
            entry["cam_from_rig_rotation"] = [qw, qx, qy, qz]
            entry["cam_from_rig_translation"] = [0.0, 0.0, 0.0]
        entry["camera_model_name"] = camera_model
        entry["camera_params"] = [fx, fy, cx, cy]
        cameras.append(entry)
    return [{"cameras": cameras}]


def write_rig_config(config: List[dict], output: Path) -> None:
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")


def materialise_nested_tree(frames: Dict[str, Dict[str, Path]], dest: Path, mode: str = "symlink") -> int:
    """Rebuild flat '{stem}_{FACE}.ext' as '{FACE}/{stem}.ext' under dest (COLMAP matches by prefix only)."""
    count = 0
    for stem, faces in sorted(frames.items()):
        for face, src in sorted(faces.items()):
            face_dir = dest / face
            face_dir.mkdir(parents=True, exist_ok=True)
            target = face_dir / f"{stem}{src.suffix}"
            _link(src, target, mode)
            count += 1
    return count


def materialise_mask_tree(mask_frames: Dict[str, Dict[str, Path]], image_frames: Dict[str, Dict[str, Path]],
                           dest: Path, mode: str = "symlink") -> Tuple[int, int]:
    """Relink masks to '{FACE}/{stem}{image_ext}.png', mirroring materialise_nested_tree()'s image tree.
    Returns (linked, orphans) -- an orphan mask has no matching image (stem/face drift)."""
    linked = orphans = 0
    for stem, faces in sorted(mask_frames.items()):
        for face, mask_src in sorted(faces.items()):
            image_src = image_frames.get(stem, {}).get(face)
            mask_name = f"{stem}{image_src.suffix}.png" if image_src is not None else f"{stem}.png.png"
            if image_src is None:
                orphans += 1
            face_dir = dest / face
            face_dir.mkdir(parents=True, exist_ok=True)
            _link(mask_src, face_dir / mask_name, mode)
            linked += 1
    return linked, orphans


def _link(src: Path, target: Path, mode: str) -> None:
    if target.exists() or target.is_symlink():
        target.unlink()
    resolved = src.resolve()
    if mode == "symlink":
        target.symlink_to(resolved)
    elif mode == "hardlink":
        target.hardlink_to(resolved)
    elif mode == "copy":
        target.write_bytes(resolved.read_bytes())
    else:
        raise ValueError(f"unknown link mode: {mode!r}")


def run_rig_stage(
    images_dir: Path, output: Path, masks_dir: Optional[Path] = None,
    face_size: Optional[int] = None, layout: str = "auto", ref_face: str = "F",
    intrinsics_convention: str = "linspace", link_tree: Optional[Path] = None,
    masks_link_tree: Optional[Path] = None, link_mode: str = "symlink",
    logger: Logger = None,
) -> Tuple[Path, Optional[Path]]:
    """Returns (link_tree_or_images_dir, masks_link_tree_or_None) -- the paths to hand COLMAP."""
    if masks_link_tree is not None and masks_dir is None:
        raise PipelineError("masks_link_tree requires masks_dir")

    frames, resolved_layout = scan_cubemap_images(images_dir, layout)
    present_faces = sorted({f for faces in frames.values() for f in faces}, key=FACES.index)
    missing_faces = [f for f in FACES if f not in present_faces]

    resolved_face_size = face_size or detect_face_size(frames)
    intrinsics = compute_intrinsics(resolved_face_size, intrinsics_convention)

    if ref_face not in present_faces:
        emit(f"warning: --rig-ref-face {ref_face} not present; using {present_faces[0]} as rig frame.", logger)
        ref_face = present_faces[0]

    partial = {s: sorted(f) for s, f in frames.items() if set(f) != set(present_faces)}

    config = build_colmap_rig_config(present_faces, ref_face, intrinsics)
    write_rig_config(config, output)

    fx, fy, cx, cy = intrinsics
    emit(f"scanned      : {images_dir}  (layout={resolved_layout})", logger)
    emit(f"frames       : {len(frames)}", logger)
    emit(f"faces present: {','.join(present_faces)}" + (f"   absent: {','.join(missing_faces)}" if missing_faces else ""), logger)
    emit(f"face size    : {resolved_face_size}px" + ("  (auto-detected)" if face_size is None else "  (override)"), logger)
    emit(f"intrinsics   : PINHOLE fx={fx:.6g} fy={fy:.6g} cx={cx:.6g} cy={cy:.6g}  [{intrinsics_convention}]", logger)
    emit(f"rig frame    : {ref_face} ({FACE_NAMES[ref_face]}) -- 1 rig definition, {len(frames)} frames, {len(present_faces)} cameras", logger)
    if partial:
        shown = list(partial.items())[:3]
        emit(f"partial frames: {len(partial)} frame(s) missing >=1 face, e.g. " + "; ".join(f"{s}=[{','.join(f)}]" for s, f in shown), logger)
    emit(f"wrote        : {output}", logger)

    final_images_path = images_dir
    if link_tree is not None:
        n = materialise_nested_tree(frames, link_tree, link_mode)
        emit(f"linked       : {n} images into {link_tree} ({link_mode})", logger)
        final_images_path = link_tree
    elif resolved_layout == "flat":
        emit("", logger)
        emit("IMPORTANT: COLMAP matches rig cameras by filename PREFIX only, so this flat", logger)
        emit("  '{stem}_{FACE}.png' naming cannot be used directly with --rig-link-tree unset.", logger)

    final_masks_path = None
    if masks_dir is not None:
        mask_frames, mask_layout = scan_cubemap_masks(masks_dir, "auto")
        n_masks = sum(len(f) for f in mask_frames.values())
        n_images = sum(len(f) for f in frames.values())
        emit("", logger)
        emit(f"masks scanned: {masks_dir}  (layout={mask_layout})", logger)
        emit(f"masks found  : {n_masks} across {len(mask_frames)} frame(s) (of {n_images} face images -- missing ones are unmasked)", logger)

        unknown_stems = set(mask_frames) - set(frames)
        if unknown_stems:
            emit(f"warning: {len(unknown_stems)} mask frame(s) have no matching image frame (stem drift?), e.g. "
                 + ", ".join(sorted(unknown_stems)[:3]), logger)

        if masks_link_tree is not None:
            linked, orphans = materialise_mask_tree(mask_frames, frames, masks_link_tree, link_mode)
            emit(f"masks linked : {linked} into {masks_link_tree} ({link_mode})", logger)
            if orphans:
                emit(f"warning: {orphans} mask(s) linked with a fallback name, no matching image found -- verify by hand.", logger)
            final_masks_path = masks_link_tree
        elif mask_layout == "flat":
            emit("", logger)
            emit("IMPORTANT: masks must mirror the nested image tree too. Set --rig-masks-link-tree.", logger)
            final_masks_path = masks_dir

    return final_images_path, final_masks_path


# ==========================================================================
# Full pipeline orchestration ("all")
# ==========================================================================

STAGE_ORDER = ["extract", "masks", "cubemap", "rig"]


def run_pipeline(args, logger: Logger = None, should_cancel: CancelChecker = None) -> None:
    project_dir: Path = args.project_dir
    frames_dir = project_dir / "01_frames"
    masks_dir = project_dir / "02_masks"
    cubemap_images_dir = project_dir / "03_cubemap" / "images"
    cubemap_masks_dir = project_dir / "03_cubemap" / "masks"
    colmap_images_dir = project_dir / "04_colmap" / "images"
    colmap_masks_dir = project_dir / "04_colmap" / "masks"
    rig_config_path = project_dir / "04_colmap" / "rig_config.json"

    start_idx = STAGE_ORDER.index(args.start_stage)
    total_stages = 4

    if start_idx <= STAGE_ORDER.index("extract"):
        if not args.video:
            raise PipelineError("--video is required when starting from the extract stage.")
        print_stage_header(1, total_stages, "Extract Sharpest Frames", logger)
        run_extract_stage(
            video=args.video, output_dir=frames_dir,
            chunk_size=args.extract_chunk_size, scale_width=args.extract_scale_width, workers=args.extract_workers,
            output_pattern=args.extract_output_pattern, output_format=args.extract_output_format,
            jpeg_quality=args.extract_jpeg_quality, reuse_metadata=args.extract_reuse_metadata,
            logger=logger, should_cancel=should_cancel,
        )
    else:
        emit(f"Skipping extract stage; using existing frames in {frames_dir}", logger)
    raise_if_cancelled(should_cancel)

    have_masks = False
    if args.skip_masks:
        emit("Skipping masks stage (--skip-masks).", logger)
    elif start_idx <= STAGE_ORDER.index("masks"):
        print_stage_header(2, total_stages, "Generate SfM Masks", logger)
        run_masks_stage(
            input_dir=frames_dir, output_dir=masks_dir,
            classes=args.mask_classes, model_name=args.mask_model, device=args.mask_device,
            conf=args.mask_conf, imgsz=args.mask_imgsz, dilate=args.mask_dilate,
            projection=args.mask_projection, face_size=args.mask_face_size, skip_faces_raw=args.mask_skip_faces,
            batch=args.mask_batch, tile_batch=args.mask_tile_batch, empty_if_none=args.mask_empty_if_none,
            logger=logger, should_cancel=should_cancel,
        )
        have_masks = True
    else:
        have_masks = masks_dir.is_dir() and any(masks_dir.glob("*.png"))
        emit(f"Skipping masks stage; {'using' if have_masks else 'no'} existing masks in {masks_dir}", logger)
    raise_if_cancelled(should_cancel)

    print_stage_header(3, total_stages, "Project to Cubemap", logger)
    run_cubemap_stage(
        frames_dir=frames_dir, output_frames_dir=cubemap_images_dir,
        masks_dir=(masks_dir if have_masks else None), output_masks_dir=(cubemap_masks_dir if have_masks else None),
        face_size=args.cubemap_face_size, skip_faces_raw=args.cubemap_skip_faces,
        jpeg_quality=args.cubemap_jpeg_quality, workers=args.cubemap_workers, resume=args.cubemap_resume,
        logger=logger, should_cancel=should_cancel,
    )
    raise_if_cancelled(should_cancel)

    have_cubemap_masks = have_masks and cubemap_masks_dir.is_dir() and any(cubemap_masks_dir.rglob("*.png"))

    print_stage_header(4, total_stages, "Generate Rig Config", logger)
    run_rig_stage(
        images_dir=cubemap_images_dir, output=rig_config_path,
        masks_dir=(cubemap_masks_dir if have_cubemap_masks else None),
        ref_face=args.rig_ref_face, intrinsics_convention=args.rig_intrinsics_convention,
        link_tree=colmap_images_dir, masks_link_tree=(colmap_masks_dir if have_cubemap_masks else None),
        link_mode=args.rig_link_mode, logger=logger,
    )

    emit("", logger)
    emit("==== PIPELINE COMPLETE ====", logger)
    emit(f"COLMAP images : {colmap_images_dir}", logger)
    if have_cubemap_masks:
        emit(f"COLMAP masks  : {colmap_masks_dir}", logger)
    emit(f"Rig config    : {rig_config_path}", logger)
    emit("", logger)
    emit("Next, run COLMAP:", logger)
    mask_flag = f" --ImageReader.mask_path {colmap_masks_dir}" if have_cubemap_masks else ""
    emit(f"  colmap feature_extractor --database_path db.db --image_path {colmap_images_dir}"
         f"{mask_flag} --ImageReader.single_camera_per_folder 1", logger)
    emit(f"  colmap rig_configurator --database_path db.db --rig_config_path {rig_config_path}", logger)
    emit("  colmap exhaustive_matcher --database_path db.db   # or vocab_tree_matcher at scale", logger)
    emit(f"  colmap mapper --database_path db.db --image_path {colmap_images_dir} --output_path {project_dir / '05_sparse'}"
         f" --Mapper.ba_refine_sensor_from_rig 0 --Mapper.ba_refine_focal_length 0 --Mapper.ba_refine_extra_params 0", logger)


# ==========================================================================
# CLI
# ==========================================================================

def _add_extract_args(p: argparse.ArgumentParser, prefix: str = "") -> None:
    d = f"--{prefix}" if prefix else "--"
    p.add_argument(f"{d}chunk-size", dest=f"{prefix}chunk_size".replace("-", "_"), type=int, default=30,
                    help="Select one sharpest frame per this many frames (default: 30)")
    p.add_argument(f"{d}scale-width", dest=f"{prefix}scale_width".replace("-", "_"), type=int, default=1920,
                    help="Resize width for sharpness analysis; 0 = original size (default: 1920)")
    p.add_argument(f"{d}workers", dest=f"{prefix}workers".replace("-", "_"), type=int, default=4,
                    help="Worker processes for sharpness analysis (default: 4)")
    p.add_argument(f"{d}output-pattern", dest=f"{prefix}output_pattern".replace("-", "_"), default="output_frame_%05d.png",
                    help="printf-style output filename pattern")
    p.add_argument(f"{d}output-format", dest=f"{prefix}output_format".replace("-", "_"), choices=("png", "jpg"), default="png")
    p.add_argument(f"{d}jpeg-quality", dest=f"{prefix}jpeg_quality".replace("-", "_"), type=int, default=95)
    p.add_argument(f"{d}regenerate-metadata", dest=f"{prefix}reuse_metadata".replace("-", "_"), action="store_false", default=True,
                    help="Ignore existing _sharpness_metadata.csv and regenerate it")


def _add_masks_args(p: argparse.ArgumentParser, prefix: str) -> None:
    d = f"--{prefix}"
    p.add_argument(f"{d}classes", dest=f"{prefix.replace('-', '_')}classes", default=",".join(COCO_MOVABLE_DEFAULTS))
    p.add_argument(f"{d}model", dest=f"{prefix.replace('-', '_')}model", default="yolo11x-seg.pt")
    p.add_argument(f"{d}device", dest=f"{prefix.replace('-', '_')}device", default="cuda")
    p.add_argument(f"{d}conf", dest=f"{prefix.replace('-', '_')}conf", type=float, default=0.35)
    p.add_argument(f"{d}imgsz", dest=f"{prefix.replace('-', '_')}imgsz", type=int, default=1280)
    p.add_argument(f"{d}dilate", dest=f"{prefix.replace('-', '_')}dilate", type=int, default=6)
    p.add_argument(f"{d}projection", dest=f"{prefix.replace('-', '_')}projection", choices=("cubemap", "equirect"), default="cubemap")
    p.add_argument(f"{d}face-size", dest=f"{prefix.replace('-', '_')}face_size", type=int, default=1024)
    p.add_argument(f"{d}skip-faces", dest=f"{prefix.replace('-', '_')}skip_faces", default="")
    p.add_argument(f"{d}batch", dest=f"{prefix.replace('-', '_')}batch", type=int, default=8)
    p.add_argument(f"{d}tile-batch", dest=f"{prefix.replace('-', '_')}tile_batch", type=int, default=8)
    p.add_argument(f"{d}empty-if-none", dest=f"{prefix.replace('-', '_')}empty_if_none", action="store_true")


def _add_cubemap_args(p: argparse.ArgumentParser, prefix: str) -> None:
    d = f"--{prefix}"
    r = lambda name: f"{prefix.replace('-', '_')}{name}"  # noqa: E731
    p.add_argument(f"{d}face-size", dest=r("face_size"), type=int, default=1600)
    p.add_argument(f"{d}skip-faces", dest=r("skip_faces"), default="")
    p.add_argument(f"{d}jpeg-quality", dest=r("jpeg_quality"), type=int, default=95)
    p.add_argument(f"{d}workers", dest=r("workers"), type=int, default=0, help="0 = all CPU cores")
    p.add_argument(f"{d}resume", dest=r("resume"), action="store_true")


def _add_rig_args(p: argparse.ArgumentParser, prefix: str) -> None:
    d = f"--{prefix}"
    r = lambda name: f"{prefix.replace('-', '_')}{name}"  # noqa: E731
    p.add_argument(f"{d}ref-face", dest=r("ref_face"), default="F", choices=list(FACES))
    p.add_argument(f"{d}intrinsics-convention", dest=r("intrinsics_convention"), choices=("linspace", "half"), default="linspace")
    p.add_argument(f"{d}link-mode", dest=r("link_mode"), choices=("symlink", "hardlink", "copy"), default="symlink")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="pipeline_360.py", description=__doc__,
                                      formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--gui-log-output", action="store_true", help=argparse.SUPPRESS)
    sub = parser.add_subparsers(dest="command", required=True)

    p_all = sub.add_parser("all", help="Run the full pipeline in one project directory.")
    p_all.add_argument("--project-dir", required=True, type=Path)
    p_all.add_argument("--video", type=Path, default=None, help="Required unless --start-stage skips extraction.")
    p_all.add_argument("--start-stage", choices=STAGE_ORDER, default="extract")
    p_all.add_argument("--skip-masks", action="store_true", help="Skip YOLO masking entirely.")
    _add_extract_args(p_all, prefix="extract-")
    _add_masks_args(p_all, prefix="mask-")
    _add_cubemap_args(p_all, prefix="cubemap-")
    _add_rig_args(p_all, prefix="rig-")

    p_extract = sub.add_parser("extract", help="Stage 1 only: extract sharpest frames from a video.")
    p_extract.add_argument("--video", required=True, type=Path)
    p_extract.add_argument("--output-dir", required=True, type=Path)
    _add_extract_args(p_extract)

    p_masks = sub.add_parser("masks", help="Stage 2 only: generate YOLO SfM masks.")
    p_masks.add_argument("--input-dir", required=True, type=Path)
    p_masks.add_argument("--output-dir", required=True, type=Path)
    _add_masks_args(p_masks, prefix="")

    p_cubemap = sub.add_parser("cubemap", help="Stage 3 only: project equirect frames+masks to cube faces.")
    p_cubemap.add_argument("--frames-dir", required=True, type=Path)
    p_cubemap.add_argument("--masks-dir", type=Path, default=None)
    p_cubemap.add_argument("--output-frames-dir", required=True, type=Path)
    p_cubemap.add_argument("--output-masks-dir", type=Path, default=None)
    _add_cubemap_args(p_cubemap, prefix="")

    p_rig = sub.add_parser("rig", help="Stage 4 only: generate the COLMAP rig config.")
    p_rig.add_argument("--images-dir", type=Path, default=None)
    p_rig.add_argument("--masks-dir", type=Path, default=None)
    p_rig.add_argument("--output", type=Path, default=Path("rig_config.json"))
    p_rig.add_argument("--face-size", type=int, default=None)
    p_rig.add_argument("--layout", choices=("auto", "flat", "nested"), default="auto")
    p_rig.add_argument("--link-tree", type=Path, default=None)
    p_rig.add_argument("--masks-link-tree", type=Path, default=None)
    p_rig.add_argument("--self-test", action="store_true")
    _add_rig_args(p_rig, prefix="")

    return parser


def _cli_logger(gui_log_output: bool) -> Logger:
    if not gui_log_output:
        return None
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    return lambda message: print(message, flush=True)


def main(argv: Optional[List[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    logger = _cli_logger(args.gui_log_output)

    try:
        if args.command == "all":
            args.project_dir.mkdir(parents=True, exist_ok=True)
            run_pipeline(args, logger=logger)

        elif args.command == "extract":
            run_extract_stage(
                video=args.video, output_dir=args.output_dir, chunk_size=args.chunk_size,
                scale_width=args.scale_width, workers=args.workers, output_pattern=args.output_pattern,
                output_format=args.output_format, jpeg_quality=args.jpeg_quality,
                reuse_metadata=args.reuse_metadata, logger=logger,
            )

        elif args.command == "masks":
            run_masks_stage(
                input_dir=args.input_dir, output_dir=args.output_dir, classes=args.classes, model_name=args.model,
                device=args.device, conf=args.conf, imgsz=args.imgsz, dilate=args.dilate, projection=args.projection,
                face_size=args.face_size, skip_faces_raw=args.skip_faces, batch=args.batch, tile_batch=args.tile_batch,
                empty_if_none=args.empty_if_none, logger=logger,
            )

        elif args.command == "cubemap":
            if args.masks_dir is not None and args.output_masks_dir is None:
                parser.error("--masks-dir given but --output-masks-dir missing")
            run_cubemap_stage(
                frames_dir=args.frames_dir, output_frames_dir=args.output_frames_dir, masks_dir=args.masks_dir,
                output_masks_dir=args.output_masks_dir, face_size=args.face_size, skip_faces_raw=args.skip_faces,
                jpeg_quality=args.jpeg_quality, workers=args.workers, resume=args.resume, logger=logger,
            )

        elif args.command == "rig":
            if args.self_test:
                return run_self_test()
            if args.images_dir is None:
                parser.error("--images-dir is required (or use --self-test)")
            if args.masks_link_tree is not None and args.masks_dir is None:
                parser.error("--masks-link-tree requires --masks-dir")
            run_rig_stage(
                images_dir=args.images_dir, output=args.output, masks_dir=args.masks_dir, face_size=args.face_size,
                layout=args.layout, ref_face=args.ref_face, intrinsics_convention=args.intrinsics_convention,
                link_tree=args.link_tree, masks_link_tree=args.masks_link_tree, link_mode=args.link_mode, logger=logger,
            )
    except PipelineCancelled:
        emit("Cancelled.", logger)
        return 130
    except PipelineError as exc:
        emit(f"Error: {exc}", logger)
        return 1
    return 0


# ==========================================================================
# Self-test (stage 4 geometry/schema only -- no video/GPU required)
# ==========================================================================

def _write_dummy_png(path: Path, width: int, height: int) -> None:
    raw = b"".join(b"\x00" + bytes(width) for _ in range(height))

    def chunk(tag: bytes, data: bytes) -> bytes:
        body = tag + data
        return struct.pack(">I", len(data)) + body + struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)

    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 0, 0, 0, 0))
                      + chunk(b"IDAT", zlib.compress(raw)) + chunk(b"IEND", b""))


def _apply(m: Mat3, v: Vec3) -> Vec3:
    return tuple(sum(m[i][k] * v[k] for k in range(3)) for i in range(3))  # type: ignore


def run_self_test() -> int:
    import tempfile
    n, ok = 64, 0

    for face in FACES:
        assert_valid_rotation(face_rotation_matrix(face), f"FACE_AXES[{face}]")
    ok += 1

    q_r = cam_from_rig_rotation("R", "F")
    expected_r = (math.sqrt(2) / 2, 0.0, -math.sqrt(2) / 2, 0.0)
    assert all(abs(a - b) < 1e-12 for a, b in zip(q_r, expected_r)), q_r
    assert abs(rotation_angle_deg(q_r) - 90.0) < 1e-9
    ok += 1

    r_mat = face_rotation_matrix("R")
    assert max(abs(a - b) for a, b in zip(_apply(mat_transpose(r_mat), (0, 0, 1)), (1, 0, 0))) < 1e-12
    for face, fwd in (("B", (0, 0, -1)), ("L", (-1, 0, 0)), ("U", (0, -1, 0)), ("D", (0, 1, 0))):
        got = _apply(mat_transpose(face_rotation_matrix(face)), (0, 0, 1))
        assert max(abs(a - b) for a, b in zip(got, fwd)) < 1e-12, (face, got)
    ok += 1

    fx, fy, cx, cy = compute_intrinsics(n, "linspace")
    assert (fx, fy, cx, cy) == ((n - 1) / 2, (n - 1) / 2, n / 2, n / 2)
    ok += 1

    with tempfile.TemporaryDirectory() as tmp:
        root = Path(tmp)
        flat, masks_flat = root / "images", root / "masks"

        for i in (1, 2):
            for face in FACES:
                _write_dummy_png(flat / f"frame_{i:05d}_{face}.png", n, n)
        for face in ("F", "R", "B", "L"):
            _write_dummy_png(flat / f"frame_{3:05d}_{face}.png", n, n)
        for face in FACES:
            _write_dummy_png(masks_flat / f"frame_{1:05d}_{face}.png.png", n, n)
        for face in ("F", "R"):
            _write_dummy_png(masks_flat / f"frame_{3:05d}_{face}.png.png", n, n)

        out, link_tree, masks_link_tree = root / "rig.json", root / "nested", root / "masks_nested"
        run_rig_stage(images_dir=flat, output=out, masks_dir=masks_flat, link_tree=link_tree,
                      masks_link_tree=masks_link_tree, link_mode="copy")
        config = json.loads(out.read_text())
        cams = config[0]["cameras"]
        assert len(cams) == 6
        refs = [c for c in cams if c.get("ref_sensor")]
        assert len(refs) == 1 and refs[0]["image_prefix"] == "F/"
        by_prefix = {c["image_prefix"]: c for c in cams}
        got_r = by_prefix["R/"]["cam_from_rig_rotation"]
        assert all(abs(a - b) < 1e-12 for a, b in zip(got_r, expected_r))
        assert (link_tree / "F" / "frame_00001.png").is_file()
        assert (masks_link_tree / "F" / "frame_00001.png.png").is_file()
        assert not (masks_link_tree / "F" / "frame_00002.png.png").exists()
        assert (masks_link_tree / "R" / "frame_00003.png.png").is_file()
        ok += 1

        partial_dir = root / "partial"
        for i in (1, 2):
            for face in ("F", "R", "B", "L"):
                _write_dummy_png(partial_dir / f"frame_{i:05d}_{face}.png", n, n)
        out2 = root / "rig_partial.json"
        run_rig_stage(images_dir=partial_dir, output=out2)
        cfg2 = json.loads(out2.read_text())
        assert len(cfg2[0]["cameras"]) == 4
        ok += 1

    print(f"\nself-test: {ok} check groups passed.")
    return 0


if __name__ == "__main__":
    multiprocessing.freeze_support()
    raise SystemExit(main())