DynamisLab/tests/test_drl_pinball_acquisition.py

import csv
import os
import subprocess
import sys
from collections import namedtuple
from pathlib import Path

import numpy as np
import pytest

from drl_pinball.acquisition import (
    accumulate_mean_fields,
    assign_phase,
    assign_periodic_phase,
    cleanup_scratch,
    complete_cycle_field_indices,
    create_scratch,
    default_reproduction_mapping,
    decode_legacy_physical_velocity,
    dual_cycle_dtw,
    pooled_phase_bins,
    prepare_role_output,
    publish_role_output,
    publish_selected_fields,
    select_phase_fields,
    smooth_center_uy,
    validate_output_storage,
    validate_modern_fields,
    write_boundary_artifacts,
    write_phase_cycle_artifacts,
)

Usage = namedtuple("Usage", "total used free")


def _storage_tree(tmp_path):
    mount = tmp_path / "optane"
    expected = mount / "DynamisLab" / "drl_pinball" / "reproduction"
    expected.mkdir(parents=True)
    mapping = tmp_path / "repo" / "src" / "drl_pinball" / "data" / "reproduction"
    mapping.parent.mkdir(parents=True)
    mapping.symlink_to(expected, target_is_directory=True)
    return mount, expected, mapping


def _validate_tmp_storage(tmp_path, **overrides):
    mount, expected, mapping = _storage_tree(tmp_path)
    kwargs = {
        "repo_mapping": mapping, "expected_root": expected, "optane_mount": mount,
        "min_free_bytes": 100, "is_mount": lambda path: path == mount,
        "disk_usage": lambda path: Usage(1000, 100, 900),
    }
    kwargs.update(overrides)
    return validate_output_storage(**kwargs), mount, expected, mapping


def test_default_reproduction_mapping_is_derivable(tmp_path):
    assert default_reproduction_mapping(tmp_path) == (
        tmp_path / "src" / "drl_pinball" / "data" / "reproduction"
    )


def test_storage_validation_defaults_output_to_repo_mapping(tmp_path):
    storage, _, expected, mapping = _validate_tmp_storage(tmp_path)
    assert storage["output_root"] == mapping
    assert storage["resolved_output_root"] == expected
def test_storage_validation_accepts_exact_mapping_and_configurable_same_device(tmp_path):

    mount, expected, mapping = _storage_tree(tmp_path)
    output = mount / "alternate" / "campaign"
    result = validate_output_storage(
        repo_mapping=mapping, output_root=output, expected_root=expected,
        optane_mount=mount, min_free_bytes=500, is_mount=lambda path: path == mount,
        disk_usage=lambda path: Usage(1000, 100, 900),
    )
    assert result["resolved_output_root"] == output
    assert result["free_bytes"] == 900


def test_storage_validation_rejects_mount_space_mapping_and_device_failures(tmp_path):
    mount, expected, mapping = _storage_tree(tmp_path)
    common = dict(repo_mapping=mapping, expected_root=expected, optane_mount=mount)
    with pytest.raises(ValueError, match="not a mount"):
        validate_output_storage(**common, is_mount=lambda path: False)
    with pytest.raises(OSError, match="free space"):
        validate_output_storage(
            **common, min_free_bytes=901, is_mount=lambda path: True,
            disk_usage=lambda path: Usage(1000, 100, 900),
        )
    wrong = tmp_path / "wrong"
    wrong.mkdir()
    mapping.unlink()
    mapping.symlink_to(wrong, target_is_directory=True)
    with pytest.raises(ValueError, match="resolve exactly"):
        validate_output_storage(**common, is_mount=lambda path: True)

    mapping.unlink()
    mapping.symlink_to(expected, target_is_directory=True)
    real_stat = os.stat
    outside = tmp_path / "outside"
    outside.mkdir()

    def different_device(path):
        value = real_stat(path)
        if Path(path) == outside:
            return type("Stat", (), {"st_dev": value.st_dev + 1})()
        return value

    with pytest.raises(ValueError, match="output root"):
        validate_output_storage(
            **common, output_root=outside / "new", is_mount=lambda path: True,
            stat=different_device,
        )


