15_07_expt/scripts/pipeline.py

"""Main entry point for the rebuilt figure-skating data pipeline."""

from __future__ import annotations

import json
import os
import pickle
import random
import time
from collections import Counter
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Iterable

import numpy as np
from tqdm import tqdm

try:
    from . import labels, preprocessing
except ImportError:
    import labels
    import preprocessing

# Imported at module load (not lazily) on purpose: gpu_pose runs a one-time LD_LIBRARY_PATH
# fix that may re-exec the process. Doing it here means any re-exec happens during startup,
# before run_pipeline does any work; a lazy import mid-run would restart and lose progress.
# Kept optional so the MediaPipe path still works on boxes without torch/ultralytics installed.
try:
    from . import gpu_pose
except ImportError:
    try:
        import gpu_pose
    except Exception:  # noqa: BLE001 -- gpu extractors simply unavailable
        gpu_pose = None


PROJECT_ROOT = Path(__file__).resolve().parents[1]
os.environ.setdefault("GLOG_minloglevel", "2")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
os.environ.setdefault("ABSL_LOGGING_MIN_LOG_LEVEL", "2")

CONFIG = {
    "datasets": {
        "skatingverse": PROJECT_ROOT / "data" / "skatingverse",
        "current": PROJECT_ROOT / "data" / "current",
        "mmfs": PROJECT_ROOT / "data" / "mmfs",
    },
    "pose_estimator": "yolo",  # "mediapipe" (CPU) or "yolo" (GPU, see yolo_* keys below)
    "yolo_weights": "yolo11n-pose.pt",  # or "yolo11s-pose.pt" / "yolo11m-pose.pt"
    "yolo_imgsz": 384,  # 640 needed for reliable skater detection on wide rink shots (256 underdetects)
    "yolo_half": True,  # fp16 inference
    "yolo_max_batch": 128,  # cap frames per predict() call to bound VRAM under parallelism
    "yolo_decode_workers": 16,  # GIL-releasing cv2 decode threads feeding the single GPU consumer
    "target_sequence_length": 128,
    "confidence_threshold": 0.3,
    "smoothing_window": 7,
    "smoothing_polyorder": 3,
    "fps": 30.0,
    "max_videos": 10000,
    "default_width": 16.0,
    "default_height": 9.0,
    "train_ratio": 0.8,
    "val_ratio": 0.1,
    "test_ratio": 0.1,
    "random_seed": 42,
    "output_dir": PROJECT_ROOT / "data" / "processed_new",
    "skeleton_cache_dir": PROJECT_ROOT / "data" / "processed_new" / "skeleton_cache",
    "num_workers": 12,
    "video_num_workers": 1,
    "suppress_mediapipe_logs": True,
    "mediapipe_model_complexity": 0,
    "extract_every_n_frames": 1,
    "max_video_frames": 300,
    "show_frame_progress": False,
}

# Hardcoded run controls. Edit these values before running this file.
DRY_RUN = False
DRY_RUN_SOURCE = "skatingverse"
NUM_DRY_RUN_SAMPLES = 1
SOURCES = None  # Example: ["current"] or ["skatingverse", "current"]; None means all sources.
POSE_ESTIMATOR = "yolo"  # "yolo" (GPU) or "mediapipe" (CPU) for video-derived sources
YOLO_WEIGHTS = "yolo11n-pose.pt"  # "yolo11n-pose.pt" / "yolo11s-pose.pt" / "yolo11m-pose.pt"
OUTPUT_DIR = CONFIG["output_dir"]
TARGET_SEQUENCE_LENGTH = CONFIG["target_sequence_length"]
NUM_WORKERS = CONFIG["num_workers"]
VIDEO_NUM_WORKERS = CONFIG["video_num_workers"]
MEDIAPIPE_MODEL_COMPLEXITY = CONFIG["mediapipe_model_complexity"]
EXTRACT_EVERY_N_FRAMES = CONFIG["extract_every_n_frames"]
MAX_VIDEO_FRAMES = CONFIG["max_video_frames"]

VIDEO_EXTENSIONS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v"}
SKELETON_EXTENSIONS = {".pkl", ".pickle", ".npz", ".npy"}


def load_pickle(path: Path):
    with path.open("rb") as f:
        return pickle.load(f)


