tests/test_training_dataset.py

"""Synthetic, GPU-free tests for fpgm.training (manifest, bundle_io, augment, dataset).

Every fixture here is a hand-built ``tmp_path`` directory tree with tiny
(16x12, 5-frame) fake ``vace/window_*`` bundles -- no real dataset, no
network, no GPU, no DiffSynth import (that stays lazy inside
``fpgm.training.train``, exercised separately/manually against the real
Wan2.1-VACE checkpoint, not here).
"""

from __future__ import annotations

import json

import cv2
import numpy as np
import pytest

from fpgm.geometry.normals import encode_normals_rgb
from fpgm.training.augment import (
    apply_augmentations,
    channel_dropout,
    mask_dilate_erode,
    pose_noise,
)
from fpgm.training.bundle_io import (
    BundleShapeError,
    compose_control_rgb,
    letterbox_to_white,
    load_window_arrays,
    normal_rgb_to_z_channel,
    read_gray_video,
    read_id_video,
    read_rgb_video,
)
from fpgm.training.dataset import (
    WindowBundleDataset,
    _largest_4k_plus_1,
    build_dataset_from_manifest,
)
from fpgm.training.manifest import discover_windows
from fpgm.training.types import (
    AugmentConfig,
    BundleAssemblyConfig,
    GateFilterConfig,
    ManifestEmptyError,
    WindowSample,
)
from fpgm.types import DataError

W, H, T = 16, 12, 5  # tiny synthetic bundle: 16x12 px, 4*1+1 = 5 frames


def _write_video(path, frames_rgb, fourcc: str):
    writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*fourcc), 15.0, (W, H), isColor=True)
    assert writer.isOpened(), f"could not open writer for {path} ({fourcc})"
    for frame in frames_rgb:
        writer.write(frame[..., ::-1])  # RGB -> BGR
    writer.release()


def _default_gates() -> dict:
    return {
        "s2_robot_alignment": True,
        "s7_drawer_prismatic_gated": True,
        "s7_support_plane_fit": True,
        "resolution_contract_hw16": True,
        "window_frame_count_4n1": True,
    }


def make_fake_window(
    root,
    episode_uuid="EP1",
    camera_serial="cam0",
    window_name="window_00000_00005",
    gates_passed: dict | None = None,
    overlaps_pose_gap: bool = False,
    rng_seed: int = 0,
    n_frames: int = T,
) -> WindowSample:
    """Build a complete, decodable fake ``vace/window_*`` bundle under ``root``."""
    rng = np.random.default_rng(rng_seed)
    window_dir = root / episode_uuid / camera_serial / "vace" / window_name
    window_dir.mkdir(parents=True, exist_ok=True)

    target = rng.integers(0, 255, size=(n_frames, H, W, 3), dtype=np.uint8)
    depth = rng.integers(0, 255, size=(n_frames, H, W), dtype=np.uint8)
    seg = rng.integers(0, 4, size=(n_frames, H, W), dtype=np.uint8)  # ids 0..3
    fg = (seg > 0).astype(np.uint8) * 255

    normals_xyz = rng.normal(size=(n_frames, H, W, 3)).astype(np.float32)
    normals_xyz /= np.linalg.norm(normals_xyz, axis=-1, keepdims=True) + 1e-8
    normal_rgb = np.stack([encode_normals_rgb(normals_xyz[t]) for t in range(n_frames)], axis=0)

    _write_video(window_dir / "target.mp4", target, "mp4v")
    _write_video(window_dir / "control_depth.mkv", np.repeat(depth[..., None], 3, axis=-1), "FFV1")
    _write_video(window_dir / "control_seg.mkv", np.repeat(seg[..., None], 3, axis=-1), "FFV1")
    _write_video(window_dir / "control_normal.mp4", normal_rgb, "mp4v")
    _write_video(window_dir / "fg_mask.mp4", np.repeat(fg[..., None], 3, axis=-1), "mp4v")

    # different aspect ratio than the control videos, on purpose
    ref_plate = rng.integers(0, 255, size=(H * 2, W * 3, 3), dtype=np.uint8)
    cv2.imwrite(str(window_dir / "ref_plate.png"), ref_plate[..., ::-1])

    (window_dir / "caption.txt").write_text("a fake caption for testing")

    payload = {
        "episode_uuid": episode_uuid,
        "camera_serial": camera_serial,
        "video_frame_start": 0,
        "video_frame_end": n_frames,
        "resolution": [W, H],
        "gates_passed": gates_passed if gates_passed is not None else _default_gates(),
        "overlaps_pose_gap": overlaps_pose_gap,
        "limitations": [],
    }
    (window_dir / "sample.json").write_text(json.dumps(payload))

    from fpgm.training.manifest import _parse_window_dir

    return _parse_window_dir(window_dir)


