import json
import tempfile
import unittest

import numpy as np
from PIL import Image


class FakePolicy:
    def __init__(self, actions=None):
        self.observation = None
        self.noise = None
        self.actions = np.ones((10, 7)) if actions is None else actions

    def infer(self, observation, *, noise=None):
        self.observation = observation
        self.noise = noise
        return {"actions": self.actions}


class InferenceTests(unittest.TestCase):
    def test_prediction_maps_ur_observation_and_writes_json(self):
        from inference import ACTION_LABELS, run_prediction

        policy = FakePolicy()
        image = Image.new("RGB", (8, 6), "red")
        result = run_prediction(
            policy, image, image, "pick up", [1, 2, 3, 4, 5, 6, 0], 2,
            "owner/model", "checkpoint",
        )
        self.assertEqual(tuple(result.actions.columns), ACTION_LABELS)
        self.assertEqual(
            policy.observation["observation/state"].tolist(),
            [1, 2, 3, 4, 5, 6, 0],
        )
        self.assertEqual(policy.observation["observation/image"].dtype, np.uint8)
        self.assertEqual(policy.observation["prompt"], "pick up")
        self.assertEqual(policy.noise.shape, (10, 32))
        with open(result.json_path, encoding="utf-8") as stream:
            self.assertEqual(json.load(stream)["seed"], 44)

    def test_same_trial_uses_same_diffusion_noise(self):
        from inference import run_prediction

        first, second = FakePolicy(), FakePolicy()
        image = Image.new("RGB", (8, 6))
        arguments = (image, image, "task", [0] * 7, 3, "model", "checkpoint")
        run_prediction(first, *arguments)
        run_prediction(second, *arguments)
        np.testing.assert_array_equal(first.noise, second.noise)

    def test_prediction_rejects_invalid_state_and_action_shape(self):
        from inference import run_prediction

        image = Image.new("RGB", (8, 6))
        with self.assertRaisesRegex(ValueError, "seven"):
            run_prediction(FakePolicy(), image, image, "task", [1, 2], 0, "m", "c")
        with self.assertRaisesRegex(RuntimeError, r"\(10, 7\)"):
            run_prediction(
                FakePolicy(np.zeros((8, 7))), image, image, "task", [0] * 7,
                0, "m", "c",
            )


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