def save_pickle(obj, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("wb") as f:
        pickle.dump(obj, f)


class suppress_stderr:
    """Temporarily redirect process stderr to os.devnull for noisy native libraries."""

    def __init__(self, enabled: bool):
        self.enabled = enabled
        self._saved_fd = None
        self._devnull_fd = None

    def __enter__(self):
        if not self.enabled:
            return self
        self._saved_fd = os.dup(2)
        self._devnull_fd = os.open(os.devnull, os.O_WRONLY)
        os.dup2(self._devnull_fd, 2)
        return self

    def __exit__(self, exc_type, exc, tb):
        if not self.enabled:
            return False
        os.dup2(self._saved_fd, 2)
        os.close(self._saved_fd)
        os.close(self._devnull_fd)
        return False


def _first_existing_key(mapping, keys):
    for key in keys:
        if key in mapping:
            return mapping[key]
    return None


def load_skeleton_file(path: Path) -> list[dict]:
    """Load skeleton samples from pickle, npy, or npz files.

    Returns a list of dicts containing at least "skeleton" and optional "label", "fps",
    "width", and "height".
    """
    if path.suffix in {".pkl", ".pickle"}:
        data = load_pickle(path)
    elif path.suffix == ".npy":
        data = np.load(path, allow_pickle=True)
    elif path.suffix == ".npz":
        data = dict(np.load(path, allow_pickle=True))
    else:
        raise ValueError(f"Unsupported skeleton file: {path}")

    if isinstance(data, np.ndarray) and data.dtype == object and data.shape == ():
        data = data.item()

    if isinstance(data, dict):
        skeletons = _first_existing_key(data, ["skeletons", "data", "x", "X", "features"])
        labels_arr = _first_existing_key(data, ["labels", "y", "Y", "targets"])
        fps_arr = _first_existing_key(data, ["fps", "frame_rates"])
        widths = _first_existing_key(data, ["widths", "video_widths", "width"])
        heights = _first_existing_key(data, ["heights", "video_heights", "height"])
        if skeletons is None and "skeleton" in data:
            skeletons = [data["skeleton"]]
        if skeletons is None:
            return []
        samples = []
        for idx, skeleton in enumerate(list(skeletons)):
            sample = {"skeleton": skeleton, "source_file": str(path), "label": path.parent.name}
            if labels_arr is not None and len(labels_arr) > idx:
                sample["label"] = labels_arr[idx].item() if hasattr(labels_arr[idx], "item") else labels_arr[idx]
            if fps_arr is not None:
                sample["fps"] = fps_arr[idx] if np.ndim(fps_arr) else fps_arr
            if widths is not None:
                sample["width"] = widths[idx] if np.ndim(widths) else widths
            if heights is not None:
                sample["height"] = heights[idx] if np.ndim(heights) else heights
            samples.append(sample)
        return samples

    if isinstance(data, (list, tuple)) and data and isinstance(data[0], dict):
        return [dict({"label": path.parent.name, "source_file": str(path)}, **item) for item in data]

    arr = np.asarray(data, dtype=object if isinstance(data, list) else None)
    if arr.ndim == 2 and arr.shape[1] in (34, 51):
        return [{"skeleton": arr, "source_file": str(path), "label": path.parent.name}]
    if arr.ndim == 3 and arr.shape[1] == 17 and arr.shape[2] in (2, 3):
        return [{"skeleton": arr, "source_file": str(path), "label": path.parent.name}]
    if arr.ndim >= 3:
        return [{"skeleton": skeleton, "source_file": str(path), "label": path.parent.name} for skeleton in arr]
    return []


def discover_skeleton_samples(dataset_dir: Path) -> list[dict]:
    samples = []
    if not dataset_dir.exists():
        return samples
    for path in dataset_dir.rglob("*"):
        if path.suffix.lower() in SKELETON_EXTENSIONS:
            samples.extend(load_skeleton_file(path))
    return samples


def iter_skeleton_files(dataset_dir: Path):
    if not dataset_dir.exists():
        return
    for path in dataset_dir.rglob("*"):
        if path.suffix.lower() in SKELETON_EXTENSIONS:
            yield path


def iter_video_files(dataset_dir: Path):
    if not dataset_dir.exists():
        return
    for path in dataset_dir.rglob("*"):
        if path.suffix.lower() in VIDEO_EXTENSIONS:
            yield path


def load_first_skeleton_sample(dataset_dir: Path) -> dict | None:
    for path in iter_skeleton_files(dataset_dir):
        samples = load_skeleton_file(path)
        if samples:
            return samples[0]
    return None


def iter_skeleton_samples(dataset_dir: Path):
    for path in iter_skeleton_files(dataset_dir):
        for sample in load_skeleton_file(path):
            yield sample


def load_skatingverse_labels(dataset_dir: Path) -> dict[str, int]:
    """Load video-stem -> class-index labels from SkatingVerse's train.txt.

    Only the train split has public ground truth. train.txt/mapping.txt/answer.txt come from
    the dataset's ModelScope repo (awei2003/1st_SkatingVerse_Dataset), not the Kaggle video-only
    mirror `download_datasets.py` pulls from, so they're copied in separately alongside
    train_videos/ and test_videos/. mapping.txt's class-index order matches labels.TAXONOMY
    exactly, so these indices can be used as-is. test_videos has no labels here: answer.txt is
    only the submission-order template for the original challenge (filenames, no classes) -- its
    ground truth was withheld by the organizers.
    """
    labels_path = dataset_dir / "train.txt"
    mapping: dict[str, int] = {}
    if not labels_path.exists():
        return mapping
    with labels_path.open("r", encoding="utf-8") as f:
        for line in f:
            parts = line.split()
            if len(parts) != 2:
                continue
            stem, idx = parts
            mapping[stem] = int(idx)
    return mapping


def discover_video_samples(dataset_dir: Path, label_lookup: dict[str, int] | None = None) -> list[dict]:
    if not dataset_dir.exists():
        return []
    samples = []
    for path in iter_video_files(dataset_dir):
        label = label_lookup.get(path.stem) if label_lookup is not None else path.parent.name
        samples.append({"video_path": path, "label": label, "source_file": str(path)})
    return samples


def discover_one_video_sample(dataset_dir: Path, label_lookup: dict[str, int] | None = None) -> dict | None:
    for path in iter_video_files(dataset_dir):
        label = label_lookup.get(path.stem) if label_lookup is not None else path.parent.name
        return {"video_path": path, "label": label, "source_file": str(path)}
    return None


def iter_video_samples(dataset_dir: Path, label_lookup: dict[str, int] | None = None):
    for path in iter_video_files(dataset_dir):
        label = label_lookup.get(path.stem) if label_lookup is not None else path.parent.name
        yield {"video_path": path, "label": label, "source_file": str(path)}


def extract_skeleton_from_video(video_path: Path, config: dict) -> dict | None:
    """Extract a COCO-17 skeleton from one video with MediaPipe and cache it."""
    cache_dir = Path(config["skeleton_cache_dir"])
    frame_stride = max(1, int(config.get("extract_every_n_frames", 1)))
    max_video_frames = config.get("max_video_frames")
    max_video_frames = None if max_video_frames is None else max(1, int(max_video_frames))
    model_complexity = int(config.get("mediapipe_model_complexity", 1))
    cache_suffix = f"mc{model_complexity}_s{frame_stride}_m{max_video_frames or 'all'}"
    cache_path = cache_dir / f"{video_path.stem}_{cache_suffix}.pkl"
    if cache_path.exists():
        return load_pickle(cache_path)

    suppress_logs = bool(config.get("suppress_mediapipe_logs", True))
    with suppress_stderr(suppress_logs):
        try:
            import cv2
            try:
                from mediapipe.python.solutions import pose as mp_pose
            except ImportError:
                import mediapipe as mp

                if not hasattr(mp, "solutions") or not hasattr(mp.solutions, "pose"):
                    raise RuntimeError(
                        "The installed mediapipe package does not expose the legacy Pose API "
                        "(mediapipe.solutions.pose). Install a MediaPipe version that includes "
                        "the solutions package, or use pre-extracted skeleton files for this run."
                    )
                mp_pose = mp.solutions.pose
        except ImportError as exc:
            raise RuntimeError("Video extraction requires opencv-python and mediapipe") from exc

        coco_indices = [0, 2, 5, 7, 8, 11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28]
        pose = mp_pose.Pose(
            static_image_mode=False,
            model_complexity=model_complexity,
            smooth_landmarks=True,
            min_detection_confidence=config["confidence_threshold"],
            min_tracking_confidence=config["confidence_threshold"],
        )
        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            pose.close()
            return None

        fps = cap.get(cv2.CAP_PROP_FPS) or config["fps"]
        width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) or config["default_width"]
        height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or config["default_height"]
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
        frames = []
        decoded_frames = 0
        processed_frames = 0
        show_frame_progress = bool(config.get("show_frame_progress", False))
        frame_iter = tqdm(
            total=total_frames if total_frames > 0 else None,
            desc=f"Extracting {video_path.name}",
            unit="frame",
            leave=False,
            disable=not show_frame_progress,
        )
        try:
            while True:
                ok, frame = cap.read()
                if not ok:
                    break
                decoded_frames += 1
                if (decoded_frames - 1) % frame_stride != 0:
                    frame_iter.update(1)
                    continue
                if max_video_frames is not None and processed_frames >= max_video_frames:
                    break
                result = pose.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                if result.pose_landmarks:
                    joints = []
                    for idx in coco_indices:
                        landmark = result.pose_landmarks.landmark[idx]
                        joints.append([landmark.x, landmark.y, landmark.visibility])
                    frames.append(joints)
                else:
                    frames.append(np.zeros((17, 3), dtype=np.float64))
                processed_frames += 1
                frame_iter.update(1)
        finally:
            frame_iter.close()
            cap.release()
            pose.close()

    if not frames:
        return None
    sample = {
        "skeleton": np.asarray(frames, dtype=np.float64),
        "fps": float(fps),
        "width": float(width),
        "height": float(height),
        "source_file": str(video_path),
        "extract_every_n_frames": frame_stride,
        "max_video_frames": max_video_frames,
        "mediapipe_model_complexity": model_complexity,
        "decoded_frames": decoded_frames,
        "processed_frames": processed_frames,
    }
    save_pickle(sample, cache_path)
    return sample


