"""Deterministic native-math oracle, consumed by Vitest. Run from repository root."""

import json
from pathlib import Path

import numpy as np

from training_server.mobile_manipulator.kernel import ROBOTS, TASK, TaskKernel, decode_action
from training_server.mobile_manipulator.motion import SafeActionController


def generate():
    rng = np.random.default_rng(2026)
    cases = []
    for config in ROBOTS:
        kernel = TaskKernel(config)
        for i in range(20):
            s = rng.uniform(-2, 2, TASK["stateSize"])
            for start in [3, 33, 40, 50]:
                q = rng.normal(size=4)
                s[start : start + 4] = q / np.linalg.norm(q)
            s[29] = rng.uniform(0, 1)
            action = rng.uniform(-1.5, 1.5, TASK["actionSize"]).astype(np.float32)
            kernel.reset()
            kernel.has_lifted = bool(i % 2)
            obs, reward, terminated, truncated, info = kernel.evaluate(s)
            cases.append(
                {
                    "robotId": config["id"],
                    "state": s.tolist(),
                    "action": action.tolist(),
                    "lifted": bool(i % 2),
                    "observation": obs.tolist(),
                    "control": decode_action(config, action).tolist(),
                    "reward": reward,
                    "terminated": terminated,
                    "truncated": truncated,
                    "info": info,
                }
            )
    target = Path(__file__).resolve().parents[2] / "contracts/fixtures/mobile-golden.json"
    target.write_text(json.dumps(cases, indent=2) + "\n")
    motion_cases = []
    for config in ROBOTS:
        for stage in ("navigate", "reach", "pick-place"):
            s = np.zeros(TASK["stateSize"])
            s[[3, 33, 40, 50]] = 1
            s[37:40] = [*TASK["navigationOffset"][:2], TASK["objectStart"][2]]
            s[29] = 1
            for i, joint in enumerate(config["armJoints"]):
                s[13 + i] = joint["neutral"]
            motion = SafeActionController(config)
            motion.reset(s)
            frames = []
            for _ in range(12):
                action = rng.uniform(-2, 2, TASK["actionSize"]).astype(np.float32)
                motion.apply(action, s, stage, False)
                frames.append(
                    dict(
                        action=action.tolist(),
                        control=motion.control.tolist(),
                        applied=motion.applied.tolist(),
                        targets=motion.targets.tolist(),
                    )
                )
            motion_cases.append(
                dict(robotId=config["id"], stage=stage, state=s.tolist(), frames=frames)
            )
    target.with_name("mobile-motion-v2-golden.json").write_text(
        json.dumps(motion_cases, indent=2) + "\n"
    )


if __name__ == "__main__":
    generate()