def test_prepare_role_output_conflict_and_scoped_overwrite(tmp_path):
    storage, _, _, _ = _validate_tmp_storage(tmp_path)
    prepared = prepare_role_output(storage, "v5", "karman_re100", "controlled")
    (prepared["role_dir"] / "old.txt").write_text("old")
    published = publish_role_output(prepared)
    sentinel = published / "old.txt"
    sibling = published.parent / "target"
    sibling.mkdir()
    (sibling / "keep.txt").write_text("keep")
    with pytest.raises(FileExistsError):
        prepare_role_output(storage, "v5", "karman_re100", "controlled")
    replaced = prepare_role_output(storage, "v5", "karman_re100", "controlled", overwrite=True)
    assert sentinel.exists()
    (replaced["role_dir"] / "new.txt").write_text("new")
    final = publish_role_output(replaced)
    assert not sentinel.exists() and (final / "new.txt").exists()
    assert (sibling / "keep.txt").read_text() == "keep"
    with pytest.raises(ValueError, match="path component"):
        prepare_role_output(storage, "v5", "karman_re100", "../target", overwrite=True)


def test_prepare_rejects_symlink_parent_without_following_it(tmp_path):
    storage, _, _, _ = _validate_tmp_storage(tmp_path)
    target = tmp_path / "external-pipeline"
    target.mkdir()
    pipeline = storage["resolved_output_root"] / "v5"
    pipeline.symlink_to(target, target_is_directory=True)
    with pytest.raises(ValueError, match="pipeline directory"):
        prepare_role_output(storage, "v5", "karman_re100", "controlled", overwrite=True)
    assert list(target.iterdir()) == []
def test_prepare_rejects_symlink_role_without_following_it(tmp_path):

    storage, _, _, _ = _validate_tmp_storage(tmp_path)
    target = tmp_path / "do-not-delete"
    target.mkdir()
    sentinel = target / "keep.txt"
    sentinel.write_text("keep")
    role = storage["resolved_output_root"] / "legacy" / "karman_re100" / "zero"
    role.parent.mkdir(parents=True)
    role.symlink_to(target, target_is_directory=True)
    with pytest.raises(ValueError, match="must not be a symlink"):
        prepare_role_output(storage, "legacy", "karman_re100", "zero", overwrite=True)
    assert sentinel.read_text() == "keep"


def test_phase_complete_interpolated_half_open():
    times = np.arange(13, dtype=float)
    sensors = np.zeros((13, 6))
    sensors[:, 3] = np.sin(2 * np.pi * (times - 0.25) / 4)
    result = assign_phase(times, sensors)
    assert np.allclose(result["crossing_times"], [0.29289322, 4.29289322, 8.29289322])
    assert result["cycle_id"].tolist() == [-1] + [0] * 4 + [1] * 4 + [-1] * 4
    assert np.isnan(result["phase"][9])


def test_crossing_exact_zero_and_plateau_are_single_crossings():
    times = np.arange(10, dtype=float)
    sensors = np.zeros((10, 6))
    sensors[:, 3] = [-1, 0, 0, 1, 1, -1, 0, 0, 1, 1]
    result = assign_phase(times, sensors)
    assert np.array_equal(result["crossing_times"], [2.0, 7.0])
    assert result["cycle_id"].tolist() == [-1, -1, 0, 0, 0, 0, 0, -1, -1, -1]
    assert result["phase"][2] == 0.0


def test_selection_is_global_eight_with_metadata_and_index_tie():
    field_times = np.concatenate((np.arange(0.5, 4.0, 0.5), np.arange(4.0, 8.0, 0.5)))
    result = select_phase_fields(field_times, np.array([0.0, 4.0, 8.0]))
    assert result["field_indices"].shape == (8,)
    assert np.array_equal(result["field_indices"], [7, 0, 1, 2, 3, 4, 5, 6])
    assert np.array_equal(result["cycle_id"], [1, 0, 0, 0, 0, 0, 0, 0])
    assert np.allclose(result["actual_phase"], np.arange(8) * np.pi / 4, atol=1e-15)
    assert np.allclose(result["phase_error"], [0.0] * 8, atol=1e-15)