def map_label_for_source(source: str, raw_label) -> tuple[str, int] | None:
    if raw_label is None:
        return None
    if source == "skatingverse":
        try:
            return labels.map_skatingverse_label(raw_label)
        except ValueError:
            return None
    if source == "mmfs":
        if isinstance(raw_label, (list, tuple)) and len(raw_label) >= 2:
            return labels.map_mmfs_label(str(raw_label[0]), str(raw_label[1]))
        return labels.map_mmfs_label(str(raw_label), "")
    return labels.map_current_dataset_label(str(raw_label))


def process_sample(sample: dict, source: str, config: dict) -> tuple[np.ndarray, int, str] | None:
    mapped = map_label_for_source(source, sample.get("label"))
    if mapped is None:
        return None
    label_name, label_idx = mapped
    width = float(sample.get("width", config["default_width"]))
    height = float(sample.get("height", config["default_height"]))
    fps = float(sample.get("fps", config["fps"]))
    # The stored skeleton keeps only every extract_every_n_frames-th frame, so consecutive
    # kept frames are `stride` original frames apart in real time. Velocities/accelerations use
    # the *effective* frame rate (fps / stride); otherwise their magnitudes are inflated by the
    # stride factor and, worse, differ in scale between strided video sources and un-strided
    # pre-extracted skeleton sources (mmfs). (Review issue #3.)
    stride = max(1, int(sample.get("extract_every_n_frames", 1) or 1))
    effective_fps = fps / stride

    skeleton = preprocessing.correct_aspect_ratio(sample["skeleton"], width, height)
    skeleton = preprocessing.filter_low_confidence(skeleton, config["confidence_threshold"])
    skeleton = preprocessing.interpolate_missing(skeleton)
    skeleton = preprocessing.smooth_skeleton(
        skeleton,
        config["smoothing_window"],
        config["smoothing_polyorder"],
    )
    skeleton = preprocessing.center_on_hips(skeleton)
    skeleton = preprocessing.normalize_scale(skeleton)
    features = preprocessing.build_feature_tensor(skeleton, effective_fps)
    features = preprocessing.resample_to_length(features, config["target_sequence_length"])
    return features.astype(np.float32), int(label_idx), label_name


