from pathlib import Path

import numpy as np

from SR_analysis.checks import policy_replay


class _Model:
    def predict(self, observation, deterministic=True):
        values = np.asarray(observation, dtype=np.float64)
        return values[:3], None


def test_replay_uses_previous_post_state_for_next_action(tmp_path, monkeypatch):
    trajectory = tmp_path / "controlled.npz"
    sensors = np.arange(24, dtype=np.float32).reshape(4, 6) / 100.0
    forces = np.zeros((4, 6), dtype=np.float32)
    actions = np.vstack((np.zeros(3), sensors[:-1, :3])).astype(np.float32)
    np.savez(trajectory, sensors=sensors, forces=forces, actions_norm=actions)

    cfg = {
        "scene_id": "karman",
        "s_dim": 12,
        "action_layout": ("front", "upper", "lower"),
    }
    norm_path = tmp_path / "norm.json"
    norm_path.write_text("{}")
    model_path = tmp_path / "model.zip"
    model_path.write_bytes(b"model")
    monkeypatch.setattr(policy_replay, "get_scene", lambda scene: cfg)
    monkeypatch.setattr(
        policy_replay,
        "load_existing_norm",
        lambda scene, config: ({"force_norm_fact": 1.0, "sens_deviation": [0] * 6, "sens_norm_fact": [1] * 6}, norm_path),
    )
    monkeypatch.setattr(policy_replay, "model_path_for_scene", lambda scene: str(model_path))
    monkeypatch.setattr(policy_replay, "load_ppo_model", lambda *args, **kwargs: _Model())
    monkeypatch.setattr(policy_replay, "normalize_raw_observation", lambda raw, norm: np.asarray(raw)[:6])

    report = policy_replay.replay("karman_re100", trajectory)

    assert report["status"] == "passed"
    assert report["n_compared"] == 3
    np.testing.assert_allclose(report["max_abs_error"], 0.0)