# --------------------------------------------------------------------------- #
# manifest
# --------------------------------------------------------------------------- #


class TestDiscoverWindows:
    def test_kept_when_all_gates_pass(self, tmp_path):
        make_fake_window(tmp_path)
        report = discover_windows(tmp_path)
        assert report.total_found == 1
        assert report.n_kept == 1
        assert report.n_dropped == 0

    def test_dropped_on_pose_gap_overlap_by_default(self, tmp_path):
        make_fake_window(tmp_path, overlaps_pose_gap=True)
        report = discover_windows(tmp_path)
        assert report.n_kept == 0
        assert report.dropped[0][1] == "overlaps_pose_gap"

    def test_pose_gap_overlap_kept_when_allowed(self, tmp_path):
        make_fake_window(tmp_path, overlaps_pose_gap=True)
        report = discover_windows(tmp_path, GateFilterConfig(allow_pose_gap_overlap=True))
        assert report.n_kept == 1

    def test_dropped_on_gate_failure(self, tmp_path):
        gates = _default_gates()
        gates["s2_robot_alignment"] = False
        make_fake_window(tmp_path, gates_passed=gates)
        report = discover_windows(tmp_path)
        assert report.n_kept == 0
        assert "gate_failed" in report.dropped[0][1]

    def test_ignored_gate_keys_tolerates_named_failure_only(self, tmp_path):
        gates = _default_gates()
        gates["s2_robot_alignment"] = False
        gates["s7_support_plane_fit"] = False
        make_fake_window(tmp_path, gates_passed=gates)
        # Ignoring only one of the two failed gates should still drop it.
        cfg = GateFilterConfig(ignored_gate_keys=("s2_robot_alignment",))
        report = discover_windows(tmp_path, cfg)
        assert report.n_kept == 0
        # Ignoring both should keep it.
        report2 = discover_windows(
            tmp_path,
            GateFilterConfig(ignored_gate_keys=("s2_robot_alignment", "s7_support_plane_fit")),
        )
        assert report2.n_kept == 1

    def test_dropped_when_sample_json_missing_known_gate_key(self, tmp_path):
        gates = _default_gates()
        del gates["s2_robot_alignment"]
        make_fake_window(tmp_path, gates_passed=gates)
        report = discover_windows(tmp_path)
        assert report.n_kept == 0
        assert "sample_json_missing_gate_keys" in report.dropped[0][1]

    def test_unknown_gate_key_dropped_only_if_required(self, tmp_path):
        gates = _default_gates()
        gates["s9_brand_new_gate"] = True
        make_fake_window(tmp_path, gates_passed=gates)
        assert discover_windows(tmp_path).n_kept == 1  # tolerated by default
        report = discover_windows(tmp_path, GateFilterConfig(require_known_gate_keys_only=True))
        assert report.n_kept == 0
        assert "unknown_gate_keys" in report.dropped[0][1]

    def test_scans_multiple_episodes_and_cameras(self, tmp_path):
        make_fake_window(
            tmp_path, episode_uuid="EP1", camera_serial="cam0", window_name="window_00000_00005"
        )
        make_fake_window(
            tmp_path, episode_uuid="EP2", camera_serial="cam1",
            window_name="window_00005_00010", rng_seed=1,
        )
        report = discover_windows(tmp_path)
        assert report.n_kept == 2

    def test_reason_counts(self, tmp_path):
        make_fake_window(
            tmp_path, episode_uuid="EP1", window_name="window_00000_00005", overlaps_pose_gap=True
        )
        make_fake_window(
            tmp_path, episode_uuid="EP2", window_name="window_00005_00010",
            overlaps_pose_gap=True, rng_seed=1,
        )
        report = discover_windows(tmp_path)
        assert report.reason_counts() == {"overlaps_pose_gap": 2}

    def test_missing_required_file_raises(self, tmp_path):
        make_fake_window(tmp_path)
        window_dir = tmp_path / "EP1" / "cam0" / "vace" / "window_00000_00005"
        (window_dir / "target.mp4").unlink()
        with pytest.raises(DataError):
            discover_windows(tmp_path)