def _array_stats(arr: np.ndarray) -> str:
    values = np.asarray(arr)
    finite = np.isfinite(values)
    if not finite.any():
        return f"shape={values.shape}, finite=0/{values.size}"
    finite_values = values[finite]
    return (
        f"shape={values.shape}, finite={finite.sum()}/{values.size}, "
        f"min={finite_values.min():.6f}, max={finite_values.max():.6f}, "
        f"mean={finite_values.mean():.6f}, std={finite_values.std():.6f}"
    )


def dry_run_sample(sample: dict, source: str, config: dict) -> bool:
    """Process one sample with detailed print logging for each important stage."""
    print("\nDry run: one-sample pipeline walkthrough")
    print("=" * 72)
    print(f"Source dataset: {source}")
    print(f"Source file: {sample.get('source_file', '<memory>')}")
    print(f"Raw label: {sample.get('label', '<missing>')}")

    mapped = map_label_for_source(source, sample.get("label"))
    print("\n1. Label mapping")
    print("   Function: map_label_for_source")
    print("   Purpose: convert the source dataset label into the canonical 28-class taxonomy.")
    if mapped is None:
        label_name = "UNMAPPED"
        label_idx = -1
        print("   Result: unmapped label. A full run would discard this sample.")
        print("   Dry-run behavior: continuing anyway so preprocessing and feature extraction can be inspected.")
    else:
        label_name, label_idx = mapped
        print(f"   Result: {sample.get('label')!r} -> {label_name!r} (index {label_idx})")

    width = float(sample.get("width", config["default_width"]))
    height = float(sample.get("height", config["default_height"]))
    fps = float(sample.get("fps", config["fps"]))
    print("\n2. Metadata defaults")
    print("   Purpose: choose geometry and timing values used by normalization and derivatives.")
    print(f"   width={width}, height={height}, fps={fps}")
    print(f"   target_sequence_length={config['target_sequence_length']}")

    skeleton = preprocessing.as_skeleton_array(sample["skeleton"])
    print("\n3. Input skeleton coercion")
    print("   Function: preprocessing.as_skeleton_array")
    print("   Purpose: accept flattened COCO-17 arrays such as (F, 51) and convert to (F, 17, 3).")
    print(f"   Output: {_array_stats(skeleton)}")

    print("\n4. Aspect ratio correction")
    print("   Function: preprocessing.correct_aspect_ratio")
    print("   Purpose: MediaPipe x and y are normalized by different image dimensions; multiply x by width/height before geometry.")
    skeleton = preprocessing.correct_aspect_ratio(skeleton, width, height)
    print(f"   Applied x multiplier: {width / height:.6f}")
    print(f"   Output: {_array_stats(skeleton)}")

    print("\n5. Confidence filtering")
    print("   Function: preprocessing.filter_low_confidence")
    print("   Purpose: mark low-confidence joint coordinates as NaN so interpolation treats them as missing observations.")
    if skeleton.shape[2] >= 3:
        low_conf = int(np.sum(skeleton[..., 2] < config["confidence_threshold"]))
        total_joints = int(np.prod(skeleton[..., 2].shape))
        print(f"   Low-confidence joints before filtering: {low_conf}/{total_joints}")
    skeleton = preprocessing.filter_low_confidence(skeleton, config["confidence_threshold"])
    print(f"   Output: {_array_stats(skeleton)}")

    print("\n6. Missing-data interpolation")
    print("   Function: preprocessing.interpolate_missing")
    print("   Purpose: linearly fill NaN coordinate gaps across time for each joint and coordinate.")
    nan_before = int(np.isnan(skeleton[..., :2]).sum())
    skeleton = preprocessing.interpolate_missing(skeleton)
    nan_after = int(np.isnan(skeleton[..., :2]).sum())
    print(f"   Coordinate NaNs before: {nan_before}; after: {nan_after}")
    print(f"   Output: {_array_stats(skeleton)}")

    print("\n7. Skeleton smoothing")
    print("   Function: preprocessing.smooth_skeleton")
    print("   Purpose: apply Savitzky-Golay smoothing to x/y only, preserving sharper motion changes better than a moving average.")
    print(f"   window={config['smoothing_window']}, polyorder={config['smoothing_polyorder']}")
    skeleton = preprocessing.smooth_skeleton(
        skeleton,
        config["smoothing_window"],
        config["smoothing_polyorder"],
    )
    print(f"   Output: {_array_stats(skeleton)}")

    print("\n8. Hip-centering")
    print("   Function: preprocessing.center_on_hips")
    print("   Purpose: subtract the per-frame midpoint between left and right hips from every joint to remove camera translation.")
    skeleton_xy = preprocessing.center_on_hips(skeleton)
    hip_mid = 0.5 * (
        skeleton_xy[:, preprocessing.LEFT_HIP, :] + skeleton_xy[:, preprocessing.RIGHT_HIP, :]
    )
    print(f"   Hip midpoint after centering: mean_abs={np.mean(np.abs(hip_mid)):.8f}")
    print(f"   Output: {_array_stats(skeleton_xy)}")

    print("\n9. Torso-length scale normalization")
    print("   Function: preprocessing.normalize_scale")
    print("   Purpose: divide each frame by torso length so coordinates are less sensitive to camera distance and resolution.")
    skeleton_xy = preprocessing.normalize_scale(skeleton_xy)
    print(f"   Output: {_array_stats(skeleton_xy)}")

    print("\n10. Joint-angle features")
    print("    Function: preprocessing.compute_joint_angles")
    print("    Purpose: compute 12 cosine angles for anatomical triplets such as elbow, shoulder, hip, knee, and ankle bends.")
    angles = preprocessing.compute_joint_angles(skeleton_xy)
    print(f"    Output: {_array_stats(angles)}")

    print("\n11. Bone-vector features")
    print("    Function: preprocessing.compute_bone_vectors")
    print("    Purpose: compute 10 parent-to-child displacement vectors, flattened to 20 values per frame.")
    bone_vectors = preprocessing.compute_bone_vectors(skeleton_xy)
    print(f"    Output: {_array_stats(bone_vectors)}")

    print("\n12. Velocity features")
    print("    Function: preprocessing.compute_velocities")
    print("    Purpose: compute first temporal derivatives with central differences; the final tensor keeps 8 key joints, 16 values.")
    velocities = preprocessing.compute_velocities(skeleton_xy, fps)
    key_velocities = velocities[:, preprocessing.KEY_JOINTS, :].reshape(velocities.shape[0], -1)
    print(f"    All-joint velocities: {_array_stats(velocities)}")
    print(f"    Key-joint velocities used: {_array_stats(key_velocities)}")

    print("\n13. Angular-velocity features")
    print("    Function: preprocessing.compute_angular_velocities")
    print("    Purpose: compute first temporal derivatives of the 12 angle-cosine features.")
    angular_velocities = preprocessing.compute_angular_velocities(angles, fps)
    print(f"    Output: {_array_stats(angular_velocities)}")

    print("\n14. Feature tensor assembly")
    print("    Function: preprocessing.build_feature_tensor")
    print("    Purpose: concatenate 34 coords + 12 angles + 20 bone vectors + 16 keypoint velocities + 12 angular velocities.")
    features = preprocessing.build_feature_tensor(skeleton_xy, fps)
    print(f"    Output: {_array_stats(features)}")
    print("    Expected per-frame feature dimension: 94")

    print("\n15. Sequence resampling")
    print("    Function: preprocessing.resample_to_length")
    print("    Purpose: linearly interpolate every feature dimension over normalized time to the fixed model length.")
    resized = preprocessing.resample_to_length(features, config["target_sequence_length"])
    print(f"    Output: {_array_stats(resized)}")

    print("\nDry run complete")
    print("=" * 72)
    print(f"Final sample label: {label_name} ({label_idx})")
    print(f"Final model input shape for one sample: {resized.astype(np.float32).shape}")
    print("No files were written.")
    return True


