"""Regenerate CPU mj_ray references: run from repository root in .venv."""

import json
import math
import sys
from pathlib import Path

import mujoco
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from task_config import OBSTACLE_TASK, deployment_metadata, validate_task_config  # noqa: E402


def generate():
    config = validate_task_config(
        OBSTACLE_TASK, {"sensorCfg": {"sensorMode": "multi_ring_raycast"}}, 42
    )
    deployment = deployment_metadata(OBSTACLE_TASK, config, 42)
    boxes = deployment["terrain"]["boxes"]
    # Low obstacle directly ahead; platform edge is a bounded-floor drop (no invented height map).
    layouts = {
        "default": boxes,
        "low": [boxes[0], {"pos": [0.9, 0, 0.025], "size": [0.2, 0.5, 0.025], "yaw": 0}],
        "edge": [boxes[0]],
    }
    cases = []
    for layout, position, euler in [
        ("default", [-5, 0, 0.32], [0, 0, 0]),
        ("default", [-3, 1, 0.42], [0.23, -0.31, 0.51]),
        ("default", [-2, -2, 0.5], [-0.35, 0.27, -0.63]),
        ("low", [0, 0, 0.32], [0, 0, 0]),
        ("low", [0, 0, 0.32], [0.2, 0.1, -0.1]),
        ("edge", [5.45, 0, 0.32], [0, 0, 0]),
    ]:
        xml = (
            "<mujoco><worldbody>"
            + "".join(
                '<geom type="box" pos="{}" size="{}"/>'.format(
                    " ".join(map(str, b["pos"])), " ".join(map(str, b["size"]))
                )
                for b in layouts[layout]
            )
            + "</worldbody></mujoco>"
        )
        model = mujoco.MjModel.from_xml_string(xml)
        data = mujoco.MjData(model)
        mujoco.mj_forward(model, data)
        q = np.zeros(4)
        mujoco.mju_euler2Quat(q, np.array(euler), "xyz")
        matrix = np.zeros(9)
        mujoco.mju_quat2Mat(matrix, q)
        matrix = matrix.reshape(3, 3)
        origin = np.array(position) + matrix @ np.array([0.3, 0, 0.05])
        distances, ids = [], []
        for pitch in [0, -20, -45]:
            for i in range(16):
                yaw = math.radians(-45 + i * 90 / 15)
                p = math.radians(pitch)
                direction = matrix @ np.array(
                    [math.cos(p) * math.cos(yaw), math.cos(p) * math.sin(yaw), math.sin(p)]
                )
                geom = np.array([-1], dtype=np.int32)
                distance = mujoco.mj_ray(model, data, origin, direction, None, 1, -1, geom)
                distances.append(distance)
                ids.append(int(geom[0]))
        cases.append(
            dict(
                layout=layout,
                position=position,
                quaternion=q.tolist(),
                distances=distances,
                hitIds=ids,
                depth=[1 if d < 0 else min(1, d / 4) for d in distances],
            )
        )
    root = Path("web_platform/src/rl/fixtures")
    (root / "multiRingGolden.json").write_text(
        json.dumps(
            dict(source=f"CPU MuJoCo {mujoco.__version__} mj_ray", layouts=layouts, cases=cases),
            indent=2,
        )
        + "\n"
    )
    (root / "multiRingDeployment.json").write_text(json.dumps(deployment, indent=2) + "\n")
    Path("/tmp/go2-multi-ring-stage5/task.json").write_text(json.dumps(config))


if __name__ == "__main__":
    generate()