"""Compare browser-exported fixed actions/rollout with the same native MJCF."""

import argparse
import json
from pathlib import Path

import mujoco
import numpy as np
from gymnasium.utils.env_checker import check_env

from .env import MobileManipulatorEnv


def validate(package, rollout, atol=2e-5):
    reference = json.loads(Path(rollout).read_text())
    env = MobileManipulatorEnv(package)
    maxima = dict(qpos=0.0, observation=0.0, control=0.0, reward=0.0)
    try:
        for action, expected in zip(reference["actions"], reference["rollout"], strict=True):
            observation, reward, _, _, info = env.step(action)
            for name, actual, target in [
                ("qpos", env.data.qpos, expected["qpos"]),
                ("observation", observation, expected["observation"]),
                ("control", env.data.ctrl, expected["ctrl"]),
                ("reward", reward, expected["reward"]),
            ]:
                error = float(np.max(np.abs(np.asarray(actual) - np.asarray(target))))
                maxima[name] = max(maxima[name], error)
                np.testing.assert_allclose(actual, target, atol=atol, rtol=0, err_msg=name)
            if info["stage"] != expected["stage"]:
                raise AssertionError("stage mismatch")
        check_env(env, skip_render_check=True)
    finally:
        env.close()
    return dict(mujoco=mujoco.__version__, steps=len(reference["actions"]), max_error=maxima)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--package", required=True)
    parser.add_argument("--rollout", required=True)
    args = parser.parse_args()
    print(json.dumps(validate(args.package, args.rollout), indent=2))


if __name__ == "__main__":
    main()