# --------------------------------------------------------------------------- #
# bundle_io
# --------------------------------------------------------------------------- #


class TestBundleIO:
    def test_gray_video_round_trip_lossless(self, tmp_path):
        depth = np.random.default_rng(0).integers(0, 255, size=(T, H, W), dtype=np.uint8)
        path = tmp_path / "d.mkv"
        _write_video(path, np.repeat(depth[..., None], 3, axis=-1), "FFV1")
        decoded = read_gray_video(path, T)
        np.testing.assert_array_equal(decoded, depth)

    def test_id_video_round_trip_exact(self, tmp_path):
        seg = np.random.default_rng(0).integers(0, 4, size=(T, H, W), dtype=np.uint8)
        path = tmp_path / "s.mkv"
        _write_video(path, np.repeat(seg[..., None], 3, axis=-1), "FFV1")
        decoded = read_id_video(path, T, max_id=3)
        np.testing.assert_array_equal(decoded, seg)

    def test_id_video_rejects_ids_above_max(self, tmp_path):
        seg = np.full((T, H, W), 5, dtype=np.uint8)
        path = tmp_path / "s.mkv"
        _write_video(path, np.repeat(seg[..., None], 3, axis=-1), "FFV1")
        with pytest.raises(BundleShapeError):
            read_id_video(path, T, max_id=3)

    def test_wrong_frame_count_raises(self, tmp_path):
        frames = np.zeros((T, H, W, 3), dtype=np.uint8)
        path = tmp_path / "v.mp4"
        _write_video(path, frames, "mp4v")
        with pytest.raises(BundleShapeError):
            read_rgb_video(path, T + 1)

    def test_letterbox_preserves_aspect_and_pads(self):
        img = np.full((10, 20, 3), 200, dtype=np.uint8)  # wide image
        out = letterbox_to_white(img, target_w=20, target_h=20, fill=255)
        assert out.shape == (20, 20, 3)
        # Top/bottom bands should be the fill color; center should not be.
        assert (out[0, 0] == 255).all()
        assert (out[10, 10] != 255).any()

    def test_compose_control_rgb_shape_mismatch_raises(self):
        depth = np.zeros((T, H, W), dtype=np.uint8)
        seg = np.zeros((T, H, W), dtype=np.uint8)
        normal = np.zeros((T, H, W + 1), dtype=np.uint8)
        with pytest.raises(BundleShapeError):
            compose_control_rgb(depth, seg, normal, max_seg_id=16)

    def test_compose_control_rgb_channel_order(self):
        depth = np.full((1, 2, 2), 10, dtype=np.uint8)
        seg = np.full((1, 2, 2), 2, dtype=np.uint8)
        normal_z = np.full((1, 2, 2), 30, dtype=np.uint8)
        out = compose_control_rgb(depth, seg, normal_z, max_seg_id=16)
        assert out.shape == (1, 2, 2, 3)
        assert (out[..., 0] == 10).all()
        assert (out[..., 1] == round(2 * 255 / 16)).all()
        assert (out[..., 2] == 30).all()

    def test_normal_z_channel_round_trips_through_encode(self):
        normals = np.zeros((1, 2, 2, 3), dtype=np.float32)
        normals[..., 2] = 1.0  # pure +Z
        rgb = np.stack([encode_normals_rgb(normals[0])], axis=0)
        z = normal_rgb_to_z_channel(rgb)
        assert z.shape == (1, 2, 2)
        assert (z == 255).all()  # +1 in Z encodes to 255

    def test_load_window_arrays_matches_shapes(self, tmp_path):
        sample = make_fake_window(tmp_path)
        arrays = load_window_arrays(sample, BundleAssemblyConfig())
        assert arrays["target"].shape == (T, H, W, 3)
        assert arrays["control"].shape == (T, H, W, 3)
        assert arrays["fg_mask"].shape == (T, H, W)
        assert arrays["fg_mask"].dtype == np.bool_
        assert arrays["ref_plate"].ndim == 3
        assert isinstance(arrays["caption"], str) and arrays["caption"]