def load_source_samples(source: str, config: dict) -> list[dict]:
    dataset_dir = Path(config["datasets"][source])
    samples = discover_skeleton_samples(dataset_dir)
    if source == "skatingverse":
        label_lookup = load_skatingverse_labels(dataset_dir)
        video_samples = discover_video_samples(dataset_dir, label_lookup)
        unlabeled = sum(1 for vs in video_samples if vs["label"] is None)
        if unlabeled:
            video_samples = [vs for vs in video_samples if vs["label"] is not None]
            print(
                f"  skipping {unlabeled} skatingverse video(s) with no known label "
                "(test_videos ground truth was never made public -- see train.txt vs answer.txt)"
            )
        max_videos = config.get("max_videos")
        if max_videos is not None and len(video_samples) > int(max_videos):
            rng = random.Random(int(config.get("random_seed", 0)))
            rng.shuffle(video_samples)  # representative subset, not filesystem-order-biased
            video_samples = video_samples[: int(max_videos)]
            print(f"  limiting to {len(video_samples)} videos (max_videos={max_videos}, "
                  f"seed={config.get('random_seed', 0)})")
        video_workers = max(1, int(config.get("video_num_workers", 1)))
        pose_estimator = str(config.get("pose_estimator", "mediapipe")).lower()
        if video_samples and pose_estimator not in ("yolo", "yolo_pose", "yolo_pose_gpu"):
            print(
                "  video extraction settings: "
                f"model_complexity={config.get('mediapipe_model_complexity')}, "
                f"stride={config.get('extract_every_n_frames')}, "
                f"max_processed_frames={config.get('max_video_frames')}, "
                f"workers={config.get('video_num_workers')}"
            )
        if pose_estimator in ("yolo", "yolo_pose", "yolo_pose_gpu"):
            if gpu_pose is None:
                raise RuntimeError(
                    "pose_estimator='yolo' requires the GPU extractor (torch + ultralytics), "
                    "but gpu_pose failed to import."
                )
            print(
                "  using YOLO-pose GPU extractor: "
                f"weights={config.get('yolo_weights')}, imgsz={config.get('yolo_imgsz')}, "
                f"half={config.get('yolo_half')}, max_batch={config.get('yolo_max_batch')}, "
                f"decode_workers={config.get('yolo_decode_workers')}"
            )
            for extracted in tqdm(
                gpu_pose.extract_all_videos_yolo(video_samples, config),
                total=len(video_samples),
                desc=f"{source} video extraction (yolo)",
                unit="video",
            ):
                samples.append(extracted)  # label already attached by extract_all_videos_yolo
        elif video_workers == 1:
            for video_sample in tqdm(video_samples, desc=f"{source} video extraction", unit="video"):
                extracted = extract_skeleton_from_video(video_sample["video_path"], config)
                if extracted is not None:
                    samples.append({**extracted, "label": video_sample["label"]})
        else:
            print(f"  extracting {len(video_samples)} video skeleton(s) with {video_workers} worker(s)")
            with ThreadPoolExecutor(max_workers=video_workers) as executor:
                futures = {
                    executor.submit(extract_skeleton_from_video, video_sample["video_path"], config): video_sample
                    for video_sample in video_samples
                }
                for future in tqdm(
                    as_completed(futures),
                    total=len(futures),
                    desc=f"{source} video extraction",
                    unit="video",
                ):
                    video_sample = futures[future]
                    try:
                        extracted = future.result()
                    except Exception as exc:
                        print(f"  skipped video {video_sample['source_file']}: {exc}")
                        continue
                    if extracted is not None:
                        samples.append({**extracted, "label": video_sample["label"]})
    return samples


