import importlib.util
import pathlib

import numpy as np
from openpi.shared import normalize

_SCRIPT_PATH = pathlib.Path(__file__).resolve().parents[1] / "scripts" / "compute_norm_stats.py"
_SPEC = importlib.util.spec_from_file_location("compute_norm_stats", _SCRIPT_PATH)
assert _SPEC is not None and _SPEC.loader is not None
compute_norm_stats = importlib.util.module_from_spec(_SPEC)
_SPEC.loader.exec_module(compute_norm_stats)


def test_online_sliding_norm_reuses_source_one_x_stats_unchanged():
    raw_actions = normalize.NormStats(
        mean=np.array([1, 2, 3, 4, 5, 6, 7], dtype=np.float32),
        std=np.array([10, 20, 30, 40, 50, 60, 70], dtype=np.float32),
        q01=np.array([-1, -2, -3, -4, -5, -6, -7], dtype=np.float32),
        q99=np.array([1, 2, 3, 4, 5, 6, 7], dtype=np.float32),
    )
    raw_state = normalize.NormStats(
        mean=np.array([1, 2], dtype=np.float32),
        std=np.array([3, 4], dtype=np.float32),
        q01=np.array([0, 1], dtype=np.float32),
        q99=np.array([2, 3], dtype=np.float32),
    )

    reused = compute_norm_stats._reuse_source_one_x_norm_stats(
        {"state": raw_state, "actions": raw_actions},
        (0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0),
    )

    assert reused["state"] is raw_state
    assert reused["actions"] is raw_actions
    np.testing.assert_allclose(reused["actions"].mean, raw_actions.mean)
    np.testing.assert_allclose(reused["actions"].std, raw_actions.std)
    np.testing.assert_allclose(reused["actions"].q01, raw_actions.q01)
    np.testing.assert_allclose(reused["actions"].q99, raw_actions.q99)