# --------------------------------------------------------------------------- #
# augment
# --------------------------------------------------------------------------- #


class TestAugment:
    def test_pose_noise_no_op_outside_mask(self):
        control = np.random.default_rng(0).integers(0, 255, size=(2, 20, 20, 3), dtype=np.uint8)
        mask = np.zeros((2, 20, 20), dtype=bool)  # nothing is foreground
        rng = np.random.default_rng(0)
        out = pose_noise(control, mask, rng, translate_px_std=5.0, rotate_deg_std=5.0)
        np.testing.assert_array_equal(out, control)

    def test_pose_noise_changes_foreground_with_nonzero_shift(self):
        control = np.zeros((1, 20, 20, 3), dtype=np.uint8)
        control[0, 5:15, 5:15] = 200  # a bright square to move
        mask = np.zeros((1, 20, 20), dtype=bool)
        mask[0, 5:15, 5:15] = True

        class _FixedRng:
            def normal(self, loc, scale):
                return loc + scale  # deterministic non-zero shift

        out = pose_noise(control, mask, _FixedRng(), translate_px_std=3.0, rotate_deg_std=0.0)
        assert not np.array_equal(out, control)

    def test_pose_noise_shape_mismatch_raises(self):
        control = np.zeros((2, 10, 10, 3), dtype=np.uint8)
        mask = np.zeros((3, 10, 10), dtype=bool)
        with pytest.raises(ValueError):
            pose_noise(control, mask, np.random.default_rng(0), 1.0, 1.0)

    def test_mask_dilate_grows_area(self):
        mask = np.zeros((1, 20, 20), dtype=bool)
        mask[0, 9:11, 9:11] = True
        area_before = mask.sum()

        class _ForceDilate:
            def choice(self, options):
                return "dilate"

            def integers(self, lo, hi):
                return 3

        out = mask_dilate_erode(mask, _ForceDilate(), max_kernel_px=5)
        assert out.sum() > area_before

    def test_mask_erode_shrinks_area(self):
        # A blob not touching the image border: cv2.erode's default border
        # handling treats the border as foreground (to avoid spurious
        # edge-erosion), so an all-True mask never shrinks -- use an interior
        # blob instead, mirroring the dilate test's fixture.
        mask = np.zeros((1, 20, 20), dtype=bool)
        mask[0, 5:15, 5:15] = True
        area_before = mask.sum()

        class _ForceErode:
            def choice(self, options):
                return "erode"

            def integers(self, lo, hi):
                return 3

        out = mask_dilate_erode(mask, _ForceErode(), max_kernel_px=5)
        assert out.sum() < area_before

    def test_mask_dilate_erode_none_is_unchanged(self):
        mask = np.random.default_rng(0).integers(0, 2, size=(1, 20, 20)).astype(bool)

        class _ForceNone:
            def choice(self, options):
                return "none"

        out = mask_dilate_erode(mask, _ForceNone(), max_kernel_px=5)
        np.testing.assert_array_equal(out, mask)

    def test_mask_dilate_erode_zero_kernel_is_unchanged(self):
        mask = np.random.default_rng(0).integers(0, 2, size=(1, 20, 20)).astype(bool)
        out = mask_dilate_erode(mask, np.random.default_rng(1), max_kernel_px=0)
        np.testing.assert_array_equal(out, mask)

    def test_channel_dropout_prob_one_zeroes_every_channel(self):
        control = np.random.default_rng(0).integers(1, 255, size=(2, 5, 5, 3), dtype=np.uint8)
        out = channel_dropout(control, np.random.default_rng(0), prob=1.0)
        assert (out == 0).all()

    def test_channel_dropout_prob_zero_is_unchanged(self):
        control = np.random.default_rng(0).integers(1, 255, size=(2, 5, 5, 3), dtype=np.uint8)
        out = channel_dropout(control, np.random.default_rng(0), prob=0.0)
        np.testing.assert_array_equal(out, control)

    def test_apply_augmentations_reproducible_for_same_seed_and_index(self):
        control = np.random.default_rng(0).integers(0, 255, size=(3, 20, 20, 3), dtype=np.uint8)
        mask = np.random.default_rng(1).integers(0, 2, size=(3, 20, 20)).astype(bool)
        cfg = AugmentConfig(seed=42)
        out1, mask1 = apply_augmentations(control, mask, cfg, sample_index=7)
        out2, mask2 = apply_augmentations(control, mask, cfg, sample_index=7)
        np.testing.assert_array_equal(out1, out2)
        np.testing.assert_array_equal(mask1, mask2)

    def test_apply_augmentations_disabled_is_identity(self):
        control = np.random.default_rng(0).integers(0, 255, size=(3, 20, 20, 3), dtype=np.uint8)
        mask = np.random.default_rng(1).integers(0, 2, size=(3, 20, 20)).astype(bool)
        cfg = AugmentConfig(
            enable_pose_noise=False, enable_mask_dilate_erode=False, enable_channel_dropout=False,
        )
        out, mask_out = apply_augmentations(control, mask, cfg, sample_index=0)
        np.testing.assert_array_equal(out, control)
        np.testing.assert_array_equal(mask_out, mask)


