neural_codec/codec_tools/codec_patch_gop/video_processor.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Video processing core - contains process_one_video and related functions."""

import os
import json
import math
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import cv2

# Optional cv_reader
try:
    from cv_reader import api as cv_api  # type: ignore
    HAS_CV_READER = True
except Exception:
    cv_api = None
    HAS_CV_READER = False

from .utils import ensure_dir, format_timestamp_ss, smart_resize, _clamp_int, _round_to_multiple
from .frame_utils import (
    frame_is_bad,
    pad_to_multiple_of_bgr,
    pad_to_multiple_of_32_bgr,
    bgr_to_residual_y_u8,
    decode_frame_bgr_at,
    decode_frames_bgr,
    detect_letterbox_bbox_bgr,
    _resize_bgr,
    _resize_gray,
    _resize_mv_and_scale,
    _bgr_to_luma_u8,
)
from .patch_utils import (
    pack_patches_to_canvases,
    save_canvases_as_jpg,
    extract_patch_rgb,
    block_to_4_patches,
    iter_blocks_in_raster,
)
from .video_probe import (
    get_total_frames_fps,
    ffprobe_video_codec_name,
    ffprobe_keyframe_frame_ids,
    ffprobe_sum_pkt_size,
    auto_max_total_patches,
)
from .energy_sampling import (
    pick_peak_frame_ids_from_pkt_size,
    build_variable_length_gops_by_energy,
    sample_frame_ids_by_energy_cdf,
    pick_windows_by_energy,
    allocate_frames_across_windows,
    enforce_time_coverage_frame_ids,
)
from .scoring import (
    mv_res_score_map,
    score_map_to_patch_scores,
    patch_scores_to_block_scores,
)


def cv_reader_fetch_bitcost(video_path: str, frame_ids: List[int]) -> List[Dict[str, Any]]:
    """Fetch bit-cost maps aligned with frame_ids (supports duplicates)."""
    if not HAS_CV_READER or not hasattr(cv_api, "read_video_cb"):
        raise RuntimeError("cv_reader.read_video_cb not available")

    frame_ids = [int(x) for x in frame_ids]
    pos_map: Dict[int, List[int]] = {}
    for i, fid in enumerate(frame_ids):
        pos_map.setdefault(fid, []).append(i)

    out: List[Optional[Dict[str, Any]]] = [None] * len(frame_ids)
    max_fid = max(frame_ids)

    def all_done() -> bool:
        for q in pos_map.values():
            if q:
                return False
        return True

    def cb(d: Dict[str, Any]):
        idx = int(d.get("frame_idx", -1))
        if idx in pos_map and pos_map[idx]:
            j = pos_map[idx].pop(0)
            out_item: Dict[str, Any] = {"frame_idx": idx}
            if "sub_mb_bit_cost" in d:
                out_item["sub_mb_bit_cost"] = np.asarray(d["sub_mb_bit_cost"])
            if "mb_bit_cost" in d:
                out_item["mb_bit_cost"] = np.asarray(d["mb_bit_cost"])
            if "ctu_bit_cost" in d:
                out_item["ctu_bit_cost"] = np.asarray(d["ctu_bit_cost"])
            if "pict_type" in d:
                out_item["pict_type"] = d["pict_type"]
            out[j] = out_item
        return (not all_done())

    cv_api.read_video_cb(str(video_path), cb, 1, int(max_fid) + 1, frame_ids, -1, 0, 0, 1)

    last = None
    for i in range(len(out)):
        if out[i] is None:
            if last is not None:
                out[i] = last
        else:
            last = out[i]
    if last is not None:
        for i in range(len(out)):
            if out[i] is None:
                out[i] = last

    return [x for x in out if x is not None]  # type: ignore


def bitcost_item_to_score_map(item: Dict[str, Any], out_h: int, out_w: int,
                              grid: str = "sub", pct: float = 99.0,
                              log_scale: bool = True,
                              codec_name: Optional[str] = None) -> np.ndarray:
    """Convert bitcost item to score map.
    
    Args:
        item: BitCost data from cv_reader
        out_h, out_w: Output dimensions
        grid: "sub", "mb", "ctu", or "adaptive"
        pct: Percentile for normalization
        log_scale: Apply log1p before normalization
        codec_name: "h264" or "hevc" for adaptive grid selection
    """
    # Handle adaptive grid selection based on codec
    effective_grid = grid
    if grid == "adaptive" or grid == "auto":
        codec = str(codec_name).lower().strip() if codec_name else ""
        if codec == "hevc" or codec == "h265":
            # H.265/HEVC: use CTU (64x64) as "macroblock" equivalent
            effective_grid = "ctu"
        else:
            # H.264/AVC: use MB (16x16)
            effective_grid = "mb"
    
    key_map = {
        "sub": "sub_mb_bit_cost",
        "mb": "mb_bit_cost",
        "ctu": "ctu_bit_cost",
    }
    key = key_map.get(effective_grid, "sub_mb_bit_cost")
    
    if key not in item:
        # Fallback: try available keys in order of preference
        for fallback_key in ["sub_mb_bit_cost", "mb_bit_cost", "ctu_bit_cost"]:
            if fallback_key in item:
                key = fallback_key
                break
        else:
            raise RuntimeError(f"cv_reader result missing bitcost data. Tried: {list(item.keys())}")

    arr = np.asarray(item[key], dtype=np.float32)
    if arr.ndim != 2:
        raise RuntimeError(f"Unexpected {key} shape: {arr.shape}")
    arr = np.maximum(arr, 0.0)
    if log_scale:
        arr = np.log1p(arr)
    s = float(np.percentile(arr, pct))
    s = max(s, 1e-6)
    arr = np.clip(arr / s, 0.0, 1.0).astype(np.float32)
    return cv2.resize(arr, (int(out_w), int(out_h)), interpolation=cv2.INTER_NEAREST)

def cv_reader_fetch_mvres(video_path: str, frame_ids: List[int]) -> List[Dict[str, Any]]:
    """
    Fetch mv + residual_y aligned with frame_ids (supports duplicates).
    Requires cv_reader.read_video_cb.
    """
    if not HAS_CV_READER or not hasattr(cv_api, "read_video_cb"):
        raise RuntimeError("cv_reader.read_video_cb not available")

    frame_ids = [int(x) for x in frame_ids]
    pos_map: Dict[int, List[int]] = {}
    for i, fid in enumerate(frame_ids):
        pos_map.setdefault(fid, []).append(i)

    out: List[Optional[Dict[str, Any]]] = [None] * len(frame_ids)
    max_fid = max(frame_ids)

    def all_done() -> bool:
        for q in pos_map.values():
            if q:
                return False
        return True

    def cb(d: Dict[str, Any]):
        idx = int(d.get("frame_idx", -1))
        if idx in pos_map and pos_map[idx]:
            j = pos_map[idx].pop(0)
            mv = np.asarray(d["motion_vector"])
            me = np.asarray(d["motion_energy"]) if "motion_energy" in d else None
            mem = np.asarray(d["motion_energy_median"]) if "motion_energy_median" in d else None
            # prefer residual_y; fallback residual
            if "residual_y" in d:
                ry = np.asarray(d["residual_y"])
            else:
                ry = np.asarray(d["residual"])
                if ry.ndim == 3:
                    ry = cv2.cvtColor(ry, cv2.COLOR_BGR2YUV)[:, :, 0]
            out_item: Dict[str, Any] = {"frame_idx": idx, "motion_vector": mv, "residual_y": ry}
            if me is not None:
                out_item["motion_energy"] = me
            if mem is not None:
                out_item["motion_energy_median"] = mem
            # Add pict_type if available
            if "pict_type" in d:
                out_item["pict_type"] = d["pict_type"]
            out[j] = out_item
        return (not all_done())

    # without_residual=0 means WITH residual (your earlier convention)
    without_residual = 0
    max_frames = int(max_fid) + 1
    # residual_rgb=0: keep residual_y only, skip expensive RGB residual conversion.
    cv_api.read_video_cb(str(video_path), cb, int(without_residual), int(max_frames), frame_ids, -1, 0, 0)

    # fill missing with nearest previous
    last = None
    for i in range(len(out)):
        if out[i] is None:
            if last is not None:
                out[i] = last
        else:
            last = out[i]
    if last is not None:
        for i in range(len(out)):
            if out[i] is None:
                out[i] = last

    return [x for x in out if x is not None]  # type: ignore


# -----------------------------
# main per-video worker
# -----------------------------
def process_one_video(
    video_path: str,
    key: str,
    out_root: str,
    seq_len: int,
    window_len_frames: int,
    num_candidates: int,
    top_k_windows: int,
    max_total_patches: int,
    patch: int,
    mv_unit_div: float,
    mv_pct: float,
    res_pct: float,
    w_mv: float,
    w_res: float,
    num_gops: int = 1,
    gop_anchor_mode: str = "sampled",
    max_total_patches_cap: int = 30000,
    mv_compensate: str = "none",  # {"none","median"}
    mv_dir_mode: str = "l0",
    w_mv_l0: float = 1.0,
    w_mv_l1: float = 1.0,
    fill_to_images: bool = True,
    force_fallback_no_cv_reader: bool = False,
    sample_id: str = "",
    caption: str = "",
    decode_backsearch_max: int = 32,
    per_video_timeout_sec: int = 0,
    group_by_time: bool = False,
    keep_blocks_within_canvas: bool = True,
    min_pixels: int = 56 * 56,
    max_pixels: int = 768 * 768,
    orig_messages: Optional[List[Dict[str, Any]]] = None,
    skip_black_frames: bool = True,
    skip_corrupt_frames: bool = True,
    black_y_mean_thr: float = 8.0,
    black_y_std_thr: float = 6.0,
    corrupt_green_frac_thr: float = 0.35,
    corrupt_g_thr: int = 180,
    corrupt_rb_thr: int = 90,
    ensure_per_second: bool = False,
    sec_stride: float = 1.0,
    min_blocks_per_second: int = 1,
    p_full_per_bucket: int = 3,
    auto_num_gops: bool = True,
    analysis_candidate_multiplier: float = 8.0,
    analysis_max_frames: int = 1024,
    block_selection_frame_penalty: float = 0.35,
    mv_source: str = "vector",
    score_source: str = "mvres",
    bitcost_grid: str = "sub",
    bitcost_pct: float = 99.0,
    bitcost_log_scale: bool = True,
    decode_backend: str = "auto",
    mask_letterbox: bool = True,
    letterbox_dark_thr: float = 16.0,
    collage_patch_order: str = "time",
) -> Tuple[str, str, Optional[Dict[str, Any]]]:
    """
    Returns (status, message, index_record). status in {"ok","skip","fail"}.
    index_record is returned for:
      - status=="ok": assets index record
      - status=="skip" with unsupported codec: unsupported codec record
    """
    # avoid OpenCV oversubscription inside multiprocessing
    try:
        cv2.setNumThreads(1)
    except Exception:
        pass
    os.environ["OPENCV_LOG_LEVEL"] = "ERROR"
    # Reduce FFmpeg/OpenCV log spam on corrupted streams (can become a major I/O bottleneck).
    # This affects OpenCV's internal FFmpeg backend in many builds.
    os.environ.setdefault("OPENCV_FFMPEG_CAPTURE_OPTIONS", "loglevel;error")

    vp = str(video_path)
    if not vp or (not Path(vp).exists()):
        return "skip", f"{key} missing video", None

    score_source = str(score_source).lower().strip()
    if score_source not in {"mvres", "bitcost"}:
        score_source = "mvres"
    auto_num_gops = bool(auto_num_gops)
    analysis_candidate_multiplier = float(max(1.0, float(analysis_candidate_multiplier)))
    analysis_max_frames = int(max(64, int(analysis_max_frames)))
    block_selection_frame_penalty = float(max(0.0, float(block_selection_frame_penalty)))
    bitcost_grid = str(bitcost_grid).lower().strip()
    if bitcost_grid not in {"sub", "mb", "ctu", "adaptive", "auto"}:
        bitcost_grid = "sub"
    decode_backend = str(decode_backend).lower().strip()
    if decode_backend not in {"auto", "opencv", "ffmpeg_native"}:
        decode_backend = "auto"

    collage_patch_order = str(collage_patch_order).lower().strip()
    if collage_patch_order == "wh":
        collage_patch_order = "score"
    if collage_patch_order not in {"time", "score"}:
        collage_patch_order = "time"

    codec_name = ffprobe_video_codec_name(vp)
    if codec_name not in {"h264", "hevc"}:
        out_dir = str(Path(out_root) / key)
        rec = {
            "kind": "unsupported_codec",
            "status": "skip",
            "key": str(key),
            "video": str(vp),
            "out_dir": str(out_dir),
            "codec_name": str(codec_name) if codec_name else "unknown",
            "supported_codecs": ["h264", "hevc"],
        }
        return "skip", f"{key} unsupported codec={rec['codec_name']}", rec

    out_dir = str(Path(out_root) / key)
    done_mark = str(Path(out_dir) / "_DONE")
    if Path(done_mark).exists():
        return "skip", f"{key} already done", None

    ensure_dir(out_dir)

    total_frames, fps, H0, W0 = get_total_frames_fps(vp)
    if total_frames <= 0:
        return "fail", f"{key} cannot read total_frames", None

    # -----------------------------
    # Adaptive sampling (Scheme 2 + 3)
    #   - seq_len adapts to output num_images (multiple of 4): seq_len = clamp(4*num_images, 64, 256)
    #   - window_len_frames/num_candidates/top_k_windows adapt to duration for more even coverage
    # This happens automatically so callers do not need to change CLI.
    # -----------------------------
    fps_use = float(fps) if (fps and fps > 0) else 30.0
    duration_sec = float(total_frames) / float(fps_use) if total_frames > 0 else 0.0

    if auto_num_gops or int(num_gops) <= 0:
        num_gops_auto = _clamp_int(round(duration_sec / 8.0), 1, 16)
        num_gops = int(min(num_gops_auto, max(1, total_frames)))
    else:
        num_gops = int(max(1, int(num_gops)))

    # Estimate padded token grid size S_full without decoding.
    # Use smart_resize on original H/W and then pad to (2*patch) to ensure hb/wb even.
    p_est = int(patch)
    pad_base_est = int(2 * p_est)
    try:
        resize_h_est, resize_w_est = smart_resize(
            height=int(H0),
            width=int(W0),
            factor=int(pad_base_est),
            min_pixels=int(min_pixels),
            max_pixels=int(max_pixels),
        )
    except Exception:
        resize_h_est, resize_w_est = int(H0), int(W0)

    pad_bottom_est = (pad_base_est - (int(resize_h_est) % pad_base_est)) % pad_base_est
    pad_right_est = (pad_base_est - (int(resize_w_est) % pad_base_est)) % pad_base_est
    H1_est = int(resize_h_est) + int(pad_bottom_est)
    W1_est = int(resize_w_est) + int(pad_right_est)
    hb_est = int(H1_est // max(1, p_est))
    wb_est = int(W1_est // max(1, p_est))
    # Ensure even grid (should hold by construction)
    if hb_est % 2 != 0:
        hb_est += 1
    if wb_est % 2 != 0:
        wb_est += 1
    S_full_est = int(max(1, hb_est * wb_est))

    # Estimate num_images from budget (auto budget if max_total_patches<=0; otherwise from provided budget).
    cap_total = int(max_total_patches_cap) if int(max_total_patches_cap) > 0 else 30000
    max_total_in = int(max_total_patches)
    if max_total_in <= 0:
        max_total_est, _dbg_est = auto_max_total_patches(
            S_full=S_full_est,
            total_frames=int(total_frames),
            fps=float(fps_use),
            cap_total=int(cap_total),
        )
    else:
        # If user provided a budget, respect it but cap it.
        max_total_est = int(min(int(max_total_in), int(cap_total)))

    num_images_est = int(max(1, int(max_total_est // max(1, S_full_est))))
    # Enforce 4-image grouping and minimum GOP coverage.
    num_images_est = max(4, int(num_images_est // 4) * 4)
    num_images_est = int(max(int(num_gops) * 4 if int(num_gops) > 0 else 4, int(num_images_est)))

    # Scheme 2: seq_len derived from num_images (candidate pool grows with output)
    seq_len_adapt = _clamp_int(float(analysis_candidate_multiplier) * float(num_images_est), 64, int(analysis_max_frames))
    # Do not exceed total_frames too aggressively; duplicates are ok, but keep it bounded.
    if int(total_frames) > 0:
        seq_len_adapt = int(min(int(seq_len_adapt), int(max(16, total_frames))))

    # Scheme 3: window params adapt to duration (more even coverage for long videos)
    if duration_sec <= 15.0:
        win_sec = 6.0
    elif duration_sec <= 60.0:
        win_sec = 10.0
    else:
        win_sec = 12.0

    window_len_frames_adapt = _round_to_multiple(float(fps_use) * float(win_sec), 16)
    window_len_frames_adapt = int(max(32, min(int(window_len_frames_adapt), int(total_frames))))

    num_candidates_adapt = int(max(5, min(20, int(math.ceil(duration_sec / 5.0)))))
    top_k_windows_adapt = int(max(2, min(8, int(math.ceil(duration_sec / 20.0)))))
    top_k_windows_adapt = int(min(int(top_k_windows_adapt), int(num_candidates_adapt)))

    # Override user-provided sampling parameters with adaptive ones.
    seq_len = int(seq_len_adapt)
    window_len_frames = int(window_len_frames_adapt)
    num_candidates = int(num_candidates_adapt)
    top_k_windows = int(top_k_windows_adapt)

    # 1) sample frames on an information-reweighted timeline.
    # High-energy regions get denser candidate frames; low-energy regions get fewer.
    # This is equivalent to constructing a new non-uniform video timeline and then
    # sampling candidate frames from that timeline.
    if int(total_frames) <= 0:
        windows = [(0, 0, 0)]
        frame_ids = [0] * int(seq_len)
        win_dbg = {
            "mode": "uniform_empty_fallback",
            "total_frames": int(total_frames),
            "fps": float(fps),
            "seq_len": int(seq_len),
            "window_len_frames_input": int(window_len_frames),
            "num_candidates_input": int(num_candidates),
            "top_k_windows_input": int(top_k_windows),
            "note": "Empty-video fallback.",
        }
    else:
        windows = [(0, int(total_frames) - 1, 0)]
        frame_ids, energy_cdf_dbg = sample_frame_ids_by_energy_cdf(
            video_path=str(vp),
            total_frames=int(total_frames),
            fps=float(fps_use),
            target_count=int(seq_len),
            bin_sec=0.5,
            smooth_bins=1,
            uniform_mix=0.15,
            max_per_bin=16,
        )
        win_dbg = {
            "mode": "pkt_energy_cdf_timeline",
            "total_frames": int(total_frames),
            "fps": float(fps),
            "seq_len": int(seq_len),
            "window_len_frames_input": int(window_len_frames),
            "num_candidates_input": int(num_candidates),
            "top_k_windows_input": int(top_k_windows),
            "note": "Candidate frames are sampled on a PB packet-size energy reweighted timeline instead of uniformly on the raw timeline.",
            "energy_cdf_debug": energy_cdf_dbg,
        }

    # ---- pkt_size dense peaks (PB only, exclude keyframes) ----
    # After the energy-CDF candidate sampling above, we still inject a small number
    # of local peak-centered frames so sudden bursts are less likely to be missed.
    pkt_dense = True
    pkt_bin_sec = 0.5
    pkt_peaks_cap = 8
    pkt_peaks_per_sec = 0.5  # ~1 peak / 2s
    pkt_peak_neighbor = 1
    pkt_smooth_bins = 1

    pkt_dbg: Optional[Dict[str, Any]] = None
    if bool(pkt_dense) and int(total_frames) > 0:
        try:
            peak_fids, pkt_dbg = pick_peak_frame_ids_from_pkt_size(
                video_path=str(vp),
                fps=float(fps_use),
                total_frames=int(total_frames),
                bin_sec=float(pkt_bin_sec),
                peaks_cap=int(pkt_peaks_cap),
                peaks_per_sec=float(pkt_peaks_per_sec),
                neighbor=int(pkt_peak_neighbor),
                smooth_bins=int(pkt_smooth_bins),
            )
            if peak_fids:
                frame_ids = sorted(set([int(x) for x in frame_ids] + [int(x) for x in peak_fids]))
                win_dbg["pkt_dense"] = True
                win_dbg["pkt_added"] = int(len(peak_fids))
                win_dbg["frame_ids_len_after_pkt"] = int(len(frame_ids))
                win_dbg["frame_ids_preview_after_pkt"] = [int(x) for x in frame_ids[:32]]
            else:
                win_dbg["pkt_dense"] = True
                win_dbg["pkt_added"] = 0
        except Exception as e:
            win_dbg["pkt_dense"] = True
            win_dbg["pkt_error"] = repr(e)

    # Save pkt debug into window_debug so meta.json carries it automatically.
    if pkt_dbg is not None:
        win_dbg["pkt_debug"] = pkt_dbg

    # Cap candidate frames to keep runtime stable; prefer keeping peak frames.
    MAX_CAND = int(max(64, int(analysis_max_frames)))
    if len(frame_ids) > MAX_CAND:
        peak_keep: set = set()
        if isinstance(pkt_dbg, dict):
            try:
                peak_keep = set(int(x) for x in pkt_dbg.get("peak_frame_ids", []) if isinstance(x, (int, float)))
            except Exception:
                peak_keep = set()
        peak_sorted = sorted([x for x in frame_ids if x in peak_keep])
        rest = [x for x in frame_ids if x not in peak_keep]
        need = MAX_CAND - len(peak_sorted)
        if need <= 0:
            frame_ids = peak_sorted[:MAX_CAND]
        else:
            if len(rest) <= need:
                frame_ids = sorted(peak_sorted + rest)
            else:
                idxs = np.linspace(0, len(rest) - 1, need, dtype=np.int32).tolist()
                picked = [rest[int(i)] for i in idxs]
                frame_ids = sorted(peak_sorted + picked)
        win_dbg["frame_ids_len_capped"] = int(len(frame_ids))
        win_dbg["frame_ids_cap"] = int(MAX_CAND)
    win_dbg["frame_ids_preview"] = [int(x) for x in frame_ids[:32]]

    # IMPORTANT: downstream uses seq_len as the final candidate-frame count after
    # information-reweighted sampling, peak injection, and candidate capping.
    seq_len = int(len(frame_ids))

    # 2) optional temporal coverage refinement
    if bool(ensure_per_second):
        frame_ids = enforce_time_coverage_frame_ids(
            frame_ids=frame_ids,
            total_frames=int(total_frames),
            fps=float(fps),
            seq_len=int(seq_len),
            stride_sec=float(sec_stride),
        )
    # ===== GOP anchor pre-alignment (sampled / keyframe / hybrid) =====
    # V2.1: build variable-length GOPs by cumulative PB packet-size energy and use
    # their anchor positions to seed GOP anchors on the sampled frame timeline.
    gop_anchor_mode = str(gop_anchor_mode).lower().strip()
    if gop_anchor_mode not in {"sampled", "keyframe", "hybrid"}:
        gop_anchor_mode = "sampled"

    num_gops_seed = int(min(max(1, int(num_gops)), max(1, len(frame_ids))))

    variable_gop_segments, variable_gop_dbg = build_variable_length_gops_by_energy(
        video_path=str(vp),
        total_frames=int(total_frames),
        fps=float(fps_use),
        target_num_gops=int(num_gops_seed),
        bin_sec=0.5,
        smooth_bins=1,
        min_span_sec=1.5,
        max_span_sec=6.0,
    )

    anchor_t_seed: List[int] = []
    if len(frame_ids) > 0:
        for seg in variable_gop_segments:
            anchor_sec = float(seg.get("anchor_sec", 0.0))
            target_fid = int(round(anchor_sec * float(fps_use)))
            if int(total_frames) > 0:
                target_fid = max(0, min(int(total_frames) - 1, int(target_fid)))
            best_t = min(range(len(frame_ids)), key=lambda ii: abs(int(frame_ids[ii]) - int(target_fid)))
            anchor_t_seed.append(int(best_t))

    if len(anchor_t_seed) == 0:
        anchor_t_seed = [0]
    anchor_t_seed = sorted(set(int(x) for x in anchor_t_seed))
    if 0 not in anchor_t_seed:
        anchor_t_seed = [0] + anchor_t_seed

    num_gops_seed = int(max(1, len(anchor_t_seed)))

    keyframe_align_dbg: Dict[str, Any] = {
        "mode": str(gop_anchor_mode),
        "anchor_t_seed": [int(x) for x in anchor_t_seed],
        "replaced": [],
        "keyframes_found": 0,
        "variable_gop_debug": variable_gop_dbg,
    }

    if gop_anchor_mode in {"keyframe", "hybrid"} and len(frame_ids) > 0:
        kf_fids = ffprobe_keyframe_frame_ids(vp, fps=float(fps), total_frames=int(total_frames))
        keyframe_align_dbg["keyframes_found"] = int(len(kf_fids))

        if len(kf_fids) > 0:
            replacements: List[Tuple[int, int, int]] = []

            if gop_anchor_mode == "keyframe":
                pick_idx = np.linspace(0, len(kf_fids) - 1, len(anchor_t_seed), dtype=np.int32).tolist()
                chosen_kfs = [int(kf_fids[int(i)]) for i in pick_idx]
            else:
                chosen_kfs = []
                kf_arr = np.asarray(kf_fids, dtype=np.int32)
                for t_anchor in anchor_t_seed:
                    target_fid = int(frame_ids[int(t_anchor)])
                    pos = int(np.searchsorted(kf_arr, target_fid, side="left"))
                    cands = []
                    if 0 <= pos < len(kf_arr):
                        cands.append(int(kf_arr[pos]))
                    if pos - 1 >= 0:
                        cands.append(int(kf_arr[pos - 1]))
                    if cands:
                        best = min(cands, key=lambda x: abs(int(x) - int(target_fid)))
                    else:
                        best = int(target_fid)
                    chosen_kfs.append(int(best))

            for t_anchor, kf_fid in zip(anchor_t_seed, chosen_kfs):
                t_anchor = int(t_anchor)
                old_fid = int(frame_ids[t_anchor])
                new_fid = int(max(0, min(int(total_frames) - 1, int(kf_fid)))) if int(total_frames) > 0 else int(kf_fid)
                frame_ids[t_anchor] = int(new_fid)
                replacements.append((int(t_anchor), int(old_fid), int(new_fid)))

            keyframe_align_dbg["replaced"] = [
                {"t": int(t), "old_fid": int(o), "new_fid": int(n)} for (t, o, n) in replacements
            ]
    # ===== end GOP anchor pre-alignment =====
    # 3) derive a stable token grid first. For bitcost mode we can analyze many more
    # candidate frames without decoding RGB up front.
    p = int(patch)
    pad_base = 2 * p
    pz = np.zeros((p, p, 3), dtype=np.uint8)

    resize_h = int(resize_h_est)
    resize_w = int(resize_w_est)
    H_dec = int(H0)
    W_dec = int(W0)
    do_resize = (int(resize_h) != int(H_dec)) or (int(resize_w) != int(W_dec))
    scale_h = float(resize_h) / float(max(1, H_dec))
    scale_w = float(resize_w) / float(max(1, W_dec))
    pad_bottom = int((pad_base - (int(resize_h) % pad_base)) % pad_base)
    pad_right = int((pad_base - (int(resize_w) % pad_base)) % pad_base)
    H1 = int(resize_h) + int(pad_bottom)
    W1 = int(resize_w) + int(pad_right)
    if H1 % p != 0 or W1 % p != 0:
        return "fail", f"{key} padded size not divisible by patch: H1={H1} W1={W1} patch={p}", None

    hb, wb = H1 // p, W1 // p
    if hb % 2 != 0 or wb % 2 != 0:
        return "fail", f"{key} patch grid not even after pad: hb={hb} wb={wb}", None

    frames_bgr_pad: List[np.ndarray] = []
    good_mask = [True] * int(len(frame_ids))

    # 5) fetch scoring features
    use_cv = (HAS_CV_READER and (not force_fallback_no_cv_reader))
    cv_reader_fallback_cause = "none"
    if not HAS_CV_READER:
        cv_reader_fallback_cause = "cv_reader_unavailable"
    elif force_fallback_no_cv_reader:
        cv_reader_fallback_cause = "force_no_cv_reader"
    items: Optional[List[Dict[str, Any]]] = None
    if use_cv:
        try:
            if score_source == "bitcost":
                items = cv_reader_fetch_bitcost(vp, frame_ids)
            else:
                items = cv_reader_fetch_mvres(vp, frame_ids)
        except Exception:
            items = None
            use_cv = False
            cv_reader_fallback_cause = "fetch_exception"
    # Safety: some corrupted streams may cause cv_reader to miss requested frame indices.
    if use_cv and items is not None and len(items) != len(frame_ids):
        items = None
        use_cv = False
        cv_reader_fallback_cause = "len_mismatch"

    residual_fallback_frames = 0
    residual_fallback_by_reason: Dict[str, int] = {
        "cv_reader_unavailable": 0,
        "force_no_cv_reader": 0,
        "fetch_exception": 0,
        "len_mismatch": 0,
        "other_cv_reader_unavailable_or_misaligned": 0,
        "frame_exception": 0,
    }

    # fallback scoring uses residual proxy centered at 128
    fused_maps: List[np.ndarray] = []
    top_lb, bottom_lb, left_lb, right_lb = 0, int(H1), 0, int(W1)
    content_mask = np.ones((int(H1), int(W1)), dtype=np.float32)
    mv_source = str(mv_source).lower().strip()
    if mv_source not in {"vector", "energy", "energy_median", "auto"}:
        mv_source = "vector"
    if score_source != "bitcost":
        frames_bgr = decode_frames_bgr(
            vp,
            frame_ids,
            backsearch_max=int(decode_backsearch_max),
            backend=decode_backend,
        )

        good_mask = []
        for fr in frames_bgr:
            bad = frame_is_bad(
                fr,
                skip_black_frames=bool(skip_black_frames),
                skip_corrupt_frames=bool(skip_corrupt_frames),
                black_y_mean_thr=float(black_y_mean_thr),
                black_y_std_thr=float(black_y_std_thr),
                solid_y_std_thr=4.0,
                solid_color_std_thr=4.0,
                solid_max_range_thr=12.0,
                corrupt_green_frac_thr=float(corrupt_green_frac_thr),
                corrupt_g_thr=int(corrupt_g_thr),
                corrupt_rb_thr=int(corrupt_rb_thr),
            )
            good_mask.append(not bool(bad))

        if (len(good_mask) > 0) and (not good_mask[0]):
            first_good = None
            for i_g in range(1, len(good_mask)):
                if good_mask[i_g]:
                    first_good = i_g
                    break
            if first_good is not None:
                frames_bgr[0], frames_bgr[first_good] = frames_bgr[first_good], frames_bgr[0]
                frame_ids[0], frame_ids[first_good] = int(frame_ids[first_good]), int(frame_ids[0])
                good_mask[0], good_mask[first_good] = good_mask[first_good], good_mask[0]
            else:
                return "skip", f"{key} all decoded frames are black/pure-color/corrupted", None

        frames_bgr_rs: List[np.ndarray] = []
        if do_resize:
            for fr in frames_bgr:
                frames_bgr_rs.append(_resize_bgr(fr, int(resize_h), int(resize_w)))
        else:
            frames_bgr_rs = frames_bgr

        for i, fr in enumerate(frames_bgr_rs):
            frp, (pb, pr) = pad_to_multiple_of_bgr(fr, pad_base)
            frames_bgr_pad.append(frp)
            if i == 0:
                pad_bottom, pad_right = pb, pr

        top_lb, bottom_lb, left_lb, right_lb = detect_letterbox_bbox_bgr(
            frames_bgr_pad[0], dark_thr=float(letterbox_dark_thr)
        )
        content_mask = np.zeros((int(H1), int(W1)), dtype=np.float32)
        content_mask[int(top_lb):int(bottom_lb), int(left_lb):int(right_lb)] = 1.0

    for t in range(len(frame_ids)):
        fr = frames_bgr_pad[t] if (score_source != "bitcost" and t < len(frames_bgr_pad)) else None
        if use_cv and items is not None and t < len(items):
            try:
                if score_source == "bitcost":
                    fused = bitcost_item_to_score_map(
                        items[t],
                        out_h=int(H1),
                        out_w=int(W1),
                        grid=bitcost_grid,
                        pct=float(bitcost_pct),
                        log_scale=bool(bitcost_log_scale),
                        codec_name=codec_name,
                    )
                else:
                    mv = np.asarray(items[t]["motion_vector"])
                    ry = np.asarray(items[t]["residual_y"])
                    me_raw = np.asarray(items[t]["motion_energy"]) if "motion_energy" in items[t] else None
                    mem_raw = np.asarray(items[t]["motion_energy_median"]) if "motion_energy_median" in items[t] else None
                    if mv_source == "energy":
                        me = me_raw
                    elif mv_source == "energy_median":
                        me = mem_raw
                    elif mv_source == "auto":
                        me = mem_raw if mem_raw is not None else me_raw
                    else:
                        me = None

                    # Resize mv/residual to match resized RGB (pixel-aligned assumption)
                    if do_resize:
                        ry = _resize_gray(ry, int(resize_h), int(resize_w))
                        mv = _resize_mv_and_scale(mv, int(resize_h), int(resize_w), scale_h=scale_h, scale_w=scale_w)
                        if me is not None:
                            me = _resize_gray(me.astype(np.float32), max(1, int(resize_h) // 4), max(1, int(resize_w) // 4))

                    # Pad residual_y to match padded RGB
                    if ry.shape[0] != H1 or ry.shape[1] != W1:
                        ry_pad = cv2.copyMakeBorder(
                            ry,
                            0,
                            pad_bottom,
                            0,
                            pad_right,
                            borderType=cv2.BORDER_CONSTANT,
                            value=128,
                        )
                    else:
                        ry_pad = ry
                    fused = mv_res_score_map(
                        mv=mv,
                        res_y=ry_pad,
                        mv_energy=me,
                        mv_unit_div=mv_unit_div,
                        mv_pct=mv_pct,
                        res_pct=res_pct,
                        w_mv=w_mv,
                        w_res=w_res,
                        mv_compensate=mv_compensate,
                        mv_dir_mode=mv_dir_mode,
                        w_mv_l0=w_mv_l0,
                        w_mv_l1=w_mv_l1,
                    )
                if bool(mask_letterbox):
                    fused = (fused * content_mask).astype(np.float32)
            except Exception:
                # Any mismatch -> fall back for this frame
                residual_fallback_frames += 1
                residual_fallback_by_reason["frame_exception"] += 1
                if fr is None:
                    fused = np.zeros((int(H1), int(W1)), dtype=np.float32)
                    if bool(mask_letterbox):
                        fused = (fused * content_mask).astype(np.float32)
                    fused_maps.append(fused.astype(np.float32))
                    continue
                ry = bgr_to_residual_y_u8(fr)
                fused = mv_res_score_map(
                    mv=np.zeros((1, 1, 2), dtype=np.float32),
                    res_y=ry,
                    mv_unit_div=mv_unit_div,
                    mv_pct=mv_pct,
                    res_pct=res_pct,
                    w_mv=0.0,
                    w_res=1.0,
                    mv_compensate="none",
                    mv_dir_mode="l0",
                    w_mv_l0=1.0,
                    w_mv_l1=1.0,
                )
                if bool(mask_letterbox):
                    fused = (fused * content_mask).astype(np.float32)
        else:
            # fallback: use "residual_y proxy" only
            residual_fallback_frames += 1
            if cv_reader_fallback_cause in residual_fallback_by_reason:
                residual_fallback_by_reason[cv_reader_fallback_cause] += 1
            else:
                residual_fallback_by_reason["other_cv_reader_unavailable_or_misaligned"] += 1
            if fr is None:
                fused = np.zeros((int(H1), int(W1)), dtype=np.float32)
                if bool(mask_letterbox):
                    fused = (fused * content_mask).astype(np.float32)
                fused_maps.append(fused.astype(np.float32))
                continue
            ry = bgr_to_residual_y_u8(fr)
            fused = mv_res_score_map(
                mv=np.zeros((1, 1, 2), dtype=np.float32),
                res_y=ry,
                mv_unit_div=mv_unit_div,
                mv_pct=mv_pct,
                res_pct=res_pct,
                w_mv=0.0,
                w_res=1.0,
                mv_compensate="none",
                mv_dir_mode="l0",
                w_mv_l0=1.0,
                w_mv_l1=1.0,
            )
            if bool(mask_letterbox):
                fused = (fused * content_mask).astype(np.float32)
        fused_maps.append(fused.astype(np.float32))

    # 6) decide token budget
    S_full = hb * wb  # patches per full image/canvas

    # How many GOP anchors (each anchor contributes one full I-frame canvas)
    p_full_per_bucket = int(max(0, int(p_full_per_bucket)))
    images_per_bucket = int(1 + p_full_per_bucket)  # I_full + P_full_* canvases
    num_gops_req = int(max(1, len(keyframe_align_dbg.get("anchor_t_seed", []))))

    # To avoid buckets degenerating (especially the last bucket), we require that each bucket
    # has enough non-anchor frames to supply P canvases.
    # Each full P canvas needs `blocks_per_full_canvas = S_full/4` blocks.
    # One frame can supply at most `max_blocks = (hb//2)*(wb//2)` blocks.
    blocks_per_full_canvas = int(S_full // 4)
    max_blocks = int((hb // 2) * (wb // 2))
    needed_blocks_per_bucket = int(p_full_per_bucket * blocks_per_full_canvas)
    min_good_frames_per_bucket = int(math.ceil(float(needed_blocks_per_bucket) / float(max(1, max_blocks))))
    min_good_frames_per_bucket = max(1, int(min_good_frames_per_bucket))

    # Need at least `min_good_frames_per_bucket` frames strictly between anchors (or anchor and end).
    # So anchor gap in t-space must be >= (min_good_frames_per_bucket + 1).
    min_anchor_gap_t = int(min_good_frames_per_bucket + 1)

    # Start from the variable-length GOP anchors built above. If spacing is too tight for
    # the current bucket budget, thin them out while preserving order and the first anchor.
    anchor_t_seed_eff = [int(x) for x in keyframe_align_dbg.get("anchor_t_seed", [0])]
    anchor_t_seed_eff = sorted(set(anchor_t_seed_eff)) if anchor_t_seed_eff else [0]
    filtered_anchor_t_seed: List[int] = []
    last_keep: Optional[int] = None
    for t_anchor in anchor_t_seed_eff:
        t_anchor = int(t_anchor)
        if last_keep is None:
            filtered_anchor_t_seed.append(int(t_anchor))
            last_keep = int(t_anchor)
            continue
        if (int(t_anchor) - int(last_keep)) >= int(min_anchor_gap_t):
            filtered_anchor_t_seed.append(int(t_anchor))
            last_keep = int(t_anchor)
    if len(filtered_anchor_t_seed) == 0:
        filtered_anchor_t_seed = [0]
    keyframe_align_dbg["anchor_t_seed_before_spacing_filter"] = [int(x) for x in anchor_t_seed_eff]
    keyframe_align_dbg["anchor_t_seed_after_spacing_filter"] = [int(x) for x in filtered_anchor_t_seed]

    # Latest valid anchor t so the last bucket still has enough frames:
    # len(frame_ids) - t_anchor - 1 >= min_good_frames_per_bucket
    max_anchor_t = int(max(0, len(frame_ids) - 1 - int(min_good_frames_per_bucket)))

    eligible_anchor_count = int(max(1, (max_anchor_t // max(1, min_anchor_gap_t)) + 1))
    if len(filtered_anchor_t_seed) > 0:
        filtered_anchor_t_seed = [int(x) for x in filtered_anchor_t_seed if int(x) <= int(max_anchor_t)]
        if len(filtered_anchor_t_seed) == 0:
            filtered_anchor_t_seed = [0]
        num_gops_eff = int(min(int(len(filtered_anchor_t_seed)), eligible_anchor_count))
        keyframe_align_dbg["anchor_t_seed_final"] = [int(x) for x in filtered_anchor_t_seed[:num_gops_eff]]
        keyframe_align_dbg["num_gops_eff"] = int(num_gops_eff)

    cap_total = int(max_total_patches_cap) if int(max_total_patches_cap) > 0 else 30000
    max_total = int(max_total_patches)
    auto_budget_dbg: Optional[Dict[str, Any]] = None

    if max_total <= 0:
        # auto budget based on resolution + duration, hard-capped by cap_total
        max_total, auto_budget_dbg = auto_max_total_patches(
            S_full=S_full, total_frames=total_frames, fps=fps, cap_total=cap_total
        )

    # Convert patch budget into image budget. Budget may reduce bucket count, but must
    # not create new buckets beyond the GOP anchors derived above.
    num_images_budget = int(max(1, int(max_total // max(1, S_full))))
    if bool(fill_to_images):
        num_images_budget = max(4, int(math.ceil(float(num_images_budget) / 4.0)) * 4)
    else:
        num_images_budget = max(1, int(num_images_budget))

    if auto_budget_dbg is not None and cap_total > 0:
        num_images_cap = max(1, int(cap_total // max(1, S_full)))
        if bool(fill_to_images):
            num_images_cap = max(4, int(num_images_cap // 4) * 4)
        num_images_budget = min(int(num_images_budget), int(max(1, num_images_cap)))

    # Keep the per-bucket canvas recipe fixed as 1I + p_full_per_bucket * P.
    # Budget only controls how many buckets survive, not how many P canvases a kept
    # bucket receives.
    max_bucket_by_budget = max(1, int(num_images_budget // max(1, images_per_bucket)))
    num_gops_eff = int(min(num_gops_req, max_bucket_by_budget, eligible_anchor_count))

    def build_even_spaced_anchors(n_anchor: int) -> List[int]:
        """Scheme A: anchors approximate equal bucket lengths in t-space.

        For n_anchor buckets over T=len(frame_ids), ideal anchors are at
          t_i ≈ round(i * T / n_anchor), i=0..n_anchor-1
        Then we enforce:
          - anchors are increasing
          - spacing >= min_anchor_gap_t
          - last anchor <= max_anchor_t (so tail bucket has enough frames)
        """
        n_anchor = int(max(1, n_anchor))
        if n_anchor <= 1:
            return [0]

        T = int(len(frame_ids))
        gap = int(max(1, min_anchor_gap_t))

        # 1) ideal equally-sized buckets in t-space
        anchors = [int(round(float(i) * float(T) / float(n_anchor))) for i in range(n_anchor)]
        anchors[0] = 0

        # 2) forward enforce spacing + clamp to max_anchor_t
        out: List[int] = [0]
        for i in range(1, n_anchor):
            t = int(anchors[i])
            # clamp by allowed max anchor
            t = min(int(max_anchor_t), max(0, t))
            # enforce spacing
            t = max(int(out[-1] + gap), int(t))
            # if clamped too hard, keep it at least out[-1]+gap, but not beyond max_anchor_t
            if t > int(max_anchor_t):
                t = int(max_anchor_t)
            out.append(int(t))

        # 3) backward pass to ensure last<=max_anchor_t and spacing
        out[-1] = min(int(out[-1]), int(max_anchor_t))
        for i in range(len(out) - 2, -1, -1):
            if int(out[i + 1]) - int(out[i]) < gap:
                out[i] = int(out[i + 1] - gap)
        out[0] = 0

        # 4) sanitize: clamp to [0, max_anchor_t] and strictly increasing
        cleaned: List[int] = []
        for t in out:
            t = min(int(max_anchor_t), max(0, int(t)))
            if not cleaned:
                cleaned.append(int(t))
            else:
                if int(t) <= int(cleaned[-1]):
                    continue
                cleaned.append(int(t))

        # 5) if we lost anchors due to constraints, fill in additional anchors with gap step
        if len(cleaned) < n_anchor:
            used = set(cleaned)
            # try to insert from left to right
            t_cand = 0
            while len(cleaned) < n_anchor and t_cand <= int(max_anchor_t):
                if t_cand not in used:
                    if not cleaned or (t_cand - cleaned[-1] >= gap):
                        cleaned.append(int(t_cand))
                        used.add(int(t_cand))
                t_cand += gap
            cleaned = sorted(cleaned)

        if not cleaned:
            cleaned = [0]
        if len(cleaned) > n_anchor:
            cleaned = cleaned[: int(n_anchor)]
        if cleaned[0] != 0:
            cleaned[0] = 0
        return cleaned

    def sample_existing_anchors_evenly(anchors: List[int], n_keep: int) -> List[int]:
        """Keep n anchors by uniformly subsampling the existing anchor sequence."""
        anchors = sorted(set(int(x) for x in anchors))
        n_keep = int(max(1, n_keep))
        if not anchors:
            return [0]
        if len(anchors) <= n_keep:
            return anchors
        if n_keep == 1:
            return [int(anchors[0])]

        idxs = np.linspace(0, len(anchors) - 1, n_keep, dtype=np.int32).tolist()
        kept = [int(anchors[int(i)]) for i in idxs]
        kept = sorted(set(kept))

        # Backfill if dedup shrank the list.
        if len(kept) < n_keep:
            used = set(kept)
            for a in anchors:
                if int(a) in used:
                    continue
                kept.append(int(a))
                used.add(int(a))
                if len(kept) >= n_keep:
                    break
            kept = sorted(kept)

        if kept[0] != int(anchors[0]):
            kept[0] = int(anchors[0])
        if len(kept) > n_keep:
            kept = kept[: int(n_keep)]
        return kept

    # Choose effective GOP anchors with guaranteed spacing >= min_anchor_gap_t.
    # Then, if bucket density is too high, reduce bucket count until each bucket
    # has enough non-anchor good frames to form P canvases without repeating blocks.

    # Use filtered variable-length GOP anchors as starting point. When budget forces us
    # to keep fewer buckets, subsample the full anchor set evenly instead of taking a prefix.
    anchor_seed_final = [int(x) for x in keyframe_align_dbg.get("anchor_t_seed_final", [0])]
    anchor_t_indices = sample_existing_anchors_evenly(anchor_seed_final, int(num_gops_eff))
    if len(anchor_t_indices) == 0:
        anchor_t_indices = [0]
    while len(anchor_t_indices) > 1:
        ok_density = True
        for bi, t_anchor in enumerate(anchor_t_indices):
            t_start = int(t_anchor)
            t_end = int(anchor_t_indices[bi + 1]) if (bi + 1) < len(anchor_t_indices) else int(len(frame_ids))
            good_cnt = 0
            for t in range(int(t_start) + 1, int(t_end)):
                if (t < len(good_mask)) and bool(good_mask[t]):
                    good_cnt += 1
            if good_cnt < int(min_good_frames_per_bucket):
                ok_density = False
                break
        if ok_density:
            break
        num_gops_eff = int(max(1, len(anchor_t_indices) - 1))
        anchor_t_indices = sample_existing_anchors_evenly(anchor_seed_final, int(num_gops_eff))
        if len(anchor_t_indices) == 0:
            anchor_t_indices = [0]

    # Post-process tail bucket: if the last bucket is much shorter than the auto-GOP
    # target, merge it back into the previous bucket to avoid tiny tail GOPs.
    if len(anchor_t_indices) >= 2:
        tail_span_t = int(len(frame_ids) - int(anchor_t_indices[-1]))
        tail_span_sec = float(tail_span_t) / float(max(fps_use, 1e-6))
        target_bucket_sec = 8.0 if bool(auto_num_gops) else max(1.5, float(duration_sec) / float(max(1, len(anchor_t_indices))))
        min_tail_merge_sec = max(3.0, 0.5 * float(target_bucket_sec))
        if tail_span_sec < float(min_tail_merge_sec):
            anchor_t_indices = anchor_t_indices[:-1]
            if len(anchor_t_indices) == 0:
                anchor_t_indices = [0]

    anchor_set = set(int(x) for x in anchor_t_indices)

    num_buckets = int(len(anchor_t_indices))
    num_images_budget = int(num_buckets * int(images_per_bucket))

    # available blocks per frame
    max_blocks = int((hb // 2) * (wb // 2))

    # 7) build output token sequence
    patches_list: List[np.ndarray] = []
    src_pos_list: List[List[int]] = []       # [frame_id, src_patch_h, src_patch_w]
    src_frameid_list: List[int] = []         # original frame id per patch
    src_ts_list: List[float] = []            # timestamp (seconds) per patch (pad tokens are -1.0)
    canvas_bucket_ids: List[int] = []
    canvas_roles: List[str] = []
    bucket_ranges_t: List[Dict[str, int]] = []

    fps_use_ts = float(fps) if (fps and fps > 0) else 30.0

    # helper to append one frame's patches (in a specific block order)
    def append_frame_patches(src_img_idx: int, frame_rgb: np.ndarray, blocks: List[Tuple[int, int]]):
        # NOTE: src_patch_position.npy first column stores the ORIGINAL frame_id (not the 0..seq_len-1 index).
        try:
            src_fid = int(frame_ids[int(src_img_idx)])
        except Exception:
            src_fid = -1
        # Place blocks in the order they appear in the input list.
        for (bh, bw) in blocks:
            coords = block_to_4_patches(bh, bw)
            for (ph, pw) in coords:
                patch_rgb = extract_patch_rgb(frame_rgb, ph, pw, patch=p)
                patches_list.append(patch_rgb.astype(np.uint8))
                # src_patch_position: [frame_id, patch_h, patch_w]
                src_pos_list.append([int(src_fid), int(ph), int(pw)])
                # keep src_frame_ids.npy consistent
                src_frameid_list.append(int(src_fid))
                if src_fid >= 0:
                    src_ts_list.append(float(src_fid) / float(fps_use_ts))
                else:
                    src_ts_list.append(-1.0)

    def append_pad_block():
        for _ in range(4):
            patches_list.append(pz)
            src_pos_list.append([-1, -1, -1])
            src_frameid_list.append(-1)
            src_ts_list.append(-1.0)

    # GOP buckets: each bucket contributes one full I canvas + K full P canvases.
    all_blocks = list(iter_blocks_in_raster(hb, wb))
    blocks_per_full_canvas = int(S_full // 4)

    selected_blocks_by_frame: Dict[int, List[Tuple[int, int]]] = {}
    selected_blocks_count_by_frame: Dict[int, int] = {}
    bucket_plans: List[Dict[str, Any]] = []

    def select_bucket_blocks_stratified(
        candidates: List[Tuple[int, int, int, float]],
        t_start: int,
        t_end: int,
        extra_images: int,
        blocks_per_canvas: int,
    ) -> Tuple[List[Tuple[int, int, int, float]], List[Dict[str, Any]]]:
        """Select P blocks using information-weighted temporal focus regions."""
        if extra_images <= 0 or blocks_per_canvas <= 0 or not candidates:
            return [], []

        t_lo = int(t_start) + 1
        t_hi = int(max(t_lo, t_end))
        chosen_all: List[Tuple[int, int, int, float]] = []
        used_keys: set = set()
        frame_mass: Dict[int, float] = {}
        frame_peak: Dict[int, float] = {}

        for (t0, _, _, sc0) in candidates:
            t0 = int(t0)
            sc0 = float(sc0)
            frame_mass[t0] = float(frame_mass.get(t0, 0.0) + max(0.0, sc0))
            frame_peak[t0] = max(float(frame_peak.get(t0, 0.0)), max(0.0, sc0))

        def _take_from_pool(pool: List[Tuple[int, int, int, float]], need: int) -> List[Tuple[int, int, int, float]]:
            out: List[Tuple[int, int, int, float]] = []
            if need <= 0:
                return out
            pool_sorted = sorted(pool, key=lambda x: (-float(x[3]), int(x[0])))
            for rec in pool_sorted:
                key = (int(rec[0]), int(rec[1]), int(rec[2]))
                if key in used_keys:
                    continue
                out.append(rec)
                used_keys.add(key)
                if len(out) >= int(need):
                    break
            return out

        frames_sorted = sorted(int(t) for t in frame_mass.keys() if int(t) >= int(t_lo) and int(t) < int(t_hi))
        focus_ranges: List[Dict[str, Any]] = []
        if not frames_sorted:
            return [], focus_ranges

        weights = np.asarray(
            [max(1e-6, float(frame_mass.get(int(t), 0.0)) + 0.35 * float(frame_peak.get(int(t), 0.0))) for t in frames_sorted],
            dtype=np.float64,
        )
        weight_sum = float(weights.sum())
        if (not np.isfinite(weight_sum)) or weight_sum <= 0.0:
            weights = np.ones((len(frames_sorted),), dtype=np.float64)
            weight_sum = float(weights.sum())
        cdf = np.cumsum(weights) / float(weight_sum)

        for pi in range(int(extra_images)):
            q_lo = float(pi) / float(max(1, extra_images))
            q_hi = float(pi + 1) / float(max(1, extra_images))
            q_mid = 0.5 * (q_lo + q_hi)
            idx_mid = int(np.searchsorted(cdf, q_mid, side="left"))
            idx_mid = max(0, min(len(frames_sorted) - 1, idx_mid))
            idx_lo = int(np.searchsorted(cdf, q_lo, side="left"))
            idx_hi = int(np.searchsorted(cdf, q_hi, side="left"))
            idx_lo = max(0, min(len(frames_sorted) - 1, idx_lo))
            idx_hi = max(idx_lo, min(len(frames_sorted) - 1, idx_hi))

            center_t = int(frames_sorted[idx_mid])
            range_lo_t = int(frames_sorted[idx_lo])
            range_hi_t = int(frames_sorted[idx_hi]) + 1
            local_pool = [
                rec for rec in candidates
                if int(rec[0]) >= int(range_lo_t) and int(rec[0]) < int(range_hi_t)
            ]
            picked = _take_from_pool(local_pool, int(blocks_per_canvas))
            if len(picked) < int(blocks_per_canvas):
                radius = int(max(1, (range_hi_t - range_lo_t)))
                expanded_lo = int(max(t_lo, center_t - radius))
                expanded_hi = int(min(t_hi, center_t + radius + 1))
                fallback_pool = [
                    rec for rec in candidates
                    if int(rec[0]) >= int(expanded_lo) and int(rec[0]) < int(expanded_hi)
                ]
                picked.extend(_take_from_pool(fallback_pool, int(blocks_per_canvas - len(picked))))
            if len(picked) < int(blocks_per_canvas):
                global_pool = [rec for rec in candidates if int(rec[0]) >= int(t_lo) and int(rec[0]) < int(t_hi)]
                picked.extend(_take_from_pool(global_pool, int(blocks_per_canvas - len(picked))))

            chosen_all.extend(picked)
            focus_ranges.append({
                "image_idx": int(pi),
                "range_t_start": int(range_lo_t),
                "range_t_end": int(max(range_lo_t + 1, range_hi_t)),
                "center_t": int(center_t),
                "frame_mass_sum": float(sum(float(frame_mass.get(int(t1), 0.0)) for t1 in frames_sorted if int(t1) >= int(range_lo_t) and int(t1) < int(range_hi_t))),
                "picked_blocks": int(len(picked)),
            })

        return chosen_all, focus_ranges

    def reorder_chosen_for_place(
        chosen_blocks: List[Tuple[int, int, int, float]],
        order_mode: str,
    ) -> List[Tuple[int, int, int, float]]:
        """Order already-selected blocks for placement inside P canvases."""
        if not chosen_blocks:
            return []
        mode = str(order_mode).lower().strip()
        if mode == "time":
            # Place by source frame_id first, then by block top-left in raster order.
            return sorted(
                chosen_blocks,
                key=lambda x: (int(x[0]), int(x[1]), int(x[2]), -float(x[3])),
            )
        # "score": keep strongest blocks first.
        return sorted(
            chosen_blocks,
            key=lambda x: (-float(x[3]), int(x[0]), int(x[1]), int(x[2])),
        )
    for bi, t_anchor in enumerate(anchor_t_indices):
        t_start = int(t_anchor)
        t_end = int(anchor_t_indices[bi + 1]) if (bi + 1) < len(anchor_t_indices) else int(len(frame_ids))
        bucket_ranges_t.append({
            "bucket_id": int(bi),
            "anchor_t": int(t_anchor),
            "p_t_start_exclusive": int(t_start),
            "p_t_end_exclusive": int(t_end),
        })

        # Build P candidates strictly inside this bucket and before next anchor.
        # Order: time asc, score desc (within each frame).
        candidates: List[Tuple[int, int, int, float]] = []  # (t, bh, bw, score)
        frame_pick_count: Dict[int, int] = {}
        bwb = int(wb // 2)
        for t in range(int(t_start) + 1, int(t_end)):
            if int(t) in anchor_set:
                continue
            if (t < len(good_mask)) and (not bool(good_mask[t])):
                continue
            ps = score_map_to_patch_scores(fused_maps[int(t)], patch=p)  # (hb,wb)
            bs = patch_scores_to_block_scores(ps)  # (hb//2, wb//2)
            flat = bs.reshape(-1).astype(np.float32)
            if flat.size <= 0:
                continue
            order = np.argsort(-flat, kind="stable")
            for fid in order.tolist():
                bh = int(int(fid) // bwb)
                bw = int(int(fid) % bwb)
                sc = float(flat[int(fid)])
                used_n = int(frame_pick_count.get(int(t), 0))
                penalty = 1.0 / math.sqrt(1.0 + float(block_selection_frame_penalty) * float(used_n))
                candidates.append((int(t), int(bh), int(bw), float(sc * penalty)))
                frame_pick_count[int(t)] = used_n + 1

        bucket_plans.append({
            "bucket_id": int(bi),
            "anchor_t": int(t_anchor),
            "t_start": int(t_start),
            "t_end": int(t_end),
            "bucket_score": float(sum(float(sc0) for (_, _, _, sc0) in candidates)),
            "candidates": [(int(t0), int(bh0), int(bw0), float(sc0)) for (t0, bh0, bw0, sc0) in candidates],
        })

    bucket_extra_images: Dict[int, int] = {
        int(bp["bucket_id"]): int(p_full_per_bucket)
        for bp in bucket_plans
    }
    bucket_image_counts: Dict[int, int] = {
        int(bp["bucket_id"]): int(images_per_bucket)
        for bp in bucket_plans
    }
    total_block_budget = int(num_buckets * int(p_full_per_bucket) * int(S_full // 4))

    selected_t_needed: set = set(int(x) for x in anchor_t_indices)
    for bucket_plan in bucket_plans:
        bi = int(bucket_plan["bucket_id"])
        extra_images = int(bucket_extra_images.get(int(bi), 0))
        chosen_for_place, focus_ranges = select_bucket_blocks_stratified(
            candidates=[(int(t0), int(bh0), int(bw0), float(sc0)) for (t0, bh0, bw0, sc0) in bucket_plan.get("candidates", [])],
            t_start=int(bucket_plan.get("t_start", int(bucket_plan["anchor_t"]))),
            t_end=int(bucket_plan.get("t_end", int(bucket_plan["anchor_t"]) + 1)),
            extra_images=int(extra_images),
            blocks_per_canvas=int(blocks_per_full_canvas),
        )
        chosen_for_place = reorder_chosen_for_place(chosen_for_place, collage_patch_order)
        bucket_plan["chosen_for_place"] = [(int(t0), int(bh0), int(bw0), float(sc0)) for (t0, bh0, bw0, sc0) in chosen_for_place]
        bucket_plan["sparse_focus_ranges"] = focus_ranges
        if extra_images <= 0:
            continue
        for (t_sel, bh_sel, bw_sel, _) in chosen_for_place:
            selected_t_needed.add(int(t_sel))
            selected_blocks_by_frame.setdefault(int(t_sel), []).append((int(bh_sel), int(bw_sel)))

    for t_sel, blks in selected_blocks_by_frame.items():
        selected_blocks_count_by_frame[int(t_sel)] = int(len(blks))

    frames_bgr_pad_by_t: Dict[int, np.ndarray] = {}
    if score_source == "bitcost":
        selected_t_sorted = sorted(int(x) for x in selected_t_needed)
        selected_frame_ids = [int(frame_ids[int(t)]) for t in selected_t_sorted]
        selected_frames_bgr = decode_frames_bgr(
            vp,
            selected_frame_ids,
            backsearch_max=int(decode_backsearch_max),
            backend=decode_backend,
        )
        good_mask = [False] * int(len(frame_ids))
        for t_sel, fr in zip(selected_t_sorted, selected_frames_bgr):
            bad = frame_is_bad(
                fr,
                skip_black_frames=bool(skip_black_frames),
                skip_corrupt_frames=bool(skip_corrupt_frames),
                black_y_mean_thr=float(black_y_mean_thr),
                black_y_std_thr=float(black_y_std_thr),
                solid_y_std_thr=4.0,
                solid_color_std_thr=4.0,
                solid_max_range_thr=12.0,
                corrupt_green_frac_thr=float(corrupt_green_frac_thr),
                corrupt_g_thr=int(corrupt_g_thr),
                corrupt_rb_thr=int(corrupt_rb_thr),
            )
            good_mask[int(t_sel)] = not bool(bad)
            fr_rs = _resize_bgr(fr, int(resize_h), int(resize_w)) if do_resize else fr
            fr_pad, _ = pad_to_multiple_of_bgr(fr_rs, pad_base)
            frames_bgr_pad_by_t[int(t_sel)] = fr_pad
        if frames_bgr_pad_by_t:
            first_t = min(frames_bgr_pad_by_t.keys())
            top_lb, bottom_lb, left_lb, right_lb = detect_letterbox_bbox_bgr(
                frames_bgr_pad_by_t[int(first_t)], dark_thr=float(letterbox_dark_thr)
            )
    else:
        frames_bgr_pad_by_t = {int(t): fr for t, fr in enumerate(frames_bgr_pad)}

    if not frames_bgr_pad_by_t:
        return "skip", f"{key} no frames available after selection", None

    num_images_actual = int(sum(bucket_image_counts.values())) if bucket_image_counts else int(num_buckets)
    max_total = int(num_images_actual * int(S_full))

    def choose_anchor_t(bucket_plan: Dict[str, Any]) -> int:
        t_anchor0 = int(bucket_plan["anchor_t"])
        if bool(good_mask[t_anchor0]):
            return int(t_anchor0)
        for (t_sel, _, _, _) in bucket_plan["chosen_for_place"]:
            if int(t_sel) < len(good_mask) and bool(good_mask[int(t_sel)]) and int(t_sel) in frames_bgr_pad_by_t:
                return int(t_sel)
        return int(t_anchor0)

    for bucket_plan in bucket_plans:
        bi = int(bucket_plan["bucket_id"])
        t_anchor_use = choose_anchor_t(bucket_plan)
        fr_anchor = frames_bgr_pad_by_t.get(int(t_anchor_use))
        if fr_anchor is not None:
            append_frame_patches(src_img_idx=int(t_anchor_use), frame_rgb=fr_anchor[:, :, ::-1], blocks=all_blocks)
        else:
            for _ in all_blocks:
                append_pad_block()
        canvas_bucket_ids.append(int(bi))
        canvas_roles.append("I_full")

        cpos = 0
        chosen_for_place = bucket_plan["chosen_for_place"]
        extra_images = int(bucket_extra_images.get(int(bi), 0))
        for pi in range(int(extra_images)):
            for _ in range(blocks_per_full_canvas):
                if cpos < len(chosen_for_place):
                    t_sel, bh_sel, bw_sel, _ = chosen_for_place[cpos]
                    fr = frames_bgr_pad_by_t.get(int(t_sel))
                    if (fr is not None) and (int(t_sel) < len(good_mask)) and bool(good_mask[int(t_sel)]):
                        append_frame_patches(src_img_idx=int(t_sel), frame_rgb=fr[:, :, ::-1], blocks=[(int(bh_sel), int(bw_sel))])
                    else:
                        append_pad_block()
                    cpos += 1
                else:
                    append_pad_block()
            canvas_bucket_ids.append(int(bi))
            canvas_roles.append(f"P_full_{int(pi) + 1}")

    # 7.5) pad/truncate to exactly max_total patches so that we can fill whole canvases
    valid_n = int(len(patches_list))
    if valid_n > int(max_total):
        patches_list = patches_list[: int(max_total)]
        src_pos_list = src_pos_list[: int(max_total)]
        src_frameid_list = src_frameid_list[: int(max_total)]
        src_ts_list = src_ts_list[: int(max_total)]
        valid_n = int(max_total)

    while len(patches_list) < int(max_total):
        patches_list.append(pz)
        src_pos_list.append([-1, -1, -1])
        src_frameid_list.append(-1)
        src_ts_list.append(-1.0)

    patches = np.stack(patches_list, axis=0).astype(np.uint8) if patches_list else np.zeros((0, p, p, 3), dtype=np.uint8)
    src_patch_position = np.asarray(src_pos_list, dtype=np.int32) if src_pos_list else np.zeros((0, 3), dtype=np.int32)
    src_frame_ids = np.asarray(src_frameid_list, dtype=np.int32) if src_frameid_list else np.zeros((0,), dtype=np.int32)
    src_timestamps = np.asarray(src_ts_list, dtype=np.float32) if src_ts_list else np.zeros((0,), dtype=np.float32)

    # Pack into full canvases
    images_rgb, patch_position, img_ptr_arr = pack_patches_to_canvases(patches, hb=hb, wb=wb, patch=p)
    frame_ids_arr = np.asarray(frame_ids, dtype=np.int32)

    # Save JPG canvases (do NOT save images.npy / patches.npy)
    jpg_files = save_canvases_as_jpg(images_rgb, out_dir=out_dir, quality=95)

    # Save required arrays/metadata
    np.save(str(Path(out_dir) / "patch_position.npy"), patch_position, allow_pickle=False)
    np.save(str(Path(out_dir) / "img_ptr.npy"), img_ptr_arr, allow_pickle=False)
    np.save(str(Path(out_dir) / "frame_ids.npy"), frame_ids_arr, allow_pickle=False)
    np.save(str(Path(out_dir) / "src_patch_position.npy"), src_patch_position, allow_pickle=False)  # [frame_id, patch_h, patch_w]
    # Compatibility alias: many downstream jsonl schemas expect `patch_positions.npy`
    np.save(str(Path(out_dir) / "patch_positions.npy"), src_patch_position, allow_pickle=False)
    np.save(str(Path(out_dir) / "src_frame_ids.npy"), src_frame_ids, allow_pickle=False)
    np.save(str(Path(out_dir) / "src_timestamps.npy"), src_timestamps, allow_pickle=False)
    # np.save(str(Path(out_dir) / "position.npy"), src_patch_position, allow_pickle=False) 

    meta = {
        "video": vp,
        "key": key,
        "out_dir": str(Path(out_root) / key),
        "mirror_src_root": None,
        "total_frames": int(total_frames),
        "fps": float(fps),
        "orig_hw": [int(H0), int(W0)],
        "pad_bottom_right": [int(pad_bottom), int(pad_right)],
        "padded_hw": [int(H1), int(W1)],
        "patch": int(p),
        "pad_base": int(pad_base),
        "hb_wb": [int(hb), int(wb)],
        "seq_len": int(seq_len),
        "auto_num_gops": bool(auto_num_gops),
        "analysis_candidate_multiplier": float(analysis_candidate_multiplier),
        "analysis_max_frames": int(analysis_max_frames),
        "windows": [{"start": int(a), "end": int(b), "pkt_sum": int(c)} for (a, b, c) in windows],
        "window_debug": win_dbg,
        "w_mv": float(w_mv),
        "w_res": float(w_res),
        "mv_compensate": str(mv_compensate),
        "mv_dir_mode": str(mv_dir_mode),
        "w_mv_l0": float(w_mv_l0),
        "w_mv_l1": float(w_mv_l1),
        "mv_source": str(mv_source),
        "use_cv_reader": bool(use_cv),
        "score_source": str(score_source),
        "bitcost_grid": str(bitcost_grid),
        "bitcost_pct": float(bitcost_pct),
        "bitcost_log_scale": bool(bitcost_log_scale),
        "decode_backend": str(decode_backend),
        "residual_fallback": {
            "used": bool(residual_fallback_frames > 0),
            "frames": int(residual_fallback_frames),
            "ratio": float(residual_fallback_frames) / float(max(1, len(frames_bgr_pad))),
            "proxy": "sobel_luma_magnitude_centered_128",
            "cv_reader_fallback_cause": str(cv_reader_fallback_cause),
            "reasons": {k: int(v) for k, v in residual_fallback_by_reason.items()},
        },
        "max_total_patches": int(max_total),
        "I_frame_patches": int(S_full),
        "num_gops_req": int(num_gops_req),
        "num_gops_eff": int(len(anchor_t_indices)),
        "gop_anchor_t_indices": [int(x) for x in anchor_t_indices],
        "gop_anchor_frame_ids": [int(frame_ids[int(x)]) for x in anchor_t_indices],
        "gop_anchor_mode": str(gop_anchor_mode),
        "images_per_bucket_max": int(images_per_bucket),
        "bucket_image_counts": {str(k): int(v) for k, v in bucket_image_counts.items()},
        "bucket_extra_images": {str(k): int(v) for k, v in bucket_extra_images.items()},
        "num_images_budget": int(num_images_budget),
        "p_full_per_bucket": int(p_full_per_bucket),
        "bucket_ranges_t": bucket_ranges_t,
        "bucket_sparse_focus_ranges": [
            {
                "bucket_id": int(bp.get("bucket_id", -1)),
                "focus_ranges": [
                    {
                        "image_idx": int(fr.get("image_idx", -1)),
                        "range_t_start": int(fr.get("range_t_start", -1)),
                        "range_t_end": int(fr.get("range_t_end", -1)),
                        "center_t": int(fr.get("center_t", -1)),
                        "frame_mass_sum": float(fr.get("frame_mass_sum", 0.0)),
                        "picked_blocks": int(fr.get("picked_blocks", 0)),
                    }
                    for fr in bp.get("sparse_focus_ranges", [])
                ],
            }
            for bp in bucket_plans
        ],
        "canvas_bucket_ids": canvas_bucket_ids,
        "canvas_roles": canvas_roles,
        "keyframe_align": keyframe_align_dbg,
        "total_patches_target": int(max_total),
        "total_patches_valid": int(valid_n),
        "num_images": int(images_rgb.shape[0]),
        "jpg_files": jpg_files,
        "patches_per_image": int(S_full),
        "auto_budget": auto_budget_dbg,
        "global_block_selection": {
            "total_block_budget": int(total_block_budget),
            "max_blocks_per_frame": int(max_blocks),
            "frame_penalty": float(block_selection_frame_penalty),
            "selected_blocks_total": int(sum(selected_blocks_count_by_frame.values()) if selected_blocks_count_by_frame else 0),
            "selected_blocks_per_frame": selected_blocks_count_by_frame,
        },
        "fill_to_images": bool(fill_to_images),
        "group_by_time": bool(group_by_time),
        "keep_blocks_within_canvas": bool(keep_blocks_within_canvas),
        "ensure_per_second": bool(ensure_per_second),
        "sec_stride": float(sec_stride),
        "min_blocks_per_second": int(min_blocks_per_second),
        "mask_letterbox": bool(mask_letterbox),
        "letterbox_dark_thr": float(letterbox_dark_thr),
        "letterbox_bbox_tblr": [int(top_lb), int(bottom_lb), int(left_lb), int(right_lb)],
        # src_patch_position now stores [frame_id, patch_h, patch_w] (pad tokens are -1)
        "src_patch_position_format": "[frame_id, patch_h, patch_w] (pad tokens are -1)",
        "position_npy": "position.npy",
        "sample_id": str(sample_id) if sample_id else "",
        "caption": str(caption) if caption else "",
        "smart_resize": {
            "min_pixels": int(min_pixels),
            "max_pixels": int(max_pixels),
            "factor": int(2 * p),
            "do_resize": bool(do_resize),
            "resize_hw": [int(resize_h), int(resize_w)],
            "scale_hw": [float(scale_h), float(scale_w)],
        },
        "decoded_selected_frames": int(len(frames_bgr_pad_by_t)),
        "has_src_timestamps": True,
        "src_timestamps_npy": "src_timestamps.npy",
    }
    meta["mirror_src_root"] = os.environ.get("MAGEVL_CODEC_MIRROR_SRC_ROOT")
    with open(str(Path(out_dir) / "meta.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

    # A compact per-sample record for fast join in step-2 (avoid per-sample filesystem probing).
    index_rec: Dict[str, Any] = {
        "video": vp,
        "key": str(key),
        "asset_dir": str(out_dir),
        "jpg": list(jpg_files),
        "patch_position": "patch_position.npy",
        "img_ptr": "img_ptr.npy",
        "src_patch_position": "src_patch_position.npy",
        "patch_positions": "patch_positions.npy",
        "src_frame_ids": "src_frame_ids.npy",
        "frame_ids": "frame_ids.npy",
        "position": "position.npy",
        "sample_id": str(sample_id) if sample_id else "",
        "caption": str(caption) if caption else "",
        # prefer meta-derived count; this equals the packed total (including pad patches)
        "num_patches": int(meta.get("max_total_patches", int(patches.shape[0]))),
        "num_images": int(meta.get("num_images", int(images_rgb.shape[0]))),
        "fps": float(meta.get("fps", fps)),
        "patch": int(meta.get("patch", p)),
        "padded_hw": meta.get("padded_hw"),
        "src_timestamps": "src_timestamps.npy",
        "orig_messages": orig_messages if isinstance(orig_messages, list) else [],
    }

    Path(done_mark).write_text("ok\n", encoding="utf-8")
    return "ok", f"{key} ok patches={int(patches.shape[0])} images={int(images_rgb.shape[0])} jpg={len(jpg_files)} use_cv_reader={int(use_cv)}", index_rec