import json
import unittest

import numpy as np

from lab import INPUTS, MODEL, PARAMS, ROOT, RampEnv, rollout


class LabTests(unittest.TestCase):
    def test_commands_use_reference_physics_and_cannot_skip_the_ramp(self):
        env = RampEnv(1, variation=False)
        try:
            self.assertEqual(env.observe().shape, (1, INPUTS))
            for _ in range(51):
                env.step(np.array([[1, 0]]))
            self.assertTrue(env.frame["qualified"])
            self.assertEqual(env.frame["mode"], "air")
            np.testing.assert_allclose(
                env.frame["velocity"],
                [-23.533519868334647, 7.483007824212656, 0],
                atol=1e-7,
            )
            self.assertTrue(all(level == 0 for level in env.frame["firing"]))
        finally:
            env.close()

        env = RampEnv(1, variation=False)
        try:
            env.step(np.array([[0, 1]]))
            self.assertTrue(env.done[0])
            self.assertFalse(env.success[0])
            self.assertFalse(env.frame["qualified"])
            self.assertGreater(env.frame["y"], 0)
        finally:
            env.close()

    def test_terminal_episodes_are_frozen(self):
        env = RampEnv(1, variation=False)
        try:
            env.step(np.array([[0, 1]]))
            before = (env.frame, env.reward.tolist(), env.observe().copy())
            env.step(np.array([[1, -1]]))
            self.assertEqual(env.frame, before[0])
            self.assertEqual(env.reward.tolist(), before[1])
            np.testing.assert_array_equal(env.observe(), before[2])
        finally:
            env.close()

    def test_untrained_replay_is_deterministic_and_does_not_invent_jets(self):
        weights = np.zeros(PARAMS)
        first = rollout(weights, seed=17, capture=True)
        self.assertEqual(first, rollout(weights, seed=17, capture=True))
        self.assertFalse(first["success"][0])
        self.assertTrue(first["frames"][-1]["done"])
        self.assertEqual(len(first["frames"][0]["activation"]), 100)
        self.assertTrue(all(not any(f["firing"]) for f in first["frames"]))

    def test_checkpoint_lands_on_unseen_starts_using_neural_activity(self):
        saved = json.loads((ROOT / "data/checkpoint.json").read_text())
        self.assertEqual(saved["model"], MODEL)
        weights = np.repeat(np.array(saved["weights"])[None], 100, axis=0)
        result = rollout(weights, seed=20261016)
        self.assertGreaterEqual(sum(result["success"]), 75)
        self.assertLess(
            sum(rollout(weights, seed=20261016, silenced=True)["success"]), 10
        )
        trial = rollout(saved["weights"], seed=42, capture=True)
        self.assertTrue(trial["success"][0])
        last = trial["frames"][-1]
        self.assertTrue(last["qualified"])
        self.assertEqual(last["mode"], "ground")
        self.assertGreater(last["up"][1], 0.95)
        self.assertLess(last["impactSpeed"], 3.2)
        self.assertTrue(any(any(f["firing"]) for f in trial["frames"] if f["air"]))
        self.assertFalse(
            any(any(f["firing"]) for f in trial["frames"] if f["mode"] == "ramp")
        )


if __name__ == "__main__":
    unittest.main()