# --------------------------------------------------------------------------- #
# dataset
# --------------------------------------------------------------------------- #


class TestLargestFourKPlusOne:
    @pytest.mark.parametrize(
        "cap,expected", [(1, 1), (4, 1), (5, 5), (8, 5), (81, 81), (41, 41), (44, 41)]
    )
    def test_values(self, cap, expected):
        assert _largest_4k_plus_1(cap) == expected

    def test_rejects_zero(self):
        with pytest.raises(ValueError):
            _largest_4k_plus_1(0)


class TestWindowBundleDataset:
    def test_empty_samples_raises(self):
        with pytest.raises(ManifestEmptyError):
            WindowBundleDataset([])

    def test_getitem_shape_and_keys(self, tmp_path):
        sample = make_fake_window(tmp_path)
        ds = WindowBundleDataset([sample], augment_cfg=AugmentConfig())
        item = ds[0]
        expected_keys = {"video", "prompt", "vace_video", "vace_reference_image", "_window_name"}
        assert set(item) >= expected_keys
        assert len(item["video"]) == T
        assert len(item["vace_video"]) == T
        assert len(item["vace_reference_image"]) == 1
        assert item["video"][0].size == (W, H)
        assert item["vace_reference_image"][0].size == (W, H)
        assert "vace_video_mask" not in item  # include_fg_mask defaults False

    def test_include_fg_mask_adds_vace_video_mask(self, tmp_path):
        sample = make_fake_window(tmp_path)
        ds = WindowBundleDataset([sample], assembly_cfg=BundleAssemblyConfig(include_fg_mask=True))
        item = ds[0]
        assert "vace_video_mask" in item
        assert len(item["vace_video_mask"]) == T

    def test_max_num_frames_truncates_to_4k_plus_1(self, tmp_path):
        sample = make_fake_window(tmp_path, n_frames=9)  # 4*2+1
        ds = WindowBundleDataset([sample], max_num_frames=6)  # largest 4k+1 <= 6 is 5
        item = ds[0]
        assert len(item["video"]) == 5
        assert len(item["vace_video"]) == 5

    def test_index_wraps_modulo_len(self, tmp_path):
        sample = make_fake_window(tmp_path)
        ds = WindowBundleDataset([sample])
        assert len(ds) == 1
        ds[5]  # should not raise

    def test_no_augment_cfg_skips_augmentation(self, tmp_path):
        sample = make_fake_window(tmp_path)
        ds_plain = WindowBundleDataset([sample], augment_cfg=None)
        item1 = ds_plain[0]
        item2 = ds_plain[0]
        frame1, frame2 = np.array(item1["vace_video"][0]), np.array(item2["vace_video"][0])
        np.testing.assert_array_equal(frame1, frame2)

    def test_build_dataset_from_manifest_raises_when_all_filtered(self, tmp_path):
        make_fake_window(tmp_path, overlaps_pose_gap=True)
        with pytest.raises(ManifestEmptyError):
            build_dataset_from_manifest(tmp_path, GateFilterConfig())

    def test_build_dataset_from_manifest_ok_when_permissive(self, tmp_path):
        make_fake_window(tmp_path, overlaps_pose_gap=True)
        ds = build_dataset_from_manifest(tmp_path, GateFilterConfig(allow_pose_gap_overlap=True))
        assert len(ds) == 1