def load_one_source_sample(source: str, config: dict) -> dict | None:
    """Load a single sample for dry-run mode without scanning/extracting every video."""
    dataset_dir = Path(config["datasets"][source])
    skeleton_sample = load_first_skeleton_sample(dataset_dir)
    if skeleton_sample is not None:
        return skeleton_sample
    if source == "skatingverse":
        video_sample = discover_one_video_sample(dataset_dir, load_skatingverse_labels(dataset_dir))
        if video_sample is not None:
            extracted = extract_skeleton_from_video(video_sample["video_path"], config)
            if extracted is not None:
                return {**extracted, "label": video_sample["label"]}
    return None


def iter_dry_run_samples(source: str, config: dict, limit: int):
    """Yield up to limit dry-run samples from exactly one source."""
    if limit <= 0:
        return
    dataset_dir = Path(config["datasets"][source])
    yielded = 0

    for sample in iter_skeleton_samples(dataset_dir):
        yield sample
        yielded += 1
        if yielded >= limit:
            return

    if source == "skatingverse":
        for video_sample in iter_video_samples(dataset_dir, load_skatingverse_labels(dataset_dir)):
            if video_sample["label"] is None:
                continue
            extracted = extract_skeleton_from_video(video_sample["video_path"], config)
            if extracted is None:
                continue
            yield {**extracted, "label": video_sample["label"]}
            yielded += 1
            if yielded >= limit:
                return


def stratified_split_indices(
    labels_arr: np.ndarray, train_ratio: float, val_ratio: float, test_ratio: float, seed: int
) -> dict[str, np.ndarray]:
    total = train_ratio + val_ratio + test_ratio
    if not np.isclose(total, 1.0):
        raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0")

    rng = random.Random(seed)
    by_class: dict[int, list[int]] = {}
    for idx, label in enumerate(labels_arr):
        by_class.setdefault(int(label), []).append(idx)

    splits = {"train": [], "val": [], "test": []}
    for indices in by_class.values():
        rng.shuffle(indices)
        n = len(indices)
        # Guarantee val/test coverage for rare classes. Plain rounding sends every class with
        # <=2 samples entirely to train (round(1*0.1)=0), so rare jump classes end up with no
        # val/test examples and can never be evaluated. Instead: n<=1 -> train only (nothing to
        # split); n==2 -> one train, one val; n>=3 -> at least one sample in each split.
        # (Review issue #4.)
        if n <= 1:
            n_train, n_val, n_test = n, 0, 0
        elif n == 2:
            n_train, n_val, n_test = 1, 1, 0
        else:
            n_train = int(round(n * train_ratio))
            n_train = min(max(n_train, 1), n - 2)  # leave >=1 each for val and test
            n_val = int(round(n * val_ratio))
            n_val = min(max(n_val, 1), n - n_train - 1)  # leave >=1 for test
            n_test = n - n_train - n_val
        splits["train"].extend(indices[:n_train])
        splits["val"].extend(indices[n_train : n_train + n_val])
        splits["test"].extend(indices[n_train + n_val :])

    np_rng = np.random.default_rng(seed)
    return {name: np_rng.permutation(values) for name, values in splits.items()}


def save_outputs(features: np.ndarray, y: np.ndarray, label_names: list[str], config: dict) -> None:
    output_dir = Path(config["output_dir"])
    output_dir.mkdir(parents=True, exist_ok=True)
    split_indices = stratified_split_indices(
        y,
        config["train_ratio"],
        config["val_ratio"],
        config["test_ratio"],
        config["random_seed"],
    )
    for split, indices in split_indices.items():
        save_pickle(features[indices], output_dir / f"{split}_features.pkl")
        save_pickle(y[indices], output_dir / f"{split}_labels.pkl")

    metadata = {
        "feature_shape": [config["target_sequence_length"], int(features.shape[-1])],
        "num_classes": len(labels.TAXONOMY),
        "taxonomy": labels.TAXONOMY,
        "label_to_idx": labels.LABEL_TO_IDX,
        "class_distribution": {labels.TAXONOMY[int(k)]: int(v) for k, v in Counter(y).items()},
        "splits": {name: int(len(indices)) for name, indices in split_indices.items()},
        "label_names": label_names,
        "config": {k: str(v) if isinstance(v, Path) else v for k, v in config.items() if k != "datasets"},
        "datasets": {k: str(v) for k, v in config["datasets"].items()},
    }
    with (output_dir / "metadata.json").open("w", encoding="utf-8") as f:
        json.dump(metadata, f, indent=2)


def process_source_sample(job: tuple[dict, str, dict]) -> tuple[str, tuple[np.ndarray, int, str] | None, str | None]:
    sample, source, config = job
    try:
        return sample.get("source_file", "<memory>"), process_sample(sample, source, config), None
    except Exception as exc:
        return sample.get("source_file", "<memory>"), None, str(exc)


def run_pipeline(config: dict = CONFIG, sources: Iterable[str] | None = None) -> bool:
    pipeline_start = time.perf_counter()
    selected_sources = list(sources or config["datasets"].keys())
    processed_features = []
    processed_labels = []
    processed_label_names = []
    skipped = Counter()

    for source in selected_sources:
        source_start = time.perf_counter()
        print(f"Loading {source} from {config['datasets'][source]}")
        source_samples = load_source_samples(source, config)
        print(f"  discovered {len(source_samples)} sample(s)")
        workers = max(1, int(config.get("num_workers", 1)))
        if workers == 1:
            for sample in tqdm(source_samples, desc=f"{source} preprocessing", unit="sample"):
                source_file, processed, error = process_source_sample((sample, source, config))
                if error is not None:
                    skipped[f"{source}:error"] += 1
                    if skipped[f"{source}:error"] <= 5:
                        print(f"  skipped {source_file}: {error}")
                    continue
                if processed is None:
                    skipped[f"{source}:unmapped"] += 1
                    continue
                features, label_idx, label_name = processed
                processed_features.append(features)
                processed_labels.append(label_idx)
                processed_label_names.append(label_name)
        else:
            print(f"  processing samples with {workers} worker(s)")
            jobs = [(sample, source, config) for sample in source_samples]
            with ProcessPoolExecutor(max_workers=workers) as executor:
                futures = [executor.submit(process_source_sample, job) for job in jobs]
                for future in tqdm(
                    as_completed(futures),
                    total=len(futures),
                    desc=f"{source} preprocessing",
                    unit="sample",
                ):
                    source_file, processed, error = future.result()
                    if error is not None:
                        skipped[f"{source}:error"] += 1
                        if skipped[f"{source}:error"] <= 5:
                            print(f"  skipped {source_file}: {error}")
                        continue
                    if processed is None:
                        skipped[f"{source}:unmapped"] += 1
                        continue
                    features, label_idx, label_name = processed
                    processed_features.append(features)
                    processed_labels.append(label_idx)
                    processed_label_names.append(label_name)
        elapsed = time.perf_counter() - source_start
        print(f"  finished {source} in {elapsed:.1f}s")

    if not processed_features:
        print("No labeled samples were processed. Check dataset paths and label files.")
        if skipped:
            print(f"Skipped summary: {dict(skipped)}")
        return False

    features = np.stack(processed_features, axis=0)
    y = np.asarray(processed_labels, dtype=np.int64)
    save_outputs(features, y, processed_label_names, config)
    total_elapsed = time.perf_counter() - pipeline_start
    print(f"Saved {features.shape[0]} samples to {Path(config['output_dir']).resolve()}")
    print(f"Feature tensor shape: {features.shape}")
    print(f"Class distribution: {dict(Counter(processed_label_names))}")
    print(f"Total pipeline time: {total_elapsed:.1f}s")
    if skipped:
        print(f"Skipped summary: {dict(skipped)}")
    return True


def run_dry_run(config: dict = CONFIG, sources: Iterable[str] | None = None) -> bool:
    if sources is None:
        source = DRY_RUN_SOURCE
    else:
        selected_sources = list(sources)
        if not selected_sources:
            print("Dry run needs exactly one source. Edit DRY_RUN_SOURCE or SOURCES.")
            return False
        source = selected_sources[0]

    sample_limit = int(config.get("num_dry_run_samples", NUM_DRY_RUN_SAMPLES))
    if sample_limit <= 0:
        print("NUM_DRY_RUN_SAMPLES must be at least 1.")
        return False

    dry_start = time.perf_counter()
    print(f"Searching for dry-run samples in {source}: {config['datasets'][source]}")
    print(f"Dry run sample limit: {sample_limit}")

    processed = 0
    for sample in iter_dry_run_samples(source, config, sample_limit):
        if "skeleton" not in sample:
            continue
        processed += 1
        print(f"  selected sample {processed}/{sample_limit}")
        print(f"  sample selection elapsed: {time.perf_counter() - dry_start:.1f}s")
        ok = dry_run_sample(sample, source, config)
        if not ok:
            return False

    if processed > 0:
        print(f"Dry run processed {processed} sample(s) in {time.perf_counter() - dry_start:.1f}s.")
        return True

    print("  no usable sample found")
    print("No sample with skeleton data was found for dry run. Check dataset paths or edit DRY_RUN_SOURCE.")
    return False


def main() -> int:
    config = dict(CONFIG)
    config["datasets"] = dict(CONFIG["datasets"])
    config["output_dir"] = Path(OUTPUT_DIR)
    config["skeleton_cache_dir"] = Path(OUTPUT_DIR) / "skeleton_cache"
    config["target_sequence_length"] = TARGET_SEQUENCE_LENGTH
    config["num_workers"] = NUM_WORKERS
    config["video_num_workers"] = VIDEO_NUM_WORKERS
    config["num_dry_run_samples"] = NUM_DRY_RUN_SAMPLES
    config["mediapipe_model_complexity"] = MEDIAPIPE_MODEL_COMPLEXITY
    config["extract_every_n_frames"] = EXTRACT_EVERY_N_FRAMES
    config["max_video_frames"] = MAX_VIDEO_FRAMES
    config["pose_estimator"] = POSE_ESTIMATOR
    config["yolo_weights"] = YOLO_WEIGHTS
    if DRY_RUN:
        return 0 if run_dry_run(config, SOURCES) else 1
    return 0 if run_pipeline(config, SOURCES) else 1


if __name__ == "__main__":
    raise SystemExit(main())