def test_smooth_center_uy_binomial_and_endpoint_copy():
    assert np.allclose(smooth_center_uy([1.0, 5.0, 1.0]), [1.0, 3.0, 1.0])
    assert np.allclose(smooth_center_uy([2.0, 4.0]), [2.0, 4.0])


def test_assign_phase_default_behavior_unchanged_with_jitter():
    times = np.arange(12, dtype=float)
    sensors = np.zeros((12, 6))
    sensors[:, 3] = [-1, 1, -0.1, 0.1, -1, 1, -1, 1, -1, 1, -1, 1]
    result = assign_phase(times, sensors)
    assert len(result["crossing_times"]) == 6
    assert set(result) == {"phase", "cycle_id", "crossing_times"}


def test_assign_periodic_phase_filters_jitter_with_minimum_interval():
    times = np.arange(20, dtype=float)
    sensors = np.zeros((20, 6))
    # Deep brief negative between rises so smoothing keeps a close false crossing.
    sensors[:, 3] = [-2, -2, -1, 1, 2, -3, 1, 2, 2, -2, -1, 1, 2, 2, -2, -1, 1, 2, 2, -2]
    raw = assign_phase(times, sensors)
    filtered = assign_periodic_phase(times, sensors, minimum_crossing_interval=4.0)
    assert len(raw["crossing_times"]) >= 4
    assert filtered["rejected_crossing_count"] >= 1
    assert 5.75 in filtered["rejected_crossing_times"]
    assert filtered["accepted_crossing_count"] == len(filtered["crossing_times"])
    assert filtered["complete_cycle_count"] == filtered["accepted_crossing_count"] - 1
    assert np.all(np.diff(filtered["crossing_times"]) >= 4.0 - 1e-12)
    assert filtered["smoothing_kernel"] == [0.25, 0.5, 0.25]
    first, last = filtered["crossing_times"][0], filtered["crossing_times"][-1]
    valid = (times >= first) & (times < last)
    assert np.all(filtered["cycle_id"][valid] >= 0)
    assert np.all(filtered["cycle_id"][~valid] < 0)


def test_complete_cycle_mean_excludes_outside_first_last_crossing():
    times = np.arange(10, dtype=float)
    crossings = np.array([2.0, 6.0, 9.0])
    selected = complete_cycle_field_indices(times, crossings)
    assert selected["field_indices"].tolist() == [2, 3, 4, 5, 6, 7, 8]
    assert selected["mean_field_count"] == 7
    assert selected["complete_cycle_count"] == 2
    ux = np.arange(10, dtype=np.float32).reshape(10, 1, 1)
    uy = (2 * np.arange(10)).astype(np.float32).reshape(10, 1, 1)
    mean = accumulate_mean_fields(ux, uy, selected["field_indices"])
    assert mean["mean_ux"].shape == (1, 1)
    assert mean["mean_uy"].shape == (1, 1)
    assert float(mean["mean_ux"][0, 0]) == pytest.approx(5.0)
    assert float(mean["mean_uy"][0, 0]) == pytest.approx(10.0)
    # Eight-slot average of selected phase snapshots must remain a different concept.
    eight = accumulate_mean_fields(ux, uy, np.array([2, 3, 4, 5, 6, 7, 8, 8]))
    assert float(eight["mean_ux"][0, 0]) != pytest.approx(float(mean["mean_ux"][0, 0]))


def test_pooled_32_bins_population_std_and_original_indices():
    result = pooled_phase_bins(
        np.array([np.pi / 32, 3 * np.pi / 32, np.pi / 32, np.nan]),
        {"sensors": np.array([[1.0, 10.0], [4.0, 40.0], [3.0, 30.0], [9.0, 90.0]])},
    )
    assert result["phase_bin_count"][:2].tolist() == [2, 1]
    assert np.allclose(result["mean"]["sensors"][0], [2, 20])
    assert np.allclose(result["std"]["sensors"][0], [1, 10])
    assert result["sorted_index"].tolist() == [0, 2, 1]
    assert np.array_equal(result["pooled"]["sensors"], [[1, 10], [3, 30], [4, 40]])


def test_named_phase_pooling_preserves_rows_with_unavailable_nan_variables():
    phase = np.array([0.1, 0.2, 0.3])
    result = pooled_phase_bins(phase, {
        "sensors": np.array([[1.0], [2.0], [3.0]]),
        "reward_raw": np.full(3, np.nan),
    })
    assert result["sorted_index"].tolist() == [0, 1, 2]
    assert result["phase_bin_count"].sum() == 3
    assert np.all(np.isnan(result["mean"]["reward_raw"]))
    assert np.any(np.isfinite(result["mean"]["sensors"]))


def test_dual_dtw_preserves_native_series_and_separates_rolling_normalized():
    step = np.arange(90)
    target = np.column_stack(
        [(channel + 1) * np.sin(2 * np.pi * step / 30 + channel) for channel in range(6)]
    )
    native = np.linspace(0.125, 0.875, 90, dtype=np.float32)
    native_bytes = native.tobytes()
    result = dual_cycle_dtw(target, target, native, cycle_length=30, epsilon=1e-9)
    assert result["native_reward_dtw"] is native
    assert result["native_reward_dtw"].tobytes() == native_bytes
    assert np.all(np.isnan(result["target_normalized_dtw"][:29]))
    assert np.all(np.isfinite(result["target_normalized_dtw"][29:]))
    assert np.allclose(result["target_normalized_dtw"][[29, 59, 89]], 1.0)
    assert result["metadata"]["lag_channel"] == 1
    assert result["metadata"]["epsilon"] == 1e-9
    assert np.array_equal(result["metadata"]["target_channels"], np.arange(6))
    assert np.array_equal(result["metadata"]["state_channels"], np.arange(6))


def test_dual_dtw_unavailable_native_is_nan_series():
    step = np.arange(60)
    target = np.column_stack([np.sin(2 * np.pi * step / 20 + channel) for channel in range(6)])
    result = dual_cycle_dtw(target, target, cycle_length=20)
    assert result["native_reward_dtw"].shape == (60,)
    assert np.all(np.isnan(result["native_reward_dtw"]))
    assert result["target_normalized_dtw"].shape == (60,)


def test_boundary_and_complete_phase_cycle_artifact_schemas(tmp_path):
    boundary = {
        "time": np.array([0.0, 1.0]),
        "sensor": np.array([[1.0, np.nan], [2.0, 3.0]]),
        "field": np.array([4.0, np.nan]),
    }
    npz_path, csv_path = write_boundary_artifacts(
        tmp_path / "raw.npz", tmp_path / "raw.csv", boundary
    )
    with np.load(npz_path, allow_pickle=False) as saved:
        assert np.isnan(saved["sensor"][0, 1])
    with csv_path.open(newline="") as stream:
        rows = list(csv.DictReader(stream))
        assert list(rows[0]) == ["time", "sensor_0", "sensor_1", "field"]
        assert rows[0]["sensor_1"] == "nan"

    folded = pooled_phase_bins(
        np.array([0.1, 0.2]), {"sensors": np.array([[1.0, 10.0], [3.0, 30.0]])}
    )
    npz_path, csv_path = write_phase_cycle_artifacts(
        tmp_path / "phase_cycle.npz", tmp_path / "phase_cycle.csv", folded
    )
    with np.load(npz_path, allow_pickle=False) as saved:
        assert set(saved.files) == {"sorted_index", "sorted_phase", "phase_center",
                                    "phase_bin_count", "sensors_pooled", "sensors_mean",
                                    "sensors_std"}
        assert np.array_equal(saved["sorted_index"], [0, 1])
        assert np.array_equal(saved["sensors_pooled"], [[1.0, 10.0], [3.0, 30.0]])
        assert len(saved["phase_bin_count"]) == 32
    with csv_path.open(newline="") as stream:
        rows = list(csv.DictReader(stream))
    assert {"sensors_0_pooled", "sensors_1_mean", "sensors_1_std", "phase_bin_count"}.issubset(rows[0])
    assert [row["representation"] for row in rows[:2]] == ["pooled", "pooled"]
    assert all(row["representation"] == "bin" for row in rows[2:])
    assert len(rows[2:]) == 32
    assert [row["original_index"] for row in rows[:2]] == ["0", "1"]


def test_publication_and_cleanup(tmp_path):
    root = tmp_path / "scratch"
    scratch = create_scratch(root)
    (scratch / "x").write_text("x")
    destination = tmp_path / "pub" / "fields.npz"
    publish_selected_fields(destination, {"ux": np.ones((2, 3, 4))})
    with pytest.raises(FileExistsError):
        publish_selected_fields(destination, {"ux": np.ones((1, 1, 1))})
    cleanup_scratch(scratch, root=root)
    assert not scratch.exists()
    with pytest.raises(ValueError):
        cleanup_scratch(root, root=root)


def test_legacy_physical_decoder_is_mask_aware_and_pressure_invariant():
    ny, nx = 2, 3
    populations = np.zeros((9, ny, nx), dtype=np.float32)
    populations[1] = 0.01
    populations[2] = -0.02
    flags = np.ones((nx, ny), dtype=np.uint8)
    flags[1, 0] = np.uint8(0b00010001)  # sensor plus FLUID remains fluid
    flags[2, 1] = np.uint8(0b00000010)
    populations[:, 1, 2] = np.nan  # ignored nonfluid garbage

    first = decode_legacy_physical_velocity(
        populations.reshape(-1), flags=flags, nx=nx, ny=ny,
        include_pressure=True,
    )
    populations[0, :, :] = 7.0  # pressure/f0 changes, velocity does not
    populations[0, 1, 2] = np.nan
    second = decode_legacy_physical_velocity(
        populations.reshape(-1), flags=flags, nx=nx, ny=ny,
        include_pressure=True,
    )

    fluid = ((flags & 1) != 0).T
    assert first["ux"].shape == (ny, nx)
    assert np.allclose(first["ux"][fluid], 0.01)
    assert np.allclose(first["uy"][fluid], -0.02)
    assert np.array_equal(first["ux"], second["ux"])
    assert np.array_equal(first["uy"], second["uy"])
    assert first["ux"][1, 2] == first["uy"][1, 2] == first["pressure"][1, 2] == 0.0
    assert not np.array_equal(first["pressure"][fluid], second["pressure"][fluid])


def test_legacy_physical_decoder_accepts_zero_sum_nonzero_momentum():
    populations = np.zeros((9, 1, 1), dtype=np.float32)
    populations[1, 0, 0] = 0.01
    populations[0, 0, 0] = -0.01
    decoded = decode_legacy_physical_velocity(
        populations.reshape(-1), flags=np.ones((1, 1), np.uint8), nx=1, ny=1,
    )
    assert decoded["ux"][0, 0] == pytest.approx(0.01)


def test_legacy_physical_decoder_rejects_nonfinite_fluid_populations():
    populations = np.zeros((9, 1, 1), dtype=np.float32)
    populations[0, 0, 0] = np.nan
    with pytest.raises(ValueError, match="fluid populations.*non-finite"):
        decode_legacy_physical_velocity(
            populations.reshape(-1), flags=np.ones((1, 1), np.uint8), nx=1, ny=1,
        )

def test_module_import_does_not_load_gpu_dependencies():
    code = (
        "import sys; import drl_pinball.acquisition; "
        "assert 'pycuda' not in sys.modules; assert 'torch' not in sys.modules"
    )
    subprocess.run([sys.executable, "-c